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

      Consectstep  consecsum  For each timestep, the current number of
                              onsecutive timsteps is counted
      Consectstep  consects   For each period of consecutive timesteps, only
                              count its lenght + last contributing timesteps

   =============================================================================
   Created:  04/08/2010 11:58:01 AM
    Author:  Ralf Mueller (ram), ralf.mueller@mpimet.mpg.de
   Company:  Max-Planck-Institute for Meteorology
   =============================================================================
 */

#include <cdi.h>

#include "process_int.h"
#include "cdo_vlist.h"
#include "param_conversion.h"

enum
{
  CONSECSUM,
  CONSECTS
};
#define SWITCHWARN "Hit default case!This should never happen (%s).\n"

static void
selEndOfPeriod(Field &periods, const Field &history, const Field &current, int isLastTimestep)
{
  double pmissval = periods.missval;
  auto &parray = periods.vec_d;
  double hmissval = history.missval;
  const auto &harray = history.vec_d;
  double cmissval = current.missval;
  const auto &carray = current.vec_d;

  auto len = gridInqSize(periods.grid);
  if (len != gridInqSize(current.grid) || (gridInqSize(current.grid) != gridInqSize(history.grid)))
    cdoAbort("Fields have different gridsize (%s)", __func__);

  if (!isLastTimestep)
    {
      if (current.nmiss || history.nmiss)
        {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
          for (size_t i = 0; i < len; i++)
            {
              if (!DBL_IS_EQUAL(harray[i], hmissval))
                {
                  if (!DBL_IS_EQUAL(carray[i], cmissval))
                    {
                      parray[i] = (DBL_IS_EQUAL(carray[i], 0.0) && IS_NOT_EQUAL(harray[i], 0.0)) ? harray[i] : pmissval;
                    }
                  else /* DBL_IS_EQUAL(carray[i], cmissval) */
                    {
                      parray[i] = (IS_NOT_EQUAL(harray[i], 0.0)) ? harray[i] : pmissval;
                    }
                }
              else /* DBL_IS_EQUAL(harray[i], hmissval) */
                {
                  parray[i] = pmissval;
                }
            }
        }
      else
        {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
          for (size_t i = 0; i < len; i++)
            parray[i] = (DBL_IS_EQUAL(carray[i], 0.0) && IS_NOT_EQUAL(harray[i], 0.0)) ? harray[i] : pmissval;
        }
    }
  else
    {
      if (current.nmiss)
        {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
          for (size_t i = 0; i < len; i++)
            if (!DBL_IS_EQUAL(carray[i], cmissval))
              {
                parray[i] = (DBL_IS_EQUAL(carray[i], 0.0)) ? pmissval : carray[i];
              }
            else /* DBL_IS_EQUAL(carray[i], cmissval) */
              {
                parray[i] = pmissval;
              }
        }
      else
        {
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
          for (size_t i = 0; i < len; i++) parray[i] = DBL_IS_EQUAL(carray[i], 0.0) ? pmissval : carray[i];
        }
    }

  periods.nmiss = varrayNumMV(len, parray, pmissval);
}

void *
Consecstat(void *process)
{
  int64_t vdate = 0, histvdate = 0;
  int vtime = 0, histvtime = 0;
  int nrecs;
  int varID, levelID;
  double refval = 0.0;

  cdoInitialize(process);
  cdoOperatorAdd("consecsum", CONSECSUM, 0, "refval");
  cdoOperatorAdd("consects", CONSECTS, 0, nullptr);
  const int operatorID = cdoOperatorID();

  if (operatorID == CONSECSUM)
    if (operatorArgc() > 0) refval = parameter2double(cdoOperatorArgv(0));

  const auto istreamID = cdoOpenRead(0);

  const auto ivlistID = cdoStreamInqVlist(istreamID);
  const auto itaxisID = vlistInqTaxis(ivlistID);
  const auto ovlistID = vlistDuplicate(ivlistID);
  const auto otaxisID = taxisDuplicate(itaxisID);
  vlistDefTaxis(ovlistID, otaxisID);

  VarList varList;
  varListInit(varList, ovlistID);

  Field field;
  field.resize(vlistGridsizeMax(ovlistID));

  const auto nvars = vlistNvars(ivlistID);
  FieldVector2D vars, hist, periods;
  fieldsFromVlist(ivlistID, vars, FIELD_VEC, 0);
  if (operatorID == CONSECTS)
    {
      fieldsFromVlist(ivlistID, hist, FIELD_VEC);
      fieldsFromVlist(ivlistID, periods, FIELD_VEC);
    }

  for (varID = 0; varID < nvars; varID++)
    {
      cdiDefKeyString(ovlistID, varID, CDI_KEY_UNITS, "steps"); // TODO
    }

  const auto ostreamID = cdoOpenWrite(1);
  cdoDefVlist(ostreamID, ovlistID);

  int itsID = 0;
  int otsID = 0;
  while ((nrecs = cdoStreamInqTimestep(istreamID, itsID)))
    {
      vdate = taxisInqVdate(itaxisID);
      vtime = taxisInqVtime(itaxisID);
      switch (operatorID)
        {
        case CONSECSUM:
          taxisDefVdate(otaxisID, vdate);
          taxisDefVtime(otaxisID, vtime);
          cdoDefTimestep(ostreamID, otsID);
          break;
        case CONSECTS:
          if (itsID != 0)
            {
              taxisDefVdate(otaxisID, histvdate);
              taxisDefVtime(otaxisID, histvtime);
              cdoDefTimestep(ostreamID, otsID - 1);
            }
          break;
        default: printf(SWITCHWARN, __func__); break;
        }

      for (int recID = 0; recID < nrecs; recID++)
        {
          cdoInqRecord(istreamID, &varID, &levelID);
          cdoReadRecord(istreamID, field.vec_d.data(), &field.nmiss);
          field.grid = varList[varID].gridID;
          field.size = varList[varID].gridsize;
          field.missval = varList[varID].missval;

          vfarsumtr(vars[varID][levelID], field, refval);

          switch (operatorID)
            {
            case CONSECSUM:
              cdoDefRecord(ostreamID, varID, levelID);
              cdoWriteRecord(ostreamID, vars[varID][levelID].vec_d.data(), vars[varID][levelID].nmiss);
              break;
            case CONSECTS:
              if (itsID != 0)
                {
                  selEndOfPeriod(periods[varID][levelID], hist[varID][levelID], vars[varID][levelID], false);
                  cdoDefRecord(ostreamID, varID, levelID);
                  cdoWriteRecord(ostreamID, periods[varID][levelID].vec_d.data(), periods[varID][levelID].nmiss);
                }
#ifdef _OPENMP
#pragma omp parallel for default(shared) schedule(static)
              for (size_t i = 0; i < vars[varID][levelID].size; i++) hist[varID][levelID].vec_d[i] = vars[varID][levelID].vec_d[i];
#else
              hist[varID][levelID].vec_d = vars[varID][levelID].vec_d;
#endif
              break;
            default: printf(SWITCHWARN, __func__); break;
            }
        }

      histvdate = vdate;
      histvtime = vtime;
      itsID++;
      otsID++;
    }

  if (operatorID == CONSECTS) /* Save the last timestep */
    {
      taxisDefVdate(otaxisID, vdate);
      taxisDefVtime(otaxisID, vtime);
      cdoDefTimestep(ostreamID, otsID - 1);

      for (varID = 0; varID < nvars; varID++)
        {
          const auto nlevels = varList[varID].nlevels;
          for (levelID = 0; levelID < nlevels; levelID++)
            {
              selEndOfPeriod(periods[varID][levelID], hist[varID][levelID], vars[varID][levelID], true);
              cdoDefRecord(ostreamID, varID, levelID);
              cdoWriteRecord(ostreamID, periods[varID][levelID].vec_d.data(), periods[varID][levelID].nmiss);
            }
        }
    }

  cdoStreamClose(istreamID);
  cdoStreamClose(ostreamID);

  cdoFinish();

  return nullptr;
}
