/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Uwe Schulzweida

*/

#include "cdo_timer.h"
#include "cdo_options.h"
#include "cdo_output.h"
#include "remap.h"
#include "remap_store_link.h"
#include "progress.h"
#include "cdo_omp.h"

// ------------------------------------------------------------------------
// This routine computes the weights for a k-nearest-neighbor interpolation
// ------------------------------------------------------------------------
void
remap_knn_weights(KnnParams const &knnParams, RemapSearch &rsearch, RemapVars &rv)
{
  if (Options::cdoVerbose) cdo_print("Called %s()", __func__);

  cdo::Progress progress;

  // Compute mappings from source to target grid

  auto *tgtGrid = rsearch.tgtGrid;
  auto tgtGridSize = rsearch.tgtGrid->size;

  bool kIsOne = (knnParams.k == 1);
  std::vector<WeightLinks> weightLinksList;
  std::vector<WeightLinks1> weightLinks1List;
  kIsOne ? weightLinks1List.resize(tgtGridSize) : weightLinksList.resize(tgtGridSize);

  cdo::timer timer;

  std::vector<long> numLinksPerValuePool(Threading::ompNumMaxThreads, -1);
  bool progressActive = (tgtGridSize > progressMinSize);
  size_t count{ 0 };
  double rcountMax{ static_cast<double>(Threading::ompNumMaxThreads) / tgtGridSize };

#ifdef _OPENMP
#pragma omp parallel default(shared)
#endif
  {
    KnnData knnData(knnParams);
    auto tid = (Threading::ompNumMaxThreads == 1) ? 0 : cdo_omp_get_thread_num();
    // Loop over target grid
#ifdef _OPENMP
#pragma omp for schedule(static)
#endif
    for (size_t cellIndex = 0; cellIndex < tgtGridSize; ++cellIndex)
    {
      if (tid == 0 && progressActive) progress.update(++count * rcountMax);

      if (!tgtGrid->mask[cellIndex]) continue;

      auto pointLL = tgtGrid->get_lonlat(cellIndex);

      // Find nearest grid points on source grid and distances to each point
      remap_search_points(rsearch, pointLL, knnData);

      // Compute weights if mask is false, eliminate those points
      auto numWeights = knnData.compute_weights(rsearch.srcGrid->mask);
      if (numWeights) { tgtGrid->cellFrac[cellIndex] = 1.0; }

      // Store the link
      if (numWeights)
        kIsOne ? store_weightLinks(knnData.m_indices[0], knnData.m_dist[0], weightLinks1List[cellIndex])
               : store_weightLinks(numWeights, knnData.m_indices, knnData.m_dist, weightLinksList[cellIndex]);

      if (!kIsOne && numWeights > 0) set_numLinksPerValue(numLinksPerValuePool[tid], static_cast<long>(numWeights));
    }
  }

  kIsOne ? weightLinks_to_remapLinks(tgtGridSize, weightLinks1List, rv)
         : weightLinks_to_remapLinks(tgtGridSize, weightLinksList, rv);

  check_numLinksPerValue(numLinksPerValuePool);
  if (kIsOne) { rv.numLinksPerValue = 1; }
  else if (numLinksPerValuePool[0] > 0) { rv.numLinksPerValue = numLinksPerValuePool[0]; }

  if (Options::cdoVerbose) cdo_print("Point search nearest: %.2f seconds", timer.elapsed());
}  // remap_knn_weights

// ------------------------------------------------------------------------------
// This routine computes and apply weights for a k-nearest-neighbor interpolation
// ------------------------------------------------------------------------------
template <typename T1, typename T2>
  requires std::is_arithmetic_v<T1> && std::is_arithmetic_v<T2>
static void
remap_knn(Varray<T1> const &srcArray, Varray<T2> &tgtArray, double srcMissval, size_t numMissVals, KnnParams const &knnParams,
          RemapSearch &rsearch)
{
  T1 missval = static_cast<T1>(srcMissval);

  if (Options::cdoVerbose) cdo_print("Called %s()", __func__);

  cdo::Progress progress;

  // Compute mappings from source to target grid

  auto *tgtGrid = rsearch.tgtGrid;
  auto tgtGridSize = rsearch.tgtGrid->size;
  auto srcGridSize = rsearch.srcGrid->size;

  auto hasTgtGridMask = !rsearch.tgtGrid->mask.empty();

  Vmask srcGridMask;
  if (numMissVals) remap_set_mask(srcArray, srcGridSize, numMissVals, srcMissval, srcGridMask);

  cdo::timer timer;

  bool progressActive = (tgtGridSize > progressMinSize);
  size_t count{ 0 };
  double rcountMax{ static_cast<double>(Threading::ompNumMaxThreads) / tgtGridSize };

#ifdef _OPENMP
#pragma omp parallel default(shared)
#endif
  {
    KnnData knnData(knnParams);
    auto tid = (Threading::ompNumMaxThreads == 1) ? 0 : cdo_omp_get_thread_num();
    // Loop over target grid
#ifdef _OPENMP
#pragma omp for schedule(static)
#endif
    for (size_t cellIndex = 0; cellIndex < tgtGridSize; ++cellIndex)
    {
      if (tid == 0 && progressActive) progress.update(++count * rcountMax);

      auto &tgtValue = tgtArray[cellIndex];
      tgtValue = missval;

      if (hasTgtGridMask && !tgtGrid->mask[cellIndex]) continue;

      auto pointLL = tgtGrid->get_lonlat(cellIndex);

      // Find nearest grid points on source grid and distances to each point
      remap_search_points(rsearch, pointLL, knnData);

      // Compute weights if mask is false, eliminate those points
      auto numWeights = numMissVals ? knnData.compute_weights(srcGridMask) : knnData.compute_weights();
      if (numWeights) { tgtValue = knnData.array_weights_sum(srcArray); }
    }
  }
  if (Options::cdoVerbose) cdo_print("Point search nearest: %.2f seconds", timer.elapsed());
}  // remap_knn

void
remap_knn(KnnParams const &knnParams, RemapSearch &remapSearch, Field const &field1, Field &field2)
{
  auto func = [&](auto const &v1, auto &v2) { remap_knn(v1, v2, field1.missval, field1.numMissVals, knnParams, remapSearch); };
  field_operation2(func, field1, field2);
}

template <typename T1, typename T2>
  requires std::is_arithmetic_v<T1> && std::is_arithmetic_v<T2>
void
intgrid_knn(Varray<T1> const &srcArray, Varray<T2> &tgtArray, int gridID1, int gridID2, double srcMissval, size_t numMissVals,
            KnnParams const &knnParams)
{
  auto mapType = RemapMethod::KNN;

  if (Options::cdoVerbose) cdo_print("Called %s()", __func__);

  RemapType remap;
  remap_set_option(RemapOption::GenerateWeights, 0);
  remap_init_grids(mapType, knnParams.extrapolate, gridID1, remap.srcGrid, gridID2, remap.tgtGrid);
  remap_search_init(mapType, remap.search, remap.srcGrid, remap.tgtGrid);

  remap_knn(srcArray, tgtArray, srcMissval, numMissVals, knnParams, remap.search);

  remap_grid_free(remap.srcGrid);
  remap_grid_free(remap.tgtGrid);
}  // intgrid_knn

void
intgrid_knn(KnnParams const &knnParams, Field const &field1, Field &field2)
{
  auto func = [&](auto const &v1, auto &v2)
  { intgrid_knn(v1, v2, field1.grid, field2.grid, field1.missval, field1.numMissVals, knnParams); };
  field_operation2(func, field1, field2);

  field2.numMissVals = field_num_mv(field2);
}

void
intgrid_1nn(Field const &field1, Field &field2)
{
  KnnParams knnParams;
  knnParams.k = 1;
  intgrid_knn(knnParams, field1, field2);
}
