/*
  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 <mpim_grid.h>
#include "remap.h"
#include "remap_utils.h"
#include "remap_store_link.h"
#include "progress.h"

// bicubic interpolation

static void
bicubic_set_weights(double xfrac, double yfrac, BicubicWeights &weights)
{
  auto xfrac1 = xfrac * xfrac * (xfrac - 1.0);
  auto xfrac2 = xfrac * (xfrac - 1.0) * (xfrac - 1.0);
  auto xfrac3 = xfrac * xfrac * (3.0 - 2.0 * xfrac);
  auto yfrac1 = yfrac * yfrac * (yfrac - 1.0);
  auto yfrac2 = yfrac * (yfrac - 1.0) * (yfrac - 1.0);
  auto yfrac3 = yfrac * yfrac * (3.0 - 2.0 * yfrac);
  // clang-format off
  weights[0][0] = (1.0 - yfrac3) * (1.0 - xfrac3);
  weights[1][0] = (1.0 - yfrac3) *        xfrac3;
  weights[2][0] =        yfrac3  *        xfrac3;
  weights[3][0] =        yfrac3  * (1.0 - xfrac3);
  weights[0][1] = (1.0 - yfrac3) *        xfrac2;
  weights[1][1] = (1.0 - yfrac3) *        xfrac1;
  weights[2][1] =        yfrac3  *        xfrac1;
  weights[3][1] =        yfrac3  *        xfrac2;
  weights[0][2] =        yfrac2  * (1.0 - xfrac3);
  weights[1][2] =        yfrac2  *        xfrac3;
  weights[2][2] =        yfrac1  *        xfrac3;
  weights[3][2] =        yfrac1  * (1.0 - xfrac3);
  weights[0][3] =        yfrac2  *        xfrac2;
  weights[1][3] =        yfrac2  *        xfrac1;
  weights[2][3] =        yfrac1  *        xfrac1;
  weights[3][3] =        yfrac1  *        xfrac2;
  // clang-format on
}

static void
renormalize_weights(std::array<double, 4> const &lats, BicubicWeights &weights)
{
  std::array<double, 4> aLats;
  for (int i = 0; i < 4; ++i) aLats[i] = std::abs(lats[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][0] = aLats[i] * invSumWeights;
    weights[i][1] = 0.0;
    weights[i][2] = 0.0;
    weights[i][3] = 0.0;
  }
}

static void
bicubic_sort_weights(std::array<size_t, 4> &indices, BicubicWeights &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
bicubic_warning()
{
  static std::atomic<bool> printWarning(true);
  if (Options::cdoVerbose || printWarning.exchange(false))
  {
    cdo_warning("Bicubic interpolation failed for some grid points - used a distance-weighted average instead!");
  }
}

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

  if (rsearch.srcGrid->rank != 2) cdo_abort("Can't do bicubic 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<WeightLinksBicubic> 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);

    SquareCorners squareCorners;
    BicubicWeights weights;  //  bicubic 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, srcGrid->mask);

    // If point found, find local xfrac/yfrac coordinates for weights
    if (searchResult > 0)
    {
      tgtGrid->cellFrac[cellIndex] = 1.0;

      auto [xfrac, yfrac] = remap_find_weights(pointLL, squareCorners);
      if (xfrac >= 0.0 && yfrac >= 0.0)
      {
        // Successfully found xfrac, yfrac - compute weights
        bicubic_set_weights(xfrac, yfrac, weights);
        bicubic_sort_weights(squareCorners.indices, weights);
        store_weightLinks(squareCorners.indices, weights, weightLinksList[cellIndex]);
      }
      else
      {
        bicubic_warning();
        searchResult = -1;
      }
    }

    // Search for bicubic 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(srcGrid->mask, squareCorners.indices, squareCorners.lats) > 0)
      {
        tgtGrid->cellFrac[cellIndex] = 1.0;
        renormalize_weights(squareCorners.lats, weights);
        bicubic_sort_weights(squareCorners.indices, weights);
        store_weightLinks(squareCorners.indices, weights, weightLinksList[cellIndex]);
      }
    }
  }

  weightLinks_to_remapLinks(tgtGridSize, weightLinksList, rv);

  rv.numLinksPerValue = 4;

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

// -----------------------------------------------------------------------
// This routine computes and apply the weights for a bicubic interpolation.
// -----------------------------------------------------------------------
template <typename T>
[[nodiscard]] static T
bicubic_remap(Varray<T> const &srcArray, std::array<size_t, 4> const &ind, BicubicWeights const &wgt,
              RemapGradients const &gradients)
{
  auto const &glat = gradients.lat;
  auto const &glon = gradients.lon;
  auto const &glatlon = gradients.latLon;

  double tgtPoint = 0.0;
  for (int i = 0; i < 4; ++i)
    tgtPoint += srcArray[ind[i]] * wgt[i][0] + glat[ind[i]] * wgt[i][1] + glon[ind[i]] * wgt[i][2] + glatlon[ind[i]] * wgt[i][3];

  return static_cast<T>(tgtPoint);
}

template <typename T1, typename T2>
  requires std::is_arithmetic_v<T1> && std::is_arithmetic_v<T2>
static void
remap_bicubic(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__);

  if (rsearch.srcGrid->rank != 2) cdo_abort("Can't do bicubic interpolation if the source grid is not a regular 2D grid!");

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

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

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

  // Compute mappings from source to target grid

  RemapGradients gradients(srcGrid->size);
  remap::gradients(srcArray, *srcGrid, srcGridMask, gradients);

  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);

    SquareCorners squareCorners;
    BicubicWeights weights;  //  bicubic 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
        bicubic_set_weights(xfrac, yfrac, weights);
        bicubic_sort_weights(squareCorners.indices, weights);
        tgtValue = bicubic_remap(srcArray, squareCorners.indices, weights, gradients);
      }
      else
      {
        bicubic_warning();
        searchResult = -1;
      }
    }

    // Search for bicubic 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)
      {
        renormalize_weights(squareCorners.lats, weights);
        bicubic_sort_weights(squareCorners.indices, weights);
        tgtValue = bicubic_remap(srcArray, squareCorners.indices, weights, gradients);
      }
    }
  }

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

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