
#include "array.h"
#include "dmemory.h"
#include "process_int.h"
#include "cdo_wtime.h"
#include "cdo_options.h"
#include "remap_vars.h"
#include "timer.h"
#include "cimdOmp.h"

/*
  -----------------------------------------------------------------------

  Performs the remapping based on weights computed elsewhere

  -----------------------------------------------------------------------
*/

static void
remapFirstOrder(double *restrict dst_array, const RemapVars &rv, const double *restrict src_array)
{
  auto num_links = rv.num_links;
  auto num_wts = rv.num_wts;
  const auto &map_wts = rv.wts;
  const auto &dst_add = rv.tgt_cell_add;
  const auto &src_add = rv.src_cell_add;
  const auto &links = rv.links;
  long links_per_value = rv.links_per_value;

  if (links.option)
    {
#ifdef HAVE_OPENMP4
#pragma omp parallel for simd default(none) shared(num_links, dst_array, dst_add)
#endif
      for (size_t n = 0; n < num_links; ++n) dst_array[dst_add[n]] = 0.;

      for (size_t j = 0; j < links.num_blks; ++j)
        {
          const size_t *restrict dst_addx = links.dst_add[j];
          const size_t *restrict src_addx = links.src_add[j];
          const size_t *restrict windex = links.w_index[j];
          auto nlinks = links.num_links[j];

#ifdef HAVE_OPENMP4
#pragma omp parallel for simd default(none) shared(nlinks, dst_array, src_array, dst_addx, src_addx, map_wts, num_wts, windex)
#endif
          for (size_t n = 0; n < nlinks; ++n)
            {
              dst_array[dst_addx[n]] += src_array[src_addx[n]] * map_wts[num_wts * windex[n]];
            }
        }
    }
  else
    {
      long lpv = links_per_value;
      if (lpv > 0)
        {
          size_t nlinks = num_links / lpv;

          if (lpv == 1)
            {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(dst_array, src_array, dst_add, src_add, map_wts, nlinks)
#endif
              for (size_t n = 0; n < nlinks; ++n)
                {
                  dst_array[dst_add[n]] = src_array[src_add[n]] * map_wts[n];
                }
            }
          else if (lpv == 4)
            {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(dst_array, src_array, dst_add, src_add, map_wts, num_wts, nlinks)
#endif
              for (size_t n = 0; n < nlinks; ++n)
                {
                  const auto noff = n * 4;
                  dst_array[dst_add[noff]] = src_array[src_add[noff]] * map_wts[num_wts * (noff)]
                                             + src_array[src_add[noff + 1]] * map_wts[num_wts * (noff + 1)]
                                             + src_array[src_add[noff + 2]] * map_wts[num_wts * (noff + 2)]
                                             + src_array[src_add[noff + 3]] * map_wts[num_wts * (noff + 3)];
                }
            }
          else
            {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(dst_array, src_array, dst_add, src_add, map_wts, num_wts, nlinks, lpv)
#endif
              for (size_t n = 0; n < nlinks; ++n)
                {
                  const auto noff = n * lpv;
                  dst_array[dst_add[noff]] = src_array[src_add[noff]] * map_wts[num_wts * noff];
                  for (size_t k = 1; k < (size_t) lpv; ++k)
                    dst_array[dst_add[noff]] += src_array[src_add[noff + k]] * map_wts[num_wts * (noff + k)];
                }
            }
        }
      else
        {
#ifdef SX
#pragma cdir nodep
#endif
          for (size_t n = 0; n < num_links; ++n) dst_array[dst_add[n]] = 0.;

          for (size_t n = 0; n < num_links; ++n)
            {
              // printf("%5zu %5zu %5zu %g # dst_add src_add n\n", dst_add[n], src_add[n], n, map_wts[num_wts*n]);
              dst_array[dst_add[n]] += src_array[src_add[n]] * map_wts[num_wts * n];
            }
        }
    }
}

static void
remapSecondOrder(double *restrict dst_array, const RemapVars &rv, const double *restrict src_array, RemapGradients &gradients)
{
  const auto &grad1 = gradients.grad_lat;
  const auto &grad2 = gradients.grad_lon;
  const auto &grad3 = gradients.grad_latlon;

  const auto num_links = rv.num_links;
  const auto num_wts = rv.num_wts;
  const auto &map_wts = rv.wts;
  const auto &dst_add = rv.tgt_cell_add;
  const auto &src_add = rv.src_cell_add;

#ifdef SX
#pragma cdir nodep
#endif
  for (size_t n = 0; n < num_links; ++n) dst_array[dst_add[n]] = 0.;

  if (num_wts == 3)
    {
      for (size_t n = 0; n < num_links; ++n)
        {
          const auto i = src_add[n];
          const auto w = &map_wts[3 * n];
          dst_array[dst_add[n]] += src_array[i] * w[0] + grad1[i] * w[1] + grad2[i] * w[2];
          // printf("%zu %zu %.5f %.5f %.5f %.5f %.5f\n", n, src_add[n], grad1[i], grad2[i], w[0], w[1], w[2]);
        }
    }
  else if (num_wts == 4)
    {
      for (size_t n = 0; n < num_links; ++n)
        {
          const auto i = src_add[n];
          const auto w = &map_wts[4 * n];
          dst_array[dst_add[n]] += src_array[i] * w[0] + grad1[i] * w[1] + grad2[i] * w[2] + grad3[i] * w[3];
        }
    }
}

void
remap(double *restrict dst_array, double missval, size_t dst_size, const RemapVars &rv, const double *restrict src_array,
      RemapGradients &gradients)
{
  /*
    Input arrays:

      dst_add    destination address for each link
      src_add    source      address for each link
      num_wts    num of weights used in remapping
      map_wts    remapping weights for each link
      src_array  array with source field to be remapped

    Optional:

      gradients  gradient arrays on source grid necessary for higher-order remappings

    Output variables:

      dst_array  array for remapped field on destination grid
  */
  extern int timer_remap;

  double start = Options::cdoVerbose ? cdo_get_wtime() : 0;

  // Check the order of the interpolation

  const bool firstOrder = gradients.grad_lat.size() == 0;

#ifdef HAVE_OPENMP4
#pragma omp parallel for simd default(none) shared(dst_size, dst_array, missval)
#endif
  for (size_t n = 0; n < dst_size; ++n) dst_array[n] = missval;

  if (Options::Timer) timer_start(timer_remap);

  if (firstOrder)  // First order remapping
    {
      remapFirstOrder(dst_array, rv, src_array);
    }
  else  // Second order remapping
    {
      remapSecondOrder(dst_array, rv, src_array, gradients);
    }

  if (Options::cdoVerbose) cdoPrint("Remap: %.2f seconds", cdo_get_wtime() - start);

  if (Options::Timer) timer_stop(timer_remap);
}

static size_t
get_max_add(size_t num_links, size_t size, const std::vector<size_t> &add)
{
  std::vector<size_t> isum(size, 0);

  for (size_t n = 0; n < num_links; ++n) isum[add[n]]++;

  size_t max_add = 0;
  for (size_t i = 0; i < size; ++i)
    if (isum[i] > max_add) max_add = isum[i];

  return max_add;
}

static size_t
binary_search_int(const std::vector<size_t> &array, size_t len, size_t value)
{
  int64_t low = 0, high = len - 1;

  while (low <= high)
    {
      const auto midpoint = low + (high - low) / 2;

      // check to see if value is equal to item in array
      if (value == array[midpoint]) return midpoint;

      if (value < array[midpoint])
        high = midpoint - 1;
      else
        low = midpoint + 1;
    }

  // item was not found
  return len;
}

/*
  -----------------------------------------------------------------------

  Performs the remapping based on weights computed elsewhere

  -----------------------------------------------------------------------
*/
void
remap_laf(double *restrict dst_array, double missval, size_t dst_size, const RemapVars &rv, const double *restrict src_array)
{
  /*
    Input:
      src_array : array with source field to be remapped

    Output:
      dst_array : array for remapped field on destination grid
  */
  auto num_links = rv.num_links;
  auto num_wts = rv.num_wts;        // num of weights used in remapping
  const auto &map_wts = rv.wts;           // remapping weights for each link
  const auto &dst_add = rv.tgt_cell_add;  // destination address for each link
  const auto &src_add = rv.src_cell_add;  // source      address for each link

  varrayFill(dst_size, dst_array, missval);

  if (num_links == 0) return;

  size_t max_cls = get_max_add(num_links, dst_size, dst_add);

#ifdef _OPENMP
  VARRAY_2D(double, src_cls2, Threading::ompNumThreads, max_cls);
  VARRAY_2D(double, src_wts2, Threading::ompNumThreads, max_cls);
#else
  Varray<double> src_cls(max_cls);
  Varray<double> src_wts(max_cls);
#endif

  for (size_t n = 0; n < num_links; ++n)
    if (DBL_IS_EQUAL(dst_array[dst_add[n]], missval)) dst_array[dst_add[n]] = 0.0;

#ifdef _OPENMP
#pragma omp parallel for default(none) shared(dst_size, src_cls2, src_wts2, num_links, dst_add, src_add, src_array, map_wts, \
                                              num_wts, dst_array, max_cls) schedule(dynamic, 1)
#endif
  for (size_t i = 0; i < dst_size; ++i)
    {
      size_t k;
      size_t ncls;
#ifdef _OPENMP
      const auto ompthID = cdo_omp_get_thread_num();
      auto &src_cls = src_cls2[ompthID];
      auto &src_wts = src_wts2[ompthID];
#endif
      varrayFill(src_cls, 0.0);
      varrayFill(src_wts, 0.0);
      /*
      ncls = 0;
      for ( n = 0; n < num_links; n++ )
        {
          if ( i == dst_add[n] )
            {
              for ( k = 0; k < ncls; k++ )
                if ( IS_EQUAL(src_array[src_add[n]], src_cls[k]) ) break;

              if ( k == ncls )
                {
                  src_cls[k] = src_array[src_add[n]];
                  ncls++;
                }

              src_wts[k] += map_wts[num_wts*n];
            }
        }
      */
      /* only for sorted dst_add! */
      {
        size_t min_add = 1, max_add = 0;

        auto n = binary_search_int(dst_add, num_links, i);

        if (n < num_links)
          {
            min_add = n;

            for (n = min_add + 1; n < num_links; ++n)
              if (i != dst_add[n]) break;

            max_add = n;

            for (n = min_add; n > 0; --n)
              if (i != dst_add[n - 1]) break;

            min_add = n;
          }

        ncls = 0;
        for (n = min_add; n < max_add; ++n)
          {
            const auto value = src_array[src_add[n]];

            for (k = 0; k < ncls; ++k)
              if (IS_EQUAL(value, src_cls[k])) break;

            if (k == ncls)
              {
                src_cls[k] = value;
                ncls++;
              }

            src_wts[k] += map_wts[num_wts * n];
          }
        // printf("i, min_add, max_add, ncls %zu %zu %zu %zu\n", i, min_add, max_add, ncls);
      }

      if (ncls)
        {
          size_t imax = 0;
          auto wts = src_wts[0];
          for (k = 1; k < ncls; ++k)
            {
              if (src_wts[k] > wts)
                {
                  wts = src_wts[k];
                  imax = k;
                }
            }

          // for (k = 0; k < ncls; ++k) printf("   k, src_wts[k],  src_cls[k] %zu %g %g\n", k, src_wts[k],  src_cls[k]);
          // printf("imax, src_wts[imax],  src_cls[imax] %zu %zu %g %g\n", i , imax, src_wts[imax],  src_cls[imax]);
          dst_array[i] = src_cls[imax];
        }
    }
}

/*
  -----------------------------------------------------------------------

  Performs the remapping based on weights computed elsewhere

  -----------------------------------------------------------------------
*/
void
remap_sum(double *restrict dst_array, double missval, size_t dst_size, const RemapVars &rv, const double *restrict src_array)
{
  const auto num_links = rv.num_links;
  const auto num_wts = rv.num_wts;        // num of weights used in remapping
  const auto &map_wts = rv.wts;           // remapping weights for each link
  const auto &dst_add = rv.tgt_cell_add;  // destination address for each link
  const auto &src_add = rv.src_cell_add;  // source      address for each link

  for (size_t n = 0; n < dst_size; ++n) dst_array[n] = missval;

#ifdef SX
#pragma cdir nodep
#endif
  for (size_t n = 0; n < num_links; ++n)
    if (DBL_IS_EQUAL(dst_array[dst_add[n]], missval)) dst_array[dst_add[n]] = 0.0;

  for (size_t n = 0; n < num_links; ++n)
    {
      // printf("%5d %5d %5d %g # dst_add src_add n\n", dst_add[n], src_add[n], n, map_wts[num_wts*n]);
      // dst_array[dst_add[n]] += src_array[src_add[n]]*map_wts[num_wts*n];
      dst_array[dst_add[n]] += src_array[src_add[n]] * map_wts[num_wts * n];
      printf("%zu %zu %zu %g %g %g\n", n, dst_add[n], src_add[n], src_array[src_add[n]], map_wts[num_wts * n],
             dst_array[dst_add[n]]);
    }
}

void
remapVarsInit(RemapMethod mapType, int remapOrder, RemapVars &rv)
{
  // Initialize all pointer
  if (!rv.pinit) rv.pinit = true;

  rv.sort_add = (mapType == RemapMethod::CONSERV_SCRIP);

  // Determine the number of weights
  rv.num_wts = (mapType == RemapMethod::CONSERV_SCRIP) ? 3 : ((mapType == RemapMethod::BICUBIC) ? 4 : 1);
  if (mapType == RemapMethod::CONSERV && remapOrder == 2) rv.num_wts = 3;

  rv.links_per_value = -1;
  rv.num_links = 0;
  rv.max_links = 0;
  rv.resize_increment = 1024;

  rv.links.option = false;
  rv.links.max_links = 0;
  rv.links.num_blks = 0;
  rv.links.num_links = nullptr;
  rv.links.src_add = nullptr;
  rv.links.dst_add = nullptr;
  rv.links.w_index = nullptr;
}

void
remapVarsEnsureSize(RemapVars &rv, size_t size)
{
  if (size >= rv.max_links)
    {
      while (size >= rv.max_links) rv.max_links += rv.resize_increment;

      rv.src_cell_add.resize(rv.max_links);
      rv.tgt_cell_add.resize(rv.max_links);
      rv.wts.resize(rv.num_wts * rv.max_links);
    }
}

void
remapVarsResize(RemapVars &rv, size_t size)
{
  rv.max_links = size;

  rv.src_cell_add.resize(rv.max_links);
  rv.tgt_cell_add.resize(rv.max_links);
  rv.wts.resize(rv.num_wts * rv.max_links);
}

void
remapVarsReorder(RemapVars &rv)
{
  size_t j, nval = 0, num_blks = 0;
  size_t n;

  size_t num_links = rv.num_links;

  printf("remapVarsReorder\n");
  printf("  num_links %zu\n", num_links);
  rv.links.option = true;

  size_t lastval = -1;
  size_t max_links = 0;
  for (n = 0; n < num_links; n++)
    {
      if (rv.tgt_cell_add[n] == lastval)
        nval++;
      else
        {
          if (nval > num_blks) num_blks = nval;
          nval = 1;
          max_links++;
          lastval = rv.tgt_cell_add[n];
        }
    }

  if (num_blks)
    {
      rv.links.max_links = max_links;
      rv.links.num_blks = num_blks;

      printf("num_links %zu  max_links %zu  num_blks %zu\n", rv.num_links, max_links, num_blks);

      rv.links.num_links = (size_t *) Malloc(num_blks * sizeof(size_t));
      rv.links.dst_add = (size_t **) Malloc(num_blks * sizeof(size_t *));
      rv.links.src_add = (size_t **) Malloc(num_blks * sizeof(size_t *));
      rv.links.w_index = (size_t **) Malloc(num_blks * sizeof(size_t *));
    }

  for (j = 0; j < num_blks; j++)
    {
      rv.links.dst_add[j] = (size_t *) Malloc(max_links * sizeof(size_t));
      rv.links.src_add[j] = (size_t *) Malloc(max_links * sizeof(size_t));
      rv.links.w_index[j] = (size_t *) Malloc(max_links * sizeof(size_t));
    }

  for (j = 0; j < num_blks; j++)
    {
      nval = 0;
      lastval = -1;
      size_t nlinks = 0;

      for (n = 0; n < num_links; n++)
        {
          if (rv.tgt_cell_add[n] == lastval)
            nval++;
          else
            {
              nval = 1;
              lastval = rv.tgt_cell_add[n];
            }

          if (nval == j + 1)
            {
              rv.links.dst_add[j][nlinks] = rv.tgt_cell_add[n];
              rv.links.src_add[j][nlinks] = rv.src_cell_add[n];
              rv.links.w_index[j][nlinks] = n;
              nlinks++;
            }
        }

      rv.links.num_links[j] = nlinks;
      printf("loop %zu  nlinks %zu\n", j + 1, nlinks);
    }
}

void
remapVarsFree(RemapVars &rv)
{
  if (rv.pinit)
    {
      rv.pinit = false;
      rv.sort_add = false;

      varrayFree(rv.src_cell_add);
      varrayFree(rv.tgt_cell_add);
      varrayFree(rv.wts);

      if (rv.links.option)
        {
          rv.links.option = false;

          if (rv.links.num_blks)
            {
              Free(rv.links.num_links);
              size_t num_blks = rv.links.num_blks;
              for (size_t i = 0; i < num_blks; ++i)
                {
                  Free(rv.links.src_add[i]);
                  Free(rv.links.dst_add[i]);
                  Free(rv.links.w_index[i]);
                }
              Free(rv.links.src_add);
              Free(rv.links.dst_add);
              Free(rv.links.w_index);
            }
        }
    }
  else
    fprintf(stderr, "%s Warning: vars not initialized!\n", __func__);

} /* remapVarsFree */

void
remapVarsCheckWeights(const RemapVars &rv)
{
  const auto num_links = rv.num_links;
  const auto num_wts = rv.num_wts;
  const auto normOpt = rv.normOpt;
  const auto &src_add = rv.src_cell_add;
  const auto &tgt_add = rv.tgt_cell_add;
  const auto &wts = rv.wts;

  for (size_t n = 0; n < num_links; ++n)
    {
      if (wts[n * num_wts] < -0.01)
        cdoPrint("Map weight < 0! grid1idx=%zu grid2idx=%zu nlink=%zu wts=%g", src_add[n], tgt_add[n], n, wts[n * num_wts]);

      if (normOpt != NormOpt::NONE && wts[n * num_wts] > 1.01)
        cdoPrint("Map weight > 1! grid1idx=%zu grid2idx=%zu nlink=%zu wts=%g", src_add[n], tgt_add[n], n, wts[n * num_wts]);
    }
}
