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

/*
   This module contains the following operators:

      Seasstat   seasrange       Seasonal range
      Seasstat   seasmin         Seasonal minimum
      Seasstat   seasmax         Seasonal maximum
      Seasstat   seassum         Seasonal sum
      Seasstat   seasmean        Seasonal mean
      Seasstat   seasavg         Seasonal average
      Seasstat   seasvar         Seasonal variance
      Seasstat   seasvar1        Seasonal variance [Normalize by (n-1)]
      Seasstat   seasstd         Seasonal standard deviation
      Seasstat   seasstd1        Seasonal standard deviation [Normalize by (n-1)]
*/

#include <cdi.h>

#include "cdo_options.h"
#include "functs.h"
#include "process_int.h"
#include "datetime.h"
#include "printinfo.h"
#include "cdo_season.h"

static void
addOperators(void)
{
  // clang-format off
  cdoOperatorAdd("seasrange", func_range, 0, nullptr);
  cdoOperatorAdd("seasmin",   func_min,   0, nullptr);
  cdoOperatorAdd("seasmax",   func_max,   0, nullptr);
  cdoOperatorAdd("seassum",   func_sum,   0, nullptr);
  cdoOperatorAdd("seasmean",  func_mean,  0, nullptr);
  cdoOperatorAdd("seasavg",   func_avg,   0, nullptr);
  cdoOperatorAdd("seasvar",   func_var,   0, nullptr);
  cdoOperatorAdd("seasvar1",  func_var1,  0, nullptr);
  cdoOperatorAdd("seasstd",   func_std,   0, nullptr);
  cdoOperatorAdd("seasstd1",  func_std1,  0, nullptr);
  // clang-format on
}

void *
Seasstat(void *process)
{
  TimeStat timestat_date = TimeStat::MEAN;
  int64_t vdate0 = 0, vdate1 = 0;
  int vtime0 = 0, vtime1 = 0;
  int nrecs;
  int seas0 = 0;
  int oldmon = 0;
  int nseason = 0;
  const char *seas_name[4];

  cdoInitialize(process);

  addOperators();

  const auto operatorID = cdoOperatorID();
  const auto operfunc = cdoOperatorF1(operatorID);

  operatorCheckArgc(0);

  const bool lrange = operfunc == func_range;
  const bool lmean = operfunc == func_mean || operfunc == func_avg;
  const bool lstd = operfunc == func_std || operfunc == func_std1;
  const bool lvarstd = operfunc == func_std || operfunc == func_var || operfunc == func_std1 || operfunc == func_var1;
  const int divisor = operfunc == func_std1 || operfunc == func_var1;
  const bool lvars2 = lvarstd || lrange;

  const auto season_start = getSeasonStart();
  getSeasonName(seas_name);

  const auto streamID1 = cdoOpenRead(0);

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

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  taxisWithBounds(taxisID2);
  if (taxisInqType(taxisID2) == TAXIS_FORECAST) taxisDefType(taxisID2, TAXIS_RELATIVE);
  vlistDefTaxis(vlistID2, taxisID2);

  const auto streamID2 = cdoOpenWrite(1);
  cdoDefVlist(streamID2, vlistID2);

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

  DateTimeList dtlist;
  dtlist.setStat(timestat_date);
  dtlist.setCalendar(taxisInqCalendar(taxisID1));

  const auto gridsizemax = vlistGridsizeMax(vlistID1);

  Field field;
  field.resize(gridsizemax);

  FieldVector2D samp1, vars1, vars2;
  fieldsFromVlist(vlistID1, samp1);
  fieldsFromVlist(vlistID1, vars1, FIELD_VEC);
  if (lvars2) fieldsFromVlist(vlistID1, vars2, FIELD_VEC);

  int tsID = 0;
  int otsID = 0;
  while (true)
    {
      long nsets = 0;
      bool newseas = false;
      while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
        {
          dtlist.taxisInqTimestep(taxisID1, nsets);
          const auto vdate = dtlist.getVdate(nsets);
          const auto vtime = dtlist.getVtime(nsets);

          const auto month = decodeMonth(vdate);
          int newmon = month;
          if (season_start == START_DEC && newmon == 12) newmon = 0;

          const auto seas = monthToSeason(month);

          if (nsets == 0)
            {
              nseason++;
              vdate0 = vdate;
              vtime0 = vtime;
              seas0 = seas;
              oldmon = newmon;
            }

          if (newmon < oldmon) newseas = true;

          if ((seas != seas0) || newseas) break;

          oldmon = newmon;

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

              if (tsID == 0)
                {
                  recList[recID].varID = varID;
                  recList[recID].levelID = levelID;
                  recList[recID].lconst = vlistInqVarTimetype(vlistID1, varID) == TIME_CONSTANT;
                }

              auto &rsamp1 = samp1[varID][levelID];
              auto &rvars1 = vars1[varID][levelID];

              const size_t fieldsize = rvars1.size;

              if (nsets == 0)
                {
                  cdoReadRecord(streamID1, rvars1.vec_d.data(), &rvars1.nmiss);
                  if (lrange)
                    {
                      vars2[varID][levelID].nmiss = rvars1.nmiss;
                      vars2[varID][levelID].vec_d = rvars1.vec_d;
                    }

                  if (rvars1.nmiss || !rsamp1.empty())
                    {
                      if (rsamp1.empty()) rsamp1.resize(fieldsize);

                      for (size_t i = 0; i < fieldsize; i++) rsamp1.vec_d[i] = !DBL_IS_EQUAL(rvars1.vec_d[i], rvars1.missval);
                    }
                }
              else
                {
                  cdoReadRecord(streamID1, field.vec_d.data(), &field.nmiss);
                  field.size = fieldsize;
                  field.grid = rvars1.grid;
                  field.missval = rvars1.missval;

                  if (field.nmiss || !rsamp1.empty())
                    {
                      if (rsamp1.empty()) rsamp1.resize(fieldsize, nsets);

                      for (size_t i = 0; i < fieldsize; i++)
                        if (!DBL_IS_EQUAL(field.vec_d[i], rvars1.missval)) rsamp1.vec_d[i]++;
                    }

                  if (lvarstd)
                    {
                      vfarsumq(vars2[varID][levelID], field);
                      vfarsum(rvars1, field);
                    }
                  else if (lrange)
                    {
                      vfarmin(vars2[varID][levelID], field);
                      vfarmax(rvars1, field);
                    }
                  else
                    {
                      vfarfun(rvars1, field, operfunc);
                    }
                }
            }

          if (nsets == 0 && lvarstd)
            for (int recID = 0; recID < maxrecs; recID++)
              {
                if (recList[recID].lconst) continue;

                const auto varID = recList[recID].varID;
                const auto levelID = recList[recID].levelID;
                vfarmoq(vars2[varID][levelID], vars1[varID][levelID]);
              }

          vdate1 = vdate;
          vtime1 = vtime;
          nsets++;
          tsID++;
        }

      if (nrecs == 0 && nsets == 0) break;

      for (int recID = 0; recID < maxrecs; recID++)
        {
          if (recList[recID].lconst) continue;

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

          if (lmean)
            {
              if (!rsamp1.empty())
                vfardiv(rvars1, rsamp1);
              else
                vfarcdiv(rvars1, (double) nsets);
            }
          else if (lvarstd)
            {
              auto &rvars2 = vars2[varID][levelID];
              if (!rsamp1.empty())
                {
                  if (lstd)
                    vfarstd(rvars1, rvars2, rsamp1, divisor);
                  else
                    vfarvar(rvars1, rvars2, rsamp1, divisor);
                }
              else
                {
                  if (lstd)
                    vfarcstd(rvars1, rvars2, nsets, divisor);
                  else
                    vfarcvar(rvars1, rvars2, nsets, divisor);
                }
            }
          else if (lrange)
            {
              Field &rvars2 = vars2[varID][levelID];
              vfarsub(rvars1, rvars2);
            }
        }

      if (Options::cdoVerbose)
        cdoPrint("season: %3d %3s  start: %s %s  end: %s %s ntimesteps: %ld", nseason, seas_name[seas0],
                 dateToString(vdate0).c_str(), timeToString(vtime0).c_str(), dateToString(vdate1).c_str(),
                 timeToString(vtime1).c_str(), nsets);

      dtlist.statTaxisDefTimestep(taxisID2, nsets);
      cdoDefTimestep(streamID2, otsID);

      if (nsets < 3)
        cdoWarning("Season %3d (%s) has only %d input time step%s!", otsID + 1, dateToString(vdate0).c_str(), nsets,
                   nsets == 1 ? "" : "s");

      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[varID][levelID];

          cdoDefRecord(streamID2, varID, levelID);
          cdoWriteRecord(streamID2, rvars1.vec_d.data(), rvars1.nmiss);
        }

      if (nrecs == 0) break;
      otsID++;
    }

  cdoStreamClose(streamID2);
  cdoStreamClose(streamID1);

  cdoFinish();

  return nullptr;
}
