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

  Copyright (C) 2003-2020 Uwe Schulzweida, <uwe.schulzweida AT mpimet.mpg.de>
  See COPYING file for copying and redistribution conditions.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; version 2 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
*/

#include "process_int.h"
#include "cdo_wtime.h"
#include <mpim_grid.h>
#include "cdo_options.h"
#include "remap.h"
#include "remap_store_link.h"
#include "progress.h"
#include "cimdOmp.h"

// bicubic interpolation

static void
bicubicSetWeights(double iw, double jw, double (&wgts)[4][4])
{
  // clang-format off
  wgts[0][0] = (1.-jw*jw*(3.-2.*jw))  * (1.-iw*iw*(3.-2.*iw));
  wgts[1][0] = (1.-jw*jw*(3.-2.*jw))  *     iw*iw*(3.-2.*iw);
  wgts[2][0] =     jw*jw*(3.-2.*jw)   *     iw*iw*(3.-2.*iw);
  wgts[3][0] =     jw*jw*(3.-2.*jw)   * (1.-iw*iw*(3.-2.*iw));
  wgts[0][1] = (1.-jw*jw*(3.-2.*jw))  *     iw*(iw-1.)*(iw-1.);
  wgts[1][1] = (1.-jw*jw*(3.-2.*jw))  *     iw*iw*(iw-1.);
  wgts[2][1] =     jw*jw*(3.-2.*jw)   *     iw*iw*(iw-1.);
  wgts[3][1] =     jw*jw*(3.-2.*jw)   *     iw*(iw-1.)*(iw-1.);
  wgts[0][2] =     jw*(jw-1.)*(jw-1.) * (1.-iw*iw*(3.-2.*iw));
  wgts[1][2] =     jw*(jw-1.)*(jw-1.) *     iw*iw*(3.-2.*iw);
  wgts[2][2] =     jw*jw*(jw-1.)      *     iw*iw*(3.-2.*iw);
  wgts[3][2] =     jw*jw*(jw-1.)      * (1.-iw*iw*(3.-2.*iw));
  wgts[0][3] =     iw*(iw-1.)*(iw-1.) *     jw*(jw-1.)*(jw-1.);
  wgts[1][3] =     iw*iw*(iw-1.)      *     jw*(jw-1.)*(jw-1.);
  wgts[2][3] =     iw*iw*(iw-1.)      *     jw*jw*(jw-1.);
  wgts[3][3] =     iw*(iw-1.)*(iw-1.) *     jw*jw*(jw-1.);
  // clang-format on
}

template <typename T>
int num_src_points(const std::vector<T> &mask, const size_t (&src_add)[4], double (&src_lats)[4]);

static void
renormalizeWeights(const double (&src_lats)[4], double (&wgts)[4][4])
{
  double sum_wgts = 0.0;  // sum of weights for normalization
  for (unsigned n = 0; n < 4; ++n) sum_wgts += std::fabs(src_lats[n]);
  for (unsigned n = 0; n < 4; ++n) wgts[n][0] = std::fabs(src_lats[n]) / sum_wgts;
  for (unsigned n = 0; n < 4; ++n) wgts[n][1] = 0.;
  for (unsigned n = 0; n < 4; ++n) wgts[n][2] = 0.;
  for (unsigned n = 0; n < 4; ++n) wgts[n][3] = 0.;
}

static void
bicubicWarning()
{
  static bool lwarn = true;

  if (Options::cdoVerbose || lwarn)
    {
      lwarn = false;
      cdoWarning("Bicubic interpolation failed for some grid points - used a distance-weighted average instead!");
    }
}

static double
bicubicRemap(const double *restrict src_array, const double (&wgts)[4][4], const size_t (&src_add)[4], const RemapGradients &gradients)
{
  const auto &glat = gradients.grad_lat;
  const auto &glon = gradients.grad_lon;
  const auto &glatlon = gradients.grad_latlon;

  double tgt_point = 0.;
  for (unsigned n = 0; n < 4; ++n)
    tgt_point += src_array[src_add[n]] * wgts[n][0] + glat[src_add[n]] * wgts[n][1] + glon[src_add[n]] * wgts[n][2]
                 + glatlon[src_add[n]] * wgts[n][3];

  return tgt_point;
}

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

  This routine computes the weights for a bicubic interpolation.

  -----------------------------------------------------------------------
*/
void
remapBicubicWeights(RemapSearch &rsearch, RemapVars &rv)
{
  auto src_grid = rsearch.srcGrid;
  auto tgt_grid = rsearch.tgtGrid;

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

  if (src_grid->rank != 2) cdoAbort("Can't do bicubic interpolation when source grid rank != 2");

  auto start = Options::cdoVerbose ? cdo_get_wtime() : 0.0;

  progress::init();

  // Compute mappings from source to target grid

  auto tgt_grid_size = tgt_grid->size;

  std::vector<WeightLinks4> weightLinks(tgt_grid_size);
  weightLinks4Alloc(tgt_grid_size, weightLinks);

  auto findex = 0.0;

  // Loop over destination grid

#ifdef _OPENMP
#pragma omp parallel for default(none) shared(findex, rsearch, weightLinks, tgt_grid_size, src_grid, tgt_grid, rv)
#endif
  for (size_t tgt_cell_add = 0; tgt_cell_add < tgt_grid_size; ++tgt_cell_add)
    {
#ifdef _OPENMP
#pragma omp atomic
#endif
      findex++;
      if (cdo_omp_get_thread_num() == 0) progress::update(0, 1, findex / tgt_grid_size);

      weightLinks[tgt_cell_add].nlinks = 0;

      if (!tgt_grid->mask[tgt_cell_add]) continue;

      const auto llpoint = remapgrid_get_lonlat(tgt_grid, tgt_cell_add);

      size_t src_add[4];   //  address for the four source points
      double src_lats[4];  //  latitudes  of four bilinear corners
      double src_lons[4];  //  longitudes of four bilinear corners
      double wgts[4][4];   //  bicubic weights for four corners

      // Find nearest square of grid points on source grid
      auto search_result = remapSearchSquare(rsearch, llpoint, src_add, src_lats, src_lons);

      // Check to see if points are mask points
      if (search_result > 0)
        {
          for (unsigned n = 0; n < 4; ++n)
            if (!src_grid->mask[src_add[n]]) search_result = 0;
        }

      // If point found, find local iw,jw coordinates for weights
      if (search_result > 0)
        {
          tgt_grid->cell_frac[tgt_cell_add] = 1.;

          double iw = 0.0, jw = 0.0;  // current guess for bilinear coordinate
          if (remapFindWeights(llpoint, src_lons, src_lats, &iw, &jw))
            {
              // Successfully found iw,jw - compute weights
              bicubicSetWeights(iw, jw, wgts);
              storeWeightLinks4(src_add, wgts, tgt_cell_add, weightLinks);
            }
          else
            {
              bicubicWarning();
              search_result = -1;
            }
        }

      // Search for bicubic failed - use a distance-weighted average instead
      // (this is typically near the pole) Distance was stored in src_lats!
      if (search_result < 0)
        {
          if (num_src_points(src_grid->mask, src_add, src_lats) > 0)
            {
              tgt_grid->cell_frac[tgt_cell_add] = 1.;
              renormalizeWeights(src_lats, wgts);
              storeWeightLinks4(src_add, wgts, tgt_cell_add, weightLinks);
            }
        }
    }

  progress::update(0, 1, 1);

  weightLinks4ToRemapLinks(tgt_grid_size, weightLinks, rv);

  if (Options::cdoVerbose) cdoPrint("Square search nearest: %.2f seconds", cdo_get_wtime() - start);
}  // remapBicubicWeights

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

  This routine computes and apply the weights for a bicubic interpolation.

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

void
remapBicubic(RemapSearch &rsearch, const double *restrict src_array, double *restrict tgt_array, double missval)
{
  auto src_grid = rsearch.srcGrid;
  auto tgt_grid = rsearch.tgtGrid;

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

  if (src_grid->rank != 2) cdoAbort("Can't do bicubic interpolation when source grid rank != 2");

  auto start = Options::cdoVerbose ? cdo_get_wtime() : 0.0;

  progress::init();

  auto tgt_grid_size = tgt_grid->size;
  auto src_grid_size = src_grid->size;

  std::vector<uint8_t> src_grid_mask(src_grid_size);
#ifdef _OPENMP
#pragma omp parallel for default(none) schedule(static) shared(src_grid_size, src_array, src_grid_mask, missval)
#endif
  for (size_t i = 0; i < src_grid_size; ++i) src_grid_mask[i] = !DBL_IS_EQUAL(src_array[i], missval);

  // Compute mappings from source to target grid

  RemapGradients gradients(src_grid->size);
  remapGradients(*src_grid, src_grid_mask, src_array, gradients);

  auto findex = 0.0;

  // Loop over destination grid

#ifdef _OPENMP
#pragma omp parallel for default(none) \
    shared(findex, rsearch, tgt_grid_size, src_grid, tgt_grid, src_array, tgt_array, missval, src_grid_mask, gradients)
#endif
  for (size_t tgt_cell_add = 0; tgt_cell_add < tgt_grid_size; ++tgt_cell_add)
    {
#ifdef _OPENMP
#pragma omp atomic
#endif
      findex++;
      if (cdo_omp_get_thread_num() == 0) progress::update(0, 1, findex / tgt_grid_size);

      tgt_array[tgt_cell_add] = missval;

      if (!tgt_grid->mask[tgt_cell_add]) continue;

      const auto llpoint = remapgrid_get_lonlat(tgt_grid, tgt_cell_add);

      size_t src_add[4];   //  address for the four source points
      double src_lats[4];  //  latitudes  of four bilinear corners
      double src_lons[4];  //  longitudes of four bilinear corners
      double wgts[4][4];   //  bicubic weights for four corners

      // Find nearest square of grid points on source grid
      auto search_result = remapSearchSquare(rsearch, llpoint, src_add, src_lats, src_lons);

      // Check to see if points are mask points
      if (search_result > 0)
        {
          for (unsigned n = 0; n < 4; ++n)
            if (!src_grid_mask[src_add[n]]) search_result = 0;
        }

      // If point found, find local iw,jw coordinates for weights
      if (search_result > 0)
        {
          tgt_grid->cell_frac[tgt_cell_add] = 1.;

          double iw = 0.0, jw = 0.0;  // current guess for bilinear coordinate
          if (remapFindWeights(llpoint, src_lons, src_lats, &iw, &jw))
            {
              // Successfully found iw,jw - compute weights
              bicubicSetWeights(iw, jw, wgts);
              sort_add_and_wgts4(4, src_add, wgts);
              tgt_array[tgt_cell_add] = bicubicRemap(src_array, wgts, src_add, gradients);
            }
          else
            {
              bicubicWarning();
              search_result = -1;
            }
        }

      // Search for bicubic failed - use a distance-weighted average instead
      // (this is typically near the pole) Distance was stored in src_lats!
      if (search_result < 0)
        {
          if (num_src_points(src_grid_mask, src_add, src_lats) > 0)
            {
              tgt_grid->cell_frac[tgt_cell_add] = 1.;
              renormalizeWeights(src_lats, wgts);
              sort_add_and_wgts4(4, src_add, wgts);
              tgt_array[tgt_cell_add] = bicubicRemap(src_array, wgts, src_add, gradients);
            }
        }
    }

  progress::update(0, 1, 1);

  if (Options::cdoVerbose) cdoPrint("Square search nearest: %.2f seconds", cdo_get_wtime() - start);
}  // remapBicubic
