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

  Copyright (C) 2003-2021 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.
*/

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

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
}

void *
Runstat(void *process)
{
  const TimeStat timestat_date = TimeStat::MEAN;
  bool runstat_nomiss = false;

  cdo_initialize(process);

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

  addOperators();

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

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

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

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

  const auto streamID1 = cdo_open_read(0);

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

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto 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);
    }

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

  const auto maxrecs = vlistNrecs(vlistID1);
  std::vector<RecordInfo> recList(maxrecs);

  DateTimeList dtlist;
  dtlist.set_stat(timestat_date);
  dtlist.set_calendar(taxisInqCalendar(taxisID1));

  VarList varList;
  varListInit(varList, vlistID1);

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

  FieldVector3D vars1(ndates + 1), vars2, samp1;
  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);
    }

  const auto gridsizemax = vlistGridsizeMax(vlistID1);
  std::vector<bool> imask(gridsizemax);

  int tsID = 0;
  int otsID = 0;
  int numSteps = 0;
  while (true)
    {
    FILL_FIRST_NDATES:
      const 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].varID = varID;
              recList[recID].levelID = levelID;
              recList[recID].lconst = (varList[varID].timetype == TIME_CONSTANT);
            }

          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)
            {
              const 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(none) shared(numSteps, imask, samp1, varID, levelID, fieldsize)
#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(none) shared(numSteps, vars1, vars2, varID, levelID, rvars1)
#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(none) shared(numSteps, vars1, vars2, varID, levelID, rvars1)
#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(none) shared(numSteps, vars1, varID, levelID, rvars1, operfunc)
#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++)
        {
          if (recList[recID].lconst) continue;

          const auto varID = recList[recID].varID;
          const auto levelID = recList[recID].levelID;
          const auto &rsamp1 = samp1[0][varID][levelID];
          auto &rvars1 = vars1[0][varID][levelID];
          const auto nsets = ndates;

          if (lmean)
            {
              if (!runstat_nomiss)
                field2_div(rvars1, rsamp1);
              else
                fieldc_div(rvars1, (double) nsets);
            }
          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], nsets, 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++)
        {
          if (otsID && recList[recID].lconst) continue;

          const auto varID = recList[recID].varID;
          const auto levelID = recList[recID].levelID;
          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];
        }
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
