/*
  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:

*/

#include <cdi.h>

#include "process_int.h"
#include "cdo_math.h"
#include "cdo_options.h"
#include "grid_healpix.h"
#include "param_conversion.h"
#include "pmlist.h"
#include "mpim_grid.h"

enum class Stat
{
  Mean = 1,
  Avg = 2
};

struct HealpixParams
{
  int fact = 1;
  int nsideIn = 0;
  int nsideOut = 0;
  HpOrder orderIn = HpOrder::Undef;
  HpOrder orderOut = HpOrder::Undef;
  Stat stat = Stat::Mean;
  double power = 0.0;
  bool doDegrade = true;
};

template <typename T>
static T
stat_avg_mv(const T *v, size_t n, T missval, double scale)
{
  double sum = 0.0;
  size_t nOut = 0;
  for (size_t i = 0; i < n; ++i)
    if (fp_is_not_equal(v[i], missval))
      {
        sum += v[i];
        nOut++;
      }

  return (nOut == n) ? (sum / nOut) * scale : missval;
}

template <typename T>
static T
stat_mean_mv(const T *v, size_t n, T missval, double scale)
{
  double sum = 0.0;
  size_t nOut = 0;
  for (size_t i = 0; i < n; ++i)
    if (fp_is_not_equal(v[i], missval))
      {
        sum += v[i];
        nOut++;
      }

  return (nOut > 0) ? (sum / nOut) * scale : missval;
}

template <typename T>
static T
stat_mean(const T *v, size_t n)
{
  double sum = 0.0;
  for (size_t i = 0; i < n; ++i) sum += v[i];
  return sum / n;
}

static double
get_scalefactor(const HealpixParams &params)
{
  double nsideQuot = static_cast<double>(params.nsideIn) / params.nsideOut;
  return (std::fabs(params.power) > 0.0) ? std::pow(nsideQuot, -params.power) : 1.0;
}

template <typename T1, typename T2>
static void
degrade(const Varray<T1> &v1, size_t gridsize2, Varray<T2> &v2, bool hasMissvals, double mv, const HealpixParams &params)
{
  T1 missval = mv;
  auto scale = get_scalefactor(params);
  size_t nvals = params.fact * params.fact;
  if (hasMissvals)
    {
      if (params.stat == Stat::Mean)
        for (size_t i = 0; i < gridsize2; ++i) v2[i] = stat_mean_mv(&v1[i * nvals], nvals, missval, scale);
      else
        for (size_t i = 0; i < gridsize2; ++i) v2[i] = stat_avg_mv(&v1[i * nvals], nvals, missval, scale);
    }
  else
    {
      for (size_t i = 0; i < gridsize2; ++i) v2[i] = stat_mean(&v1[i * nvals], nvals) * scale;
    }
}

static void
hp_degrade(const Field &field1, Field &field2, const HealpixParams &params)
{
  auto hasMissvals = (field1.numMissVals > 0);
  auto func = [&](const auto &v1, auto &v2, auto size, double mv) { degrade(v1, size, v2, hasMissvals, mv, params); };
  field_operation2(func, field1, field2, field2.gridsize, field1.missval);
  if (hasMissvals) field_num_mv(field2);
}

template <typename T1, typename T2>
static void
upgrade(size_t gridsize1, const Varray<T1> &v1, Varray<T2> &v2, bool hasMissvals, double mv, const HealpixParams &params)
{
  T1 missval = mv;
  auto scale = get_scalefactor(params);
  size_t nvals = params.fact * params.fact;
  if (hasMissvals)
    {
      for (size_t i = 0; i < gridsize1; ++i)
        for (size_t k = 0; k < nvals; ++k) v2[i * nvals + k] = fp_is_equal(v1[i], missval) ? missval : v1[i] * scale;
    }
  else
    {
      for (size_t i = 0; i < gridsize1; ++i)
        for (size_t k = 0; k < nvals; ++k) v2[i * nvals + k] = v1[i] * scale;
    }
}

static void
hp_upgrade(const Field &field1, Field &field2, const HealpixParams &params)
{
  auto hasMissvals = (field1.numMissVals > 0);
  auto func = [&](const auto &v1, auto &v2, auto size, double mv) { upgrade(size, v1, v2, hasMissvals, mv, params); };
  field_operation2(func, field1, field2, field1.gridsize, field1.missval);
  if (hasMissvals) field_num_mv(field2);
}

template <typename T>
static void
ring_to_nested(int nside, size_t gridsize, Varray<T> &v)
{
  Varray<T> vtmp = v;
  hp_ring_to_nested(nside, gridsize, vtmp.data(), v.data());
}

static void
ring_to_nested(Field &field, int nside)
{
  auto func = [&](auto &v, auto gridsize) { ring_to_nested(nside, gridsize, v); };
  field_operation(func, field, field.gridsize);
}

template <typename T>
static void
nested_to_ring(int nside, size_t gridsize, Varray<T> &v)
{
  Varray<T> vtmp = v;
  hp_nested_to_ring(nside, gridsize, vtmp.data(), v.data());
}

static void
nested_to_ring(Field &field, int nside)
{
  auto func = [&](auto &v, auto gridsize) { nested_to_ring(nside, gridsize, v); };
  field_operation(func, field, field.gridsize);
}

static Stat
set_stat(const std::string &statString)
{
  if (statString == "mean")
    return Stat::Mean;
  else if (statString == "avg")
    return Stat::Avg;
  else
    cdo_abort("Parameter value stat=%s unsupported!", statString);

  return Stat::Mean;
}

static HealpixParams
get_parameter(void)
{
  HealpixParams params;

  auto pargc = cdo_operator_argc();
  if (pargc)
    {
      const auto &pargv = cdo_get_oper_argv();

      KVList kvlist;
      kvlist.name = cdo_module_name();
      if (kvlist.parse_arguments(pargv) != 0) cdo_abort("Parse error!");
      if (Options::cdoVerbose) kvlist.print();

      for (const auto &kv : kvlist)
        {
          const auto &key = kv.key;
          if (kv.nvalues > 1) cdo_abort("Too many values for parameter key >%s<!", key);
          if (kv.nvalues < 1) cdo_abort("Missing value for parameter key >%s<!", key);
          const auto &value = kv.values[0];

          // clang-format off
          if      (key == "nside") params.nsideOut = parameter_to_int(value);
          else if (key == "order") params.orderOut = hp_get_order(parameter_to_word(value));
          else if (key == "fact")  params.fact     = parameter_to_int(value);
          else if (key == "stat")  params.stat     = set_stat(parameter_to_word(value));
          else if (key == "power") params.power    = parameter_to_double(value);
          else cdo_abort("Invalid parameter key >%s<!", key);
          // clang-format on
        }
    }

  return params;
}

static void
verify_parameter(const HealpixParams &params)
{
  if (params.fact > 1 && params.nsideOut > 0) cdo_abort("Parameter 'fact' can't be combined with 'nside'!");
}

static int
define_healpix_grid(size_t gridsize, int nside, HpOrder order)
{
  auto orderString = (order == HpOrder::Ring) ? "ring" : "nested";
  auto projection = "healpix";
  auto gridID = gridCreate(GRID_PROJECTION, gridsize);
  cdiDefKeyString(gridID, CDI_GLOBAL, CDI_KEY_DIMNAME, "cells");
  cdiDefKeyString(gridID, CDI_GLOBAL, CDI_KEY_GRIDMAP_VARNAME, projection);
  cdiDefKeyString(gridID, CDI_GLOBAL, CDI_KEY_GRIDMAP_NAME, projection);
  cdiDefAttTxt(gridID, CDI_GLOBAL, "grid_mapping_name", (int) strlen(projection), projection);
  cdiDefAttInt(gridID, CDI_GLOBAL, "healpix_nside", CDI_DATATYPE_INT32, 1, &nside);
  cdiDefAttTxt(gridID, CDI_GLOBAL, "healpix_order", (int) strlen(orderString), orderString);

  return gridID;
}

static int
hp_define_grid(int gridID1, HealpixParams &params)
{
  int gridIDout = -1;

  auto hpParams = cdo::get_healpix_params(gridID1);
  auto nside = hpParams.get_nside();
  auto order = hpParams.get_order();
  params.nsideIn = nside;
  params.orderIn = order;

  if (!cdo::is_power_of_two(params.nsideIn)) cdo_abort("Input healpix: nside must be a power of two!");

  if (params.nsideOut == 0)
    {
      auto fact = params.fact;
      params.nsideOut = (fact > 1) ? (params.doDegrade ? nside / fact : nside * fact) : nside;
    }
  else
    {
      if (params.doDegrade)
        {
          if (params.nsideOut > params.nsideIn) cdo_abort("Parameter nside must be less than input nside=%d!", params.nsideIn);
          params.fact = params.nsideIn / params.nsideOut;
        }
      else
        {
          if (params.nsideOut < params.nsideIn) cdo_abort("Parameter nside must be greater than input nside=%d!", params.nsideIn);
          params.fact = params.nsideOut / params.nsideIn;
        }
    }

  if (!cdo::is_power_of_two(params.nsideOut)) cdo_abort("Parameter nside must be a power of two!");

  if (params.orderOut == HpOrder::Undef) params.orderOut = params.orderIn;

  size_t gridsize = 12 * params.nsideOut * params.nsideOut;
  gridIDout = define_healpix_grid(gridsize, params.nsideOut, params.orderOut);

  return gridIDout;
}

class Healpix : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Healpix",
    .operators = { { "hpupgrade", HealpixHelp }, { "hpdegrade", HealpixHelp } },
    .aliases = {},
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Healpix> registration = RegisterEntry<Healpix>(module);

  int HPDEGRADE;
  CdoStreamID streamID1;
  CdoStreamID streamID2;
  int taxisID1;
  int taxisID2;
  int vlistID2;

  bool doDegrade;

  Field field1;
  Field field2;

  HealpixParams params;

  VarList varList1;
  VarList varList2;

public:
  void
  init() override
  {
    HPDEGRADE = module.get_id("hpdegrade");

    auto operatorID = cdo_operator_id();
    doDegrade = (operatorID == HPDEGRADE);

    params = get_parameter();
    params.doDegrade = doDegrade;
    verify_parameter(params);

    streamID1 = cdo_open_read(0);
    auto vlistID1 = cdo_stream_inq_vlist(streamID1);
    taxisID1 = vlistInqTaxis(vlistID1);

    auto ngrids = vlistNumGrids(vlistID1);
    if (ngrids > 1) cdo_abort("Too many different grids!");

    auto gridID = vlistGrid(vlistID1, 0);
    if (!is_healpix_grid(gridID)) cdo_abort("Input grid is not healpix!");

    vlistID2 = vlistDuplicate(vlistID1);
    taxisID2 = taxisDuplicate(taxisID1);
    vlistDefTaxis(vlistID2, taxisID2);

    auto gridID2 = hp_define_grid(gridID, params);
    for (int index = 0; index < ngrids; ++index) vlistChangeGridIndex(vlistID2, index, gridID2);

    varList1 = VarList(vlistID1);
    varList2 = VarList(vlistID2);

    streamID2 = cdo_open_write(1);
    cdo_def_vlist(streamID2, vlistID2);
  }

  void
  run() override
  {
    int tsID1 = 0;
    while (true)
      {
        auto numFields = cdo_stream_inq_timestep(streamID1, tsID1);
        if (numFields == 0) break;

        cdo_taxis_copy_timestep(taxisID2, taxisID1);
        cdo_def_timestep(streamID2, tsID1);

        for (int fieldID = 0; fieldID < numFields; ++fieldID)
          {
            auto [varID, levelID] = cdo_inq_field(streamID1);
            field1.init(varList1.vars[varID]);
            cdo_read_field(streamID1, field1);

            field2.init(varList2.vars[varID]);

            if (params.orderIn == HpOrder::Ring) ring_to_nested(field1, params.nsideIn);

            doDegrade ? hp_degrade(field1, field2, params) : hp_upgrade(field1, field2, params);

            if (params.orderOut == HpOrder::Ring) nested_to_ring(field2, params.nsideOut);

            cdo_def_field(streamID2, varID, levelID);
            cdo_write_field(streamID2, field2);
          }

        tsID1++;
      }
  }

  void
  close() override
  {
    cdo_stream_close(streamID1);
    cdo_stream_close(streamID2);

    vlistDestroy(vlistID2);
  }
};
