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

  Author: Uwe Schulzweida

*/

#include <algorithm>
#include <atomic>

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

// bilinear interpolation

static inline void
limit_dphi_bounds(double &dphi)
{
  constexpr double bound = 3.0 * PIH;
  if (dphi > bound) { dphi -= PI2; }
  else if (dphi < -bound) { dphi += PI2; }
}

[[nodiscard]] std::pair<double, double>
remap_find_weights(PointLonLat const &pointLL, SquareCorners const &squareCorners)
{
  constexpr double converge = 1.0e-10;  // Convergence criterion
  constexpr int MaxIteration = 100;     // Max iteration count for i, j iteration

  auto const &lons = squareCorners.lons;
  auto const &lats = squareCorners.lats;

  // Iterate to find xfrac,yfrac for bilinear approximation

  // some latitude  differences
  auto dth1 = lats[1] - lats[0];
  auto dth2 = lats[3] - lats[0];
  auto dth3 = lats[2] - lats[1] - dth2;

  // some longitude differences
  auto dph1 = lons[1] - lons[0];
  auto dph2 = lons[3] - lons[0];
  auto dph3 = lons[2] - lons[1];

  limit_dphi_bounds(dph1);
  limit_dphi_bounds(dph2);
  limit_dphi_bounds(dph3);
  dph3 -= dph2;

  // current guess for bilinear coordinate
  double xguess = 0.5;
  double yguess = 0.5;

  auto dlon = pointLL.lon() - lons[0];
  auto dlat = pointLL.lat() - lats[0];
  int iter = 0;  // iteration counters
  for (iter = 0; iter < MaxIteration; ++iter)
  {
    auto xy = xguess * yguess;

    auto dthp = dlat - dth1 * xguess - dth2 * yguess - dth3 * xy;
    auto dphp = dlon;
    limit_dphi_bounds(dphp);
    dphp -= dph1 * xguess + dph2 * yguess + dph3 * xy;

    auto mat1 = dth1 + dth3 * yguess;
    auto mat2 = dth2 + dth3 * xguess;
    auto mat3 = dph1 + dph3 * yguess;
    auto mat4 = dph2 + dph3 * xguess;
    auto inv_determinant = 1.0 / (mat1 * mat4 - mat2 * mat3);
    auto deli = (dthp * mat4 - dphp * mat2) * inv_determinant;
    auto delj = (dphp * mat1 - dthp * mat3) * inv_determinant;

    if (std::abs(deli) < converge && std::abs(delj) < converge) break;

    xguess += deli;
    yguess += delj;
  }

  if (iter >= MaxIteration) { xguess = yguess = -1.0; }

  return { xguess, yguess };
}

static void
bilinear_set_weights(double xfrac, double yfrac, std::array<double, 4> &weights)
{
  // clang-format off
  weights[0] = (1.0 - xfrac) * (1.0 - yfrac);
  weights[1] =        xfrac  * (1.0 - yfrac);
  weights[2] =        xfrac  *        yfrac;
  weights[3] = (1.0 - xfrac) *        yfrac;
  // clang-format on
}

static void
renormalize_weights(std::array<double, 4> const &srcLats, std::array<double, 4> &weights)
{
  std::array<double, 4> aLats;
  for (int i = 0; i < 4; ++i) aLats[i] = std::abs(srcLats[i]);
  // inverse sum of weights for normalization
  auto invSumWeights = 1.0 / (aLats[0] + aLats[1] + aLats[2] + aLats[3]);
  for (int i = 0; i < 4; ++i) { weights[i] = aLats[i] * invSumWeights; }
}

static void
bilinear_sort_weights(std::array<size_t, 4> &indices, std::array<double, 4> &weights)
{
  auto cmp_swap = [&](size_t a, size_t b)
  {
    if (indices[a] > indices[b]) [[unlikely]]
    {
      std::swap(indices[a], indices[b]);
      std::swap(weights[a], weights[b]);
    }
  };
  cmp_swap(0, 2);
  cmp_swap(1, 3);
  cmp_swap(0, 1);
  cmp_swap(2, 3);
  cmp_swap(1, 2);
}

static void
bilinear_warning()
{
  static std::atomic<bool> printWarning(true);
  if (Options::cdoVerbose || printWarning.exchange(false))
  {
    cdo_warning("Bilinear interpolation failed for some grid points - used a distance-weighted average instead!");
  }
}

static void
remap_bilinear_weights_regular(RemapSearch &rsearch, Vmask const &srcGridMask, PointLonLat const &pointLL, double &tgtCellFrac,
                               WeightLinksBilinear &weightLinks)
{
  SquareCorners squareCorners;
  std::array<double, 4> weights;  //  bilinear weights for four corners

  // Find nearest square of grid points on source grid
  auto searchResult = remap_search_square(rsearch, pointLL, squareCorners);

  // Check to see if points are mask points
  if (searchResult > 0) searchResult = remap_check_mask_indices(squareCorners.indices, srcGridMask);

  // If point found, find local xfrac/yfrac coordinates for weights
  if (searchResult > 0)
  {
    tgtCellFrac = 1.0;

    auto [xfrac, yfrac] = remap_find_weights(pointLL, squareCorners);
    if (xfrac >= 0.0 && yfrac >= 0.0)
    {
      // Successfully found xfrac, yfrac - compute weights
      bilinear_set_weights(xfrac, yfrac, weights);
      bilinear_sort_weights(squareCorners.indices, weights);
      store_weightLinks(squareCorners.indices, weights, weightLinks);
    }
    else
    {
      bilinear_warning();
      searchResult = -1;
    }
  }

  // Search for bilinear failed - use a distance-weighted average instead
  // (this is typically near the pole) Distance was stored in srcLats!
  if (searchResult < 0)
  {
    if (num_src_points(srcGridMask, squareCorners.indices, squareCorners.lats) > 0)
    {
      tgtCellFrac = 1.0;
      renormalize_weights(squareCorners.lats, weights);
      bilinear_sort_weights(squareCorners.indices, weights);
      store_weightLinks(squareCorners.indices, weights, weightLinks);
    }
  }
}

static void
remap_bilinear_weights_healpix(RemapSearch const &rsearch, Vmask const &srcGridMask, PointLonLat const &pointLL,
                               double &tgtCellFrac, WeightLinksBilinear &weightLinks)
{
  std::array<size_t, 4> indices;  // indices for the four source points
  std::array<double, 4> weights;  // bilinear weights for four corners

  hp_bilinear_interpolate_weights(rsearch.srcGrid->hpParams, pointLL.lon(), pointLL.lat(), indices, weights);

  // Check to see if points are mask points
  auto searchResult = remap_check_mask_indices(indices, srcGridMask);
  if (searchResult > 0)
  {
    tgtCellFrac = 1.0;
    bilinear_sort_weights(indices, weights);
    store_weightLinks(indices, weights, weightLinks);
  }
}

// -----------------------------------------------------------------------
// This routine computes the weights for a bilinear interpolation.
// -----------------------------------------------------------------------
void
remap_bilinear_weights(RemapSearch &rsearch, RemapVars &rv)
{
  if (Options::cdoVerbose) cdo_print("Called %s()", __func__);

  auto isHealpixGrid = (rsearch.srcGrid->type == RemapGridType::HealPix);
  if (!isHealpixGrid && rsearch.srcGrid->rank != 2)
    cdo_abort("Can't do bilinear interpolation if the source grid is not a regular 2D grid!");

  cdo::timer timer;
  cdo::Progress progress;

  // Compute mappings from source to target grid

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

  std::vector<WeightLinksBilinear> weightLinksList(tgtGridSize);

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

  // Loop over target grid

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

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

    auto pointLL = tgtGrid->get_lonlat(cellIndex);

    auto &tgtCellFrac = tgtGrid->cellFrac[cellIndex];
    if (isHealpixGrid)
      remap_bilinear_weights_healpix(rsearch, srcGrid->mask, pointLL, tgtCellFrac, weightLinksList[cellIndex]);
    else
      remap_bilinear_weights_regular(rsearch, srcGrid->mask, pointLL, tgtCellFrac, weightLinksList[cellIndex]);
  }

  weightLinks_to_remapLinks(tgtGridSize, weightLinksList, rv);

  rv.numLinksPerValue = 4;

  if (Options::cdoVerbose) cdo_print("%s: %.2f seconds", __func__, timer.elapsed());
}  // remap_bilinear_weights

template <typename T>
[[nodiscard]] static inline T
bilinear_remap(Varray<T> const &srcArray, std::array<size_t, 4> const &ind, std::array<double, 4> const &wgt)
{
  return srcArray[ind[0]] * wgt[0] + srcArray[ind[1]] * wgt[1] + srcArray[ind[2]] * wgt[2] + srcArray[ind[3]] * wgt[3];
}

template <typename T1, typename T2>
  requires std::is_arithmetic_v<T1> && std::is_arithmetic_v<T2>
static void
remap_bilinear_regular(RemapSearch &rsearch, Varray<T1> const &srcArray, Vmask const &srcGridMask, PointLonLat const &pointLL,
                       T2 &tgtValue)
{
  SquareCorners squareCorners;
  std::array<double, 4> weights;  //  bilinear weights for four corners

  // Find nearest square of grid points on source grid
  auto searchResult = remap_search_square(rsearch, pointLL, squareCorners);

  // Check to see if points are mask points
  if (searchResult > 0) searchResult = remap_check_mask_indices(squareCorners.indices, srcGridMask);

  // If point found, find local xfrac/yfrac coordinates for weights
  if (searchResult > 0)
  {
    auto [xfrac, yfrac] = remap_find_weights(pointLL, squareCorners);
    if (xfrac >= 0.0 && yfrac >= 0.0)
    {
      // Successfully found xfrac, yfrac - compute weights
      bilinear_set_weights(xfrac, yfrac, weights);
      bilinear_sort_weights(squareCorners.indices, weights);
      tgtValue = bilinear_remap(srcArray, squareCorners.indices, weights);
    }
    else
    {
      bilinear_warning();
      searchResult = -1;
    }
  }

  // Search for bilinear failed - use a distance-weighted average instead
  // (this is typically near the pole) Distance was stored in srcLats!
  if (searchResult < 0)
  {
    if (srcGridMask.size() == 0 || num_src_points(srcGridMask, squareCorners.indices, squareCorners.lats) > 0)
    {
      renormalize_weights(squareCorners.lats, weights);
      bilinear_sort_weights(squareCorners.indices, weights);
      tgtValue = bilinear_remap(srcArray, squareCorners.indices, weights);
    }
  }
}

template <typename T1, typename T2>
  requires std::is_arithmetic_v<T1> && std::is_arithmetic_v<T2>
static void
remap_bilinear_healpix(RemapSearch const &rsearch, Varray<T1> const &srcArray, Vmask const &srcGridMask, PointLonLat const &pointLL,
                       T2 &tgtValue)
{
  std::array<size_t, 4> indices;  // indices for the four source points
  std::array<double, 4> weights;  // bilinear weights for four corners

  hp_bilinear_interpolate_weights(rsearch.srcGrid->hpParams, pointLL.lon(), pointLL.lat(), indices, weights);

  // Check to see if points are mask points
  auto searchResult = remap_check_mask_indices(indices, srcGridMask);
  if (searchResult > 0)
  {
    bilinear_sort_weights(indices, weights);
    tgtValue = bilinear_remap(srcArray, indices, weights);
  }
}

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

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

  auto isHealpixGrid = (rsearch.srcGrid->type == RemapGridType::HealPix);
  if (!isHealpixGrid && rsearch.srcGrid->rank != 2)
    cdo_abort("Can't do bilinear interpolation if the source grid is not a regular 2D grid!");

  cdo::timer timer;
  cdo::Progress progress;

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

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

  // Compute mappings from source to target grid

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

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

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

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

    auto pointLL = tgtGrid->get_lonlat(cellIndex);

    if (isHealpixGrid)
      remap_bilinear_healpix(rsearch, srcArray, srcGridMask, pointLL, tgtValue);
    else
      remap_bilinear_regular(rsearch, srcArray, srcGridMask, pointLL, tgtValue);
  }

  if (Options::cdoVerbose) cdo_print("%s: %.2f seconds", __func__, timer.elapsed());
}  // remap_bilinear

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