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

  Author: Uwe Schulzweida

*/

/*
   This module contains the following operators:

      Remapweights  gencon          Generate first order conservative remap weights
      Remapweights  gencon2         Generate second order conservative remap weights
      Remapweights  genbil          Generate bilinear interpolation weights
      Remapweights  genbic          Generate bicubic interpolation weights
      Remapweights  gendis          Generate distance-weighted averaging weights
      Remapweights  gennn           Generate nearest neighbor weights
      Remapweights  genlaf          Generate largest area fraction weights
*/

#include <algorithm>

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include "remap_utils.h"
#include <mpim_grid.h>
#include "griddes.h"
#include "cdo_options.h"
#include "util_string.h"

static void
add_operators(void)
{
  // clang-format off
  cdo_operator_add("gencon",       GENCON,       0, nullptr);
  cdo_operator_add("genycon2test", GENYCON2,     0, nullptr);
  cdo_operator_add("genscon",      GENSCON,      0, nullptr);
  cdo_operator_add("gencon2",      GENCON2,      0, nullptr);
  cdo_operator_add("genbil",       GENBIL,       0, nullptr);
  cdo_operator_add("genbic",       GENBIC,       0, nullptr);
  cdo_operator_add("gendis",       GENDIS,       0, nullptr);
  cdo_operator_add("gennn",        GENNN,        0, nullptr);
  cdo_operator_add("genlaf",       GENLAF,       0, nullptr);
  // clang-format on
}

class ModuleRemapweights
{
  RemapSwitches remapSwitches;
  int numNeighbors = 0;

  CdoStreamID streamID1;
  int vlistID1;

  int gridID2;

  Field field1;

  int operfunc;

  RemapType remap;
  RemapMethod mapType;

  VarList varList1;
  std::vector<bool> remapGrids;

  RemapParams remapParams;

  bool useMask;
  bool extrapolateIsSet;
  bool remapExtrapolate;
  bool needGradients;

  bool remap_genweights = true;

  int remapOrder;
  NormOpt normOpt;

public:
  void
  init(void *process)
  {
    cdo_initialize(process);

    add_operators();

    auto operatorID = cdo_operator_id();
    operfunc = cdo_operator_f1(operatorID);
    auto writeRemapWeightsOnly = true;

    remap_set_int(REMAP_WRITE_REMAP, writeRemapWeightsOnly);

    remapParams = remap_get_params();
    remap_set_params(remapParams);
    extrapolateIsSet = (remapParams.extrapolate != -1);
    remapExtrapolate = extrapolateIsSet ? (bool) remapParams.extrapolate : remap_func_is_dist(operfunc);

    if (Options::cdoVerbose) cdo_print("Extrapolation %s!", remapExtrapolate ? "enabled" : "disabled");

    {
      operator_input_arg("grid description file or name");
      if (operfunc == GENDIS && cdo_operator_argc() == 2)
        {
          auto inum = parameter_to_int(cdo_operator_argv(1));
          if (inum < 1) cdo_abort("Number of nearest neighbors out of range (>0)!");
          numNeighbors = inum;
        }
      else { operator_check_argc(1); }
    }

    gridID2 = cdo_define_grid(cdo_operator_argv(0));
    if (gridInqType(gridID2) == GRID_GENERIC) cdo_abort("Unsupported target grid type (generic)!");

    streamID1 = cdo_open_read(0);

    vlistID1 = cdo_stream_inq_vlist(streamID1);

    varList_init(varList1, vlistID1);

    remapGrids = remap_set_grids(vlistID1, varList1);

    auto numRemapGrids = std::count_if(remapGrids.begin(), remapGrids.end(), [](auto flag) { return (flag == true); });
    if (numRemapGrids == 0) cdo_abort("No remappable grid found!");

    remapSwitches = remap_operfunc_to_maptype(operfunc);
    if (numNeighbors) remapSwitches.numNeighbors = numNeighbors;

    mapType = remapSwitches.mapType;
    remapOrder = remapSwitches.remapOrder;
    useMask = !(!remap_genweights
                && (mapType == RemapMethod::BILINEAR || mapType == RemapMethod::BICUBIC || mapType == RemapMethod::DISTWGT
                    || mapType == RemapMethod::CONSERV));

    remap_set_int(REMAP_GENWEIGHTS, (int) remap_genweights);

    normOpt = NormOpt(NormOpt::NONE);
    if (mapType == RemapMethod::CONSERV_SCRIP || mapType == RemapMethod::CONSERV) normOpt = remap_get_normOpt();

    needGradients = (mapType == RemapMethod::BICUBIC);
    if ((mapType == RemapMethod::CONSERV_SCRIP || mapType == RemapMethod::CONSERV) && remapOrder == 2)
      {
        if (Options::cdoVerbose) cdo_print("Second order remapping");
        needGradients = true;
      }
  }

  void
  run()
  {
    Varray<short> imask;

    int tsID = 0;
    auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);

    for (int recID = 0; recID < nrecs; ++recID)
      {
        int varID, levelID;
        cdo_inq_record(streamID1, &varID, &levelID);
        field1.init(varList1[varID]);
        cdo_read_record(streamID1, field1);
        auto nmiss1 = useMask ? field1.nmiss : 0;

        auto gridID = varList1[varID].gridID;
        auto missval = varList1[varID].missval;
        auto gridsize = varList1[varID].gridsize;

        if (remapGrids[vlistGridIndex(vlistID1, gridID)])
          {
            if (mapType != RemapMethod::CONSERV_SCRIP && mapType != RemapMethod::CONSERV && gridInqType(gridID) == GRID_GME)
              cdo_abort("Only conservative remapping is available to remap between GME grids!");

            if (gridIsCircular(gridID) && !extrapolateIsSet) remapExtrapolate = true;

            remap_set_mask(gridsize, field1, nmiss1, missval, imask);

            //  remap.srcGrid.luse_cell_area = false;
            //  remap.tgtGrid.luse_cell_area = false;

            auto numSearchBins = remapParams.numSearchBins;
            if (gridInqType(gridID) != GRID_UNSTRUCTURED && !remapParams.numSearchBinsIsSet)
              {
                numSearchBins
                    = (!remapExtrapolate && mapType == RemapMethod::DISTWGT) ? 1 : remap_gen_numbins(gridInqYsize(gridID));
              }

            remap_set_int(REMAP_NUM_SRCH_BINS, numSearchBins);

            remap.vars.normOpt = normOpt;
            remap.vars.pinit = false;

            if ((mapType == RemapMethod::BILINEAR || mapType == RemapMethod::BICUBIC)
                && (gridInqType(gridID) == GRID_GME || gridInqType(gridID) == GRID_UNSTRUCTURED))
              cdo_abort("Bilinear/bicubic interpolation doesn't support unstructured source grids!");

            // Initialize grid information for both grids
            remap_init_grids(mapType, remapExtrapolate, gridID, remap.srcGrid, gridID2, remap.tgtGrid);
            remap_search_init(mapType, remap.search, remap.srcGrid, remap.tgtGrid);

            remap.gridID = gridID;
            remap.nmiss = nmiss1;

            if (gridInqType(gridID) == GRID_GME)
              {
                for (size_t i = 0, j = 0; i < gridsize; ++i)
                  if (remap.srcGrid.vgpm[i]) imask[j++] = imask[i];
              }

            varray_copy(remap.srcGrid.size, imask, remap.srcGrid.mask);

            if (mapType == RemapMethod::CONSERV_SCRIP || mapType == RemapMethod::CONSERV)
              {
                varray_fill(remap.srcGrid.cell_area, 0.0);
                varray_fill(remap.srcGrid.cell_frac, 0.0);
                varray_fill(remap.tgtGrid.cell_area, 0.0);
              }
            varray_fill(remap.tgtGrid.cell_frac, 0.0);

            // initialize some remapping variables
            remap_vars_init(mapType, remapOrder, remap.vars);

            remap_print_info(operfunc, remap_genweights, remap.srcGrid, remap.tgtGrid, nmiss1, remapSwitches.numNeighbors);

            if (needGradients && remap.srcGrid.rank != 2 && remapOrder == 2)
              {
                cdo_abort("Second order remapping is not available for unstructured grids!");
              }

            remap_gen_weights(remapParams.sortMode, remapSwitches, remap);

            break;
          }
      }

    remap_write_data_scrip(cdo_get_stream_name(1), remapSwitches, remap.srcGrid, remap.tgtGrid, remap.vars);
  }

  void
  close()
  {

    cdo_stream_close(streamID1);

    remap_vars_free(remap.vars);
    remap_grid_free(remap.srcGrid);
    remap_grid_free(remap.tgtGrid);
    remap_search_free(remap.search);

    cdo_finish();
  }
};

void *
Remapweights(void *process)
{
  ModuleRemapweights remapweights;
  remapweights.init(process);
  remapweights.run();
  remapweights.close();

  return nullptr;
}
