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

      Runstat    runrange        Running range
      Runstat    runmin          Running minimum
      Runstat    runmax          Running maximum
      Runstat    runsum          Running sum
      Runstat    runmean         Running mean
      Runstat    runavg          Running average
      Runstat    runvar          Running variance
      Runstat    runvar1         Running variance [Normalize by (n-1)]
      Runstat    runstd          Running standard deviation
      Runstat    runstd1         Running standard deviation [Normalize by (n-1)]
*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include "datetime.h"
#include "field_functions.h"

static void
addOperators(void)
{
  // clang-format off
  cdo_operator_add("runrange", FieldFunc_Range, 0, nullptr);
  cdo_operator_add("runmin",   FieldFunc_Min,   0, nullptr);
  cdo_operator_add("runmax",   FieldFunc_Max,   0, nullptr);
  cdo_operator_add("runsum",   FieldFunc_Sum,   0, nullptr);
  cdo_operator_add("runmean",  FieldFunc_Mean,  0, nullptr);
  cdo_operator_add("runavg",   FieldFunc_Avg,   0, nullptr);
  cdo_operator_add("runvar",   FieldFunc_Var,   0, nullptr);
  cdo_operator_add("runvar1",  FieldFunc_Var1,  0, nullptr);
  cdo_operator_add("runstd",   FieldFunc_Std,   0, nullptr);
  cdo_operator_add("runstd1",  FieldFunc_Std1,  0, nullptr);
  // clang-format on
}

class ModuleRunstat
{
  TimeStat timestatDate{ TimeStat::MEAN };
  bool runstat_nomiss = false;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID1;
  int taxisID2;

  DateTimeList dtlist;
  VarList varList;
  std::vector<RecordInfo> recList;
  FieldVector3D vars1;
  FieldVector3D vars2;
  FieldVector3D samp1;
  std::vector<bool> imask;

  int ndates;
  bool lrange;
  bool lvarstd;
  bool lvars2;
  bool lmean;
  bool lstd;
  int operfunc;
  int maxrecs;
  int divisor;

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

    auto envstr = getenv("RUNSTAT_NOMISS");
    if (envstr)
      {
        char *endptr;
        auto envval = (int) strtol(envstr, &endptr, 10);
        if (envval == 1) runstat_nomiss = true;
      }

    addOperators();

    auto operatorID = cdo_operator_id();
    operfunc = cdo_operator_f1(operatorID);  // used in omp loop

    auto lminmax = (operfunc == FieldFunc_Min || operfunc == FieldFunc_Max);
    lrange = (operfunc == FieldFunc_Range);
    lmean = (operfunc == FieldFunc_Mean || operfunc == FieldFunc_Avg);
    lstd = (operfunc == FieldFunc_Std || operfunc == FieldFunc_Std1);
    lvarstd = (lstd || operfunc == FieldFunc_Var || operfunc == FieldFunc_Var1);
    lvars2 = (lvarstd || lrange);
    divisor = (operfunc == FieldFunc_Std1 || operfunc == FieldFunc_Var1);

    operator_input_arg("number of timesteps");
    operator_check_argc(1);
    ndates = parameter_to_int(cdo_operator_argv(0));

    streamID1 = cdo_open_read(0);

    auto vlistID1 = cdo_stream_inq_vlist(streamID1);
    auto vlistID2 = vlistDuplicate(vlistID1);

    if (!lminmax) vlist_unpack(vlistID2);

    taxisID1 = vlistInqTaxis(vlistID1);
    taxisID2 = taxisDuplicate(taxisID1);
    taxisWithBounds(taxisID2);
    vlistDefTaxis(vlistID2, taxisID2);
    // Number of timestep will be reduced compared to the input error handling in case of not enough timesteps is done per record
    auto nsteps = vlistNtsteps(vlistID1);
    if (nsteps != -1)
      {
        nsteps -= ndates - 1;
        if (nsteps > 0) vlistDefNtsteps(vlistID2, nsteps);
      }

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

    maxrecs = vlistNrecs(vlistID1);

    recList = std::vector<RecordInfo>(maxrecs);

    dtlist.set_stat(timestatDate);
    dtlist.set_calendar(taxisInqCalendar(taxisID1));

    varList_init(varList, vlistID1);

    int VARS_MEMTYPE = 0;
    if ((operfunc == FieldFunc_Min) || (operfunc == FieldFunc_Max)) VARS_MEMTYPE = FIELD_NAT;

    vars1 = FieldVector3D(ndates + 1);
    if (!runstat_nomiss) samp1.resize(ndates + 1);
    if (lvars2) vars2.resize(ndates + 1);

    for (int its = 0; its < ndates; its++)
      {
        if (!runstat_nomiss) fields_from_vlist(vlistID1, samp1[its], FIELD_VEC);
        fields_from_vlist(vlistID1, vars1[its], FIELD_VEC | VARS_MEMTYPE);
        if (lvars2) fields_from_vlist(vlistID1, vars2[its], FIELD_VEC);
      }

    auto gridsizemax = vlistGridsizeMax(vlistID1);
    imask = std::vector<bool>(gridsizemax);
  }
  void
  run()
  {

    auto field2_stdvar_func = lstd ? field2_std : field2_var;
    auto fieldc_stdvar_func = lstd ? fieldc_std : fieldc_var;

    int tsID = 0;
    int otsID = 0;
    int numSteps = 0;
    while (true)
      {
      FILL_FIRST_NDATES:
        auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
        if (nrecs == 0)
          {
            if (tsID < ndates)
              cdo_abort("File has less then %d timesteps!", ndates);
            else
              break;
          }

        numSteps = (tsID < ndates) ? tsID : ndates - 1;

        dtlist.taxis_inq_timestep(taxisID1, numSteps);

        for (int recID = 0; recID < nrecs; ++recID)
          {
            int varID, levelID;
            cdo_inq_record(streamID1, &varID, &levelID);

            if (tsID == 0) recList[recID].set(varID, levelID);

            auto &rvars1 = vars1[numSteps][varID][levelID];

            auto fieldsize = rvars1.size;  // used in omp loop

            cdo_read_record(streamID1, rvars1);
            if (lrange)
              {
                vars2[numSteps][varID][levelID].nmiss = rvars1.nmiss;
                vars2[numSteps][varID][levelID].vec_d = rvars1.vec_d;
              }

            if (runstat_nomiss && rvars1.nmiss) cdo_abort("Missing values supported was swichted off by env. RUNSTAT_NOMISS!");

            if (!runstat_nomiss)
              {
                auto missval = rvars1.missval;

                if (rvars1.memType == MemType::Float)
                  for (size_t i = 0; i < fieldsize; ++i) imask[i] = !DBL_IS_EQUAL(rvars1.vec_f[i], missval);
                else
                  for (size_t i = 0; i < fieldsize; ++i) imask[i] = !DBL_IS_EQUAL(rvars1.vec_d[i], missval);

                for (size_t i = 0; i < fieldsize; ++i) samp1[numSteps][varID][levelID].vec_d[i] = (double) imask[i];

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
                for (int inp = 0; inp < numSteps; ++inp)
                  {
                    auto &samp = samp1[inp][varID][levelID].vec_d;
                    for (size_t i = 0; i < fieldsize; ++i)
                      if (imask[i]) samp[i]++;
                  }
              }

            if (lvarstd)
              {
                field2_moq(vars2[numSteps][varID][levelID], vars1[numSteps][varID][levelID]);
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
                for (int inp = 0; inp < numSteps; ++inp)
                  {
                    field2_sumsumq(vars1[inp][varID][levelID], vars2[inp][varID][levelID], rvars1);
                  }
              }
            else if (lrange)
              {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
                for (int inp = 0; inp < numSteps; ++inp)
                  {
                    field2_maxmin(vars1[inp][varID][levelID], vars2[inp][varID][levelID], rvars1);
                  }
              }
            else
              {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
                for (int inp = 0; inp < numSteps; ++inp) { field2_function(vars1[inp][varID][levelID], rvars1, operfunc); }
              }
          }

        tsID++;  // don't move this line

        if (tsID < ndates) goto FILL_FIRST_NDATES;

        for (int recID = 0; recID < maxrecs; ++recID)
          {
            auto [varID, levelID] = recList[recID].get();
            if (varList[varID].isConstant) continue;

            const auto &rsamp1 = samp1[0][varID][levelID];
            auto &rvars1 = vars1[0][varID][levelID];
            auto numSets = ndates;

            if (lmean)
              {
                if (!runstat_nomiss)
                  field2_div(rvars1, rsamp1);
                else
                  fieldc_div(rvars1, (double) numSets);
              }
            else if (lvarstd)
              {
                if (!runstat_nomiss)
                  field2_stdvar_func(rvars1, vars2[0][varID][levelID], rsamp1, divisor);
                else
                  fieldc_stdvar_func(rvars1, vars2[0][varID][levelID], numSets, divisor);
              }
            else if (lrange) { field2_sub(rvars1, vars2[0][varID][levelID]); }
          }

        dtlist.stat_taxis_def_timestep(taxisID2, ndates);
        cdo_def_timestep(streamID2, otsID);

        for (int recID = 0; recID < maxrecs; ++recID)
          {
            auto [varID, levelID] = recList[recID].get();
            if (otsID && varList[varID].isConstant) continue;

            auto &rvars1 = vars1[0][varID][levelID];

            cdo_def_record(streamID2, varID, levelID);
            cdo_write_record(streamID2, rvars1);
          }

        otsID++;

        dtlist.shift();

        vars1[ndates] = vars1[0];
        if (!runstat_nomiss) samp1[ndates] = samp1[0];
        if (lvars2) vars2[ndates] = vars2[0];

        for (int inp = 0; inp < ndates; ++inp)
          {
            vars1[inp] = vars1[inp + 1];
            if (!runstat_nomiss) samp1[inp] = samp1[inp + 1];
            if (lvars2) vars2[inp] = vars2[inp + 1];
          }
      }
  }

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

    cdo_finish();
  }
};

void *
Runstat(void *process)
{
  ModuleRunstat runstat;
  runstat.init(process);
  runstat.run();
  runstat.close();

  return nullptr;
}
