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

      Regres      regres           Regression
*/

#include <cdi.h>

#include "process_int.h"
#include "cdo_vlist.h"
#include "cdo_options.h"
#include "datetime.h"
#include "pmlist.h"
#include "param_conversion.h"

// Same code as Trend!

static void
regresGetParameter(bool &tstepIsEqual)
{
  const auto pargc = cdo_operator_argc();
  if (pargc)
    {
      auto pargv = cdo_get_oper_argv();

      KVList kvlist;
      kvlist.name = "TREND";
      if (kvlist.parse_arguments(pargc, 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.c_str());
          if (kv.nvalues < 1) cdo_abort("Missing value for parameter key >%s<!", key.c_str());
          const auto &value = kv.values[0];

          // clang-format off
          if      (key == "equal")  tstepIsEqual = parameter_to_bool(value);
          else cdo_abort("Invalid parameter key >%s<!", key.c_str());
          // clang-format on
        }
    }
}

void *
Regres(void *process)
{
  cdo_initialize(process);

  auto tstepIsEqual = true;
  regresGetParameter(tstepIsEqual);

  const auto streamID1 = cdo_open_read(0);

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

  vlistDefNtsteps(vlistID2, 1);

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

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

  VarList varList;
  varListInit(varList, vlistID1);

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

  const auto gridsizemax = vlistGridsizeMax(vlistID1);

  Field field1, field2;
  field1.resize(gridsizemax);
  field2.resize(gridsizemax);

  constexpr size_t nwork = 5;
  FieldVector2D work[nwork];
  for (auto &w : work) fields_from_vlist(vlistID1, w, FIELD_VEC, 0);

  const auto calendar = taxisInqCalendar(taxisID1);

  CheckTimeInc checkTimeInc;
  JulianDate juldate0;
  double deltat1 = 0;
  int64_t vdate = 0;
  int vtime = 0;
  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      vdate = taxisInqVdate(taxisID1);
      vtime = taxisInqVtime(taxisID1);

      if (tstepIsEqual) check_time_increment(tsID, calendar, vdate, vtime, checkTimeInc);
      const auto zj = tstepIsEqual ? (double) tsID : delta_time_step_0(tsID, calendar, vdate, vtime, juldate0, deltat1);

      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);
            }

          size_t nmiss;
          cdo_read_record(streamID1, field1.vec_d.data(), &nmiss);

          const auto gridsize = varList[varID].gridsize;
          const auto missval = varList[varID].missval;

          auto &sumj = work[0][varID][levelID].vec_d;
          auto &sumjj = work[1][varID][levelID].vec_d;
          auto &sumjx = work[2][varID][levelID].vec_d;
          auto &sumx = work[3][varID][levelID].vec_d;
          auto &zn = work[4][varID][levelID].vec_d;

          for (size_t i = 0; i < gridsize; i++)
            if (!DBL_IS_EQUAL(field1.vec_d[i], missval))
              {
                sumj[i] += zj;
                sumjj[i] += zj * zj;
                sumjx[i] += zj * field1.vec_d[i];
                sumx[i] += field1.vec_d[i];
                zn[i]++;
              }
        }

      tsID++;
    }

  taxisDefVdate(taxisID2, vdate);
  taxisDefVtime(taxisID2, vtime);
  // cdo_def_timestep(streamID2, 0);
  cdo_def_timestep(streamID3, 0);

  for (int recID = 0; recID < maxrecs; recID++)
    {
      const auto varID = recList[recID].varID;
      const auto levelID = recList[recID].levelID;
      const auto gridsize = varList[varID].gridsize;
      const auto missval = varList[varID].missval;
      const auto missval1 = missval;
      const auto missval2 = missval;
      field1.size = gridsize;
      field1.missval = missval;
      field2.size = gridsize;
      field2.missval = missval;

      const auto &sumj = work[0][varID][levelID].vec_d;
      const auto &sumjj = work[1][varID][levelID].vec_d;
      const auto &sumjx = work[2][varID][levelID].vec_d;
      const auto &sumx = work[3][varID][levelID].vec_d;
      const auto &zn = work[4][varID][levelID].vec_d;

      for (size_t i = 0; i < gridsize; i++)
        {
          const auto temp1 = SUBMN(sumjx[i], DIVMN(MULMN(sumj[i], sumx[i]), zn[i]));
          const auto temp2 = SUBMN(sumjj[i], DIVMN(MULMN(sumj[i], sumj[i]), zn[i]));

          field2.vec_d[i] = DIVMN(temp1, temp2);
          field1.vec_d[i] = SUBMN(DIVMN(sumx[i], zn[i]), MULMN(DIVMN(sumj[i], zn[i]), field2.vec_d[i]));
        }

      cdo_def_record(streamID3, varID, levelID);
      cdo_write_record(streamID3, field2.vec_d.data(), field_num_miss(field2));
    }

  cdo_stream_close(streamID3);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
