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

  Author: Uwe Schulzweida

*/

#include <atomic>

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

// Interpolation using a distance-weighted average

// -----------------------------------------------------------------------
// This routine computes the weights for a k-nearest-neighbor interpolation
// -----------------------------------------------------------------------
void
remap_knn_weights(const KnnParams &knnParams, RemapSearch &rsearch, RemapVars &rv)
{
  auto srcGrid = rsearch.srcGrid;
  auto tgtGrid = rsearch.tgtGrid;

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

  cdo::Progress progress;

  // Compute mappings from source to target grid

  auto tgtGridSize = tgtGrid->size;

  std::vector<WeightLinks> weightLinks(tgtGridSize);
  weight_links_alloc(knnParams.k, tgtGridSize, weightLinks);

  std::vector<KnnData> knnDataList;
  knnDataList.reserve(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; ++i) knnDataList.push_back(KnnData(knnParams));

  cdo::timer timer;

  // Loop over target grid

  std::atomic<long> numLinksPerValue{ -1 };
  std::atomic<size_t> atomicCount{ 0 };

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t tgtCellIndex = 0; tgtCellIndex < tgtGridSize; ++tgtCellIndex)
    {
      atomicCount++;
      auto ompthID = cdo_omp_get_thread_num();
      if (ompthID == 0 && tgtGridSize > progressMinSize) progress.update((double) atomicCount / tgtGridSize);

      weightLinks[tgtCellIndex].nlinks = 0;

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

      auto &knnData = knnDataList[ompthID];
      auto pointLL = remapgrid_get_lonlat(tgtGrid, tgtCellIndex);

      // 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(srcGrid->mask);

      if (numWeights) tgtGrid->cellFrac[tgtCellIndex] = 1.0;

      // Store the link
      store_weightlinks(0, numWeights, knnData.m_indices.data(), knnData.m_dist.data(), tgtCellIndex, weightLinks);

      if (knnParams.k > 1 && numWeights > 0)
        {
          if (numLinksPerValue == -1)
            numLinksPerValue = numWeights;
          else if (numLinksPerValue > 0 && numLinksPerValue != (long) numWeights)
            numLinksPerValue = 0;
        }
    }

  weight_links_to_remap_links(0, tgtGridSize, weightLinks, rv);

  if (knnParams.k == 1)
    rv.numLinksPerValue = 1;
  else if (numLinksPerValue > 0)
    rv.numLinksPerValue = numLinksPerValue;

  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>
static void
remap_knn(const Varray<T1> &srcArray, Varray<T2> &tgtArray, double srcMissval, size_t numMissVals, const KnnParams &knnParams,
          RemapSearch &rsearch)
{
  T1 missval = srcMissval;
  auto srcGrid = rsearch.srcGrid;
  auto tgtGrid = rsearch.tgtGrid;

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

  cdo::Progress progress;

  // Compute mappings from source to target grid

  auto tgtGridSize = tgtGrid->size;
  auto srcGridSize = srcGrid->size;

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

  std::vector<KnnData> knnDataList;
  knnDataList.reserve(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; ++i) knnDataList.push_back(KnnData(knnParams));

  cdo::timer timer;

  // Loop over target grid

  std::atomic<size_t> atomicCount{ 0 };

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t tgtCellIndex = 0; tgtCellIndex < tgtGridSize; ++tgtCellIndex)
    {
      atomicCount++;
      auto ompthID = cdo_omp_get_thread_num();
      if (ompthID == 0 && tgtGridSize > progressMinSize) progress.update((double) atomicCount / tgtGridSize);

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

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

      auto &knnData = knnDataList[ompthID];
      auto pointLL = remapgrid_get_lonlat(tgtGrid, tgtCellIndex);

      // 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 = (srcGridMask.size() > 0) ? 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(const KnnParams &knnParams, RemapSearch &remapSearch, const Field &field1, Field &field2)
{
  auto func = [&](const auto &v1, auto &v2, double missval, size_t numMissVals) {
    remap_knn(v1, v2, missval, numMissVals, knnParams, remapSearch);
  };
  field_operation2(func, field1, field2, field1.missval, field1.numMissVals);
}

void
intgrid_knn(KnnParams knnParams, const Field &field1, Field &field2)
{
  auto mapType = RemapMethod::KNN;
  auto gridID1 = field1.grid;
  auto gridID2 = field2.grid;
  auto srcMissval = field1.missval;
  auto tgtMissval = field2.missval;
  const auto &srcArray = field1.vec_d;
  auto &tgtArray = field2.vec_d;

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

  cdo::Progress progress;

  // Interpolate from source to target grid

  RemapType remap;

  remap_set_int(REMAP_GENWEIGHTS, 0);
  remap_init_grids(mapType, knnParams.extrapolate, gridID1, remap.srcGrid, gridID2, remap.tgtGrid);

  auto srcGridSize = remap.srcGrid.size;
  auto tgtGridSize = remap.tgtGrid.size;

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

  std::vector<KnnData> knnDataList;
  knnDataList.reserve(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; ++i) knnDataList.push_back(KnnData(knnParams));

  remap_search_init(mapType, remap.search, remap.srcGrid, remap.tgtGrid);

  cdo::timer timer;

  // Loop over target grid

  std::atomic<size_t> numMissVals{ 0 };
  std::atomic<size_t> atomicCount{ 0 };

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t tgtCellIndex = 0; tgtCellIndex < tgtGridSize; ++tgtCellIndex)
    {
      atomicCount++;
      auto ompthID = cdo_omp_get_thread_num();
      if (ompthID == 0 && tgtGridSize > progressMinSize) progress.update((double) atomicCount / tgtGridSize);

      auto &tgtValue = tgtArray[tgtCellIndex];
      tgtValue = tgtMissval;

      // if (!tgt_mask[tgtCellIndex]) continue;

      auto &knnData = knnDataList[ompthID];
      auto pointLL = remapgrid_get_lonlat(&remap.tgtGrid, tgtCellIndex);

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

      // Compute weights if mask is false, eliminate those points
      auto numWeights = (srcGridMask.size() > 0) ? knnData.compute_weights(srcGridMask) : knnData.compute_weights();
      if (numWeights)
        tgtValue = knnData.array_weights_sum(srcArray);
      else
        numMissVals++;
    }

  field2.numMissVals = numMissVals;

  remap_grid_free(remap.srcGrid);
  remap_grid_free(remap.tgtGrid);

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

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