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

      Vertstat   vertrange       Vertical range
      Vertstat   vertmin         Vertical minimum
      Vertstat   vertmax         Vertical maximum
      Vertstat   vertsum         Vertical sum
      Vertstat   vertint         Vertical integral
      Vertstat   vertmean        Vertical mean
      Vertstat   vertavg         Vertical average
      Vertstat   vertvar         Vertical variance
      Vertstat   vertvar1        Vertical variance [Normalize by (n-1)]
      Vertstat   vertstd         Vertical standard deviation
      Vertstat   vertstd1        Vertical standard deviation [Normalize by (n-1)]
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "cdo_zaxis.h"
#include "param_conversion.h"
#include "pmlist.h"
#include "cdi_lockedIO.h"

#define IS_SURFACE_LEVEL(zaxisID) (zaxisInqType(zaxisID) == ZAXIS_SURFACE && zaxisInqSize(zaxisID) == 1)

int
get_surface_ID(int vlistID)
{
  int surfID = -1;

  const auto nzaxis = vlistNzaxis(vlistID);
  for (int index = 0; index < nzaxis; ++index)
    {
      const auto zaxisID = vlistZaxis(vlistID, index);
      if (IS_SURFACE_LEVEL(zaxisID))
        {
          surfID = vlistZaxis(vlistID, index);
          break;
        }
    }

  if (surfID == -1) surfID = zaxis_from_name("surface");

  return surfID;
}

static void
setSurfaceID(const int vlistID, const int surfID)
{
  const auto nzaxis = vlistNzaxis(vlistID);
  for (int index = 0; index < nzaxis; ++index)
    {
      const auto zaxisID = vlistZaxis(vlistID, index);
      if (zaxisID != surfID || !IS_SURFACE_LEVEL(zaxisID)) vlistChangeZaxisIndex(vlistID, index, surfID);
    }
}

static void
vertstat_get_parameter(bool &weights, bool &genbounds)
{
  const auto pargc = cdo_operator_argc();
  if (pargc)
    {
      const auto pargv = cdo_get_oper_argv();

      KVList kvlist;
      kvlist.name = "VERTSTAT";
      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 == "weights")   weights = parameter_to_bool(value);
          else if (key == "genbounds") genbounds = parameter_to_bool(value);
          else cdo_abort("Invalid parameter key >%s<!", key.c_str());
          // clang-format on
        }
    }
}

void *
Vertstat(void *process)
{
  struct VertInfo
  {
    int zaxisID = -1;
    int status = -1;
    int numlevel = 0;
    Varray<double> thickness;
    Varray<double> weights;
  };

  cdo_initialize(process);

  // clang-format off
                 cdo_operator_add("vertrange", FieldFunc_Range, 0, nullptr);
                 cdo_operator_add("vertmin",   FieldFunc_Min,   0, nullptr);
                 cdo_operator_add("vertmax",   FieldFunc_Max,   0, nullptr);
                 cdo_operator_add("vertsum",   FieldFunc_Sum,   0, nullptr);
  int VERTINT  = cdo_operator_add("vertint",   FieldFunc_Sum,   1, nullptr);
                 cdo_operator_add("vertmean",  FieldFunc_Mean,  1, nullptr);
                 cdo_operator_add("vertavg",   FieldFunc_Avg,   1, nullptr);
                 cdo_operator_add("vertvar",   FieldFunc_Var,   1, nullptr);
                 cdo_operator_add("vertvar1",  FieldFunc_Var1,  1, nullptr);
                 cdo_operator_add("vertstd",   FieldFunc_Std,   1, nullptr);
                 cdo_operator_add("vertstd1",  FieldFunc_Std1,  1, nullptr);
  // clang-format on

  const auto operatorID = cdo_operator_id();
  const auto operfunc = cdo_operator_f1(operatorID);
  const bool needWeights = cdo_operator_f2(operatorID);

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

  // int applyWeights = lmean;

  const auto streamID1 = cdo_open_read(0);
  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);

  vlistClearFlag(vlistID1);
  const auto nvars = vlistNvars(vlistID1);
  for (int varID = 0; varID < nvars; varID++) vlistDefFlag(vlistID1, varID, 0, true);

  const auto vlistID2 = vlistCreate();
  cdo_vlist_copy_flag(vlistID2, vlistID1);
  vlistDefNtsteps(vlistID2, vlistNtsteps(vlistID1));

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

  const auto surfID = get_surface_ID(vlistID1);
  setSurfaceID(vlistID2, surfID);

  const auto nzaxis = vlistNzaxis(vlistID1);
  std::vector<VertInfo> vert(nzaxis);
  if (needWeights)
    {
      auto useweights = true;
      auto genbounds = false;
      vertstat_get_parameter(useweights, genbounds);

      if (!useweights)
        {
          genbounds = false;
          cdo_print("Using constant vertical weights!");
        }

      for (int index = 0; index < nzaxis; ++index)
        {
          const auto zaxisID = vlistZaxis(vlistID1, index);
          const auto nlev = zaxisInqSize(zaxisID);
          vert[index].numlevel = 0;
          vert[index].status = 0;
          vert[index].zaxisID = zaxisID;
          // if (nlev > 1)
          {
            vert[index].numlevel = nlev;
            vert[index].thickness.resize(nlev);
            vert[index].weights.resize(nlev);
            vert[index].status = get_layer_thickness(useweights, genbounds, index, zaxisID, nlev, vert[index].thickness.data(),
                                                     vert[index].weights.data());
          }
          if (!useweights) vert[index].status = 3;
        }
    }

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

  VarList varList;
  varListInit(varList, vlistID1);

  Field field;

  FieldVector vars1(nvars), samp1(nvars), vars2;
  if (lvarstd || lrange) vars2.resize(nvars);

  for (int varID = 0; varID < nvars; varID++)
    {
      samp1[varID].grid = varList[varID].gridID;
      samp1[varID].missval = varList[varID].missval;
      samp1[varID].memType = MemType::Double;
      vars1[varID].grid = varList[varID].gridID;
      vars1[varID].missval = varList[varID].missval;
      vars1[varID].memType = MemType::Double;
      vars1[varID].resize(varList[varID].gridsize);
      if (lvarstd || lrange)
        {
          vars2[varID].grid = varList[varID].gridID;
          vars2[varID].missval = varList[varID].missval;
          vars2[varID].memType = MemType::Double;
          vars2[varID].resize(varList[varID].gridsize);
        }
    }

  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      cdo_taxis_copy_timestep(taxisID2, taxisID1);
      cdo_def_timestep(streamID2, tsID);

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

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

          rvars1.nsamp++;
          if (lrange) vars2[varID].nsamp++;

          const auto gridsize = varList[varID].gridsize;
          const auto zaxisID = varList[varID].zaxisID;
          const auto nlev = varList[varID].nlevels;

          auto layer_weight = 1.0;
          auto layer_thickness = 1.0;
          if (needWeights)
            {
              for (int index = 0; index < nzaxis; ++index)
                if (vert[index].zaxisID == zaxisID)
                  {
                    if (vert[index].status == 0 && tsID == 0 && levelID == 0 && nlev > 1)
                      {
                        cdo_warning("Layer bounds not available, using constant vertical weights for variable %s!",
                                    varList[varID].name);
                      }
                    else
                      {
                        layer_weight = vert[index].weights[levelID];
                        layer_thickness = vert[index].thickness[levelID];
                      }

                    break;
                  }
            }

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

              if (operatorID == VERTINT && IS_NOT_EQUAL(layer_thickness, 1.0)) fieldc_mul(rvars1, layer_thickness);
              if (lmean && IS_NOT_EQUAL(layer_weight, 1.0)) fieldc_mul(rvars1, layer_weight);

              if (lvarstd)
                {
                  if (IS_NOT_EQUAL(layer_weight, 1.0))
                    {
                      field2_moqw(vars2[varID], rvars1, layer_weight);
                      fieldc_mul(rvars1, layer_weight);
                    }
                  else
                    {
                      field2_moq(vars2[varID], rvars1);
                    }
                }

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

                  for (size_t i = 0; i < gridsize; i++)
                    rsamp1.vec_d[i] = (DBL_IS_EQUAL(rvars1.vec_d[i], rvars1.missval)) ? 0.0 : layer_weight;
                }
            }
          else
            {
              field.init(varList[varID]);
              cdo_read_record(streamID1, field);

              if (operatorID == VERTINT && IS_NOT_EQUAL(layer_thickness, 1.0)) fieldc_mul(field, layer_thickness);
              if (lmean && IS_NOT_EQUAL(layer_weight, 1.0)) fieldc_mul(field, layer_weight);

              if (field.nmiss || !rsamp1.empty())
                {
                  if (rsamp1.empty()) rsamp1.resize(gridsize, rvars1.nsamp);

                  if (field.memType == MemType::Float)
                    {
                      for (size_t i = 0; i < gridsize; i++)
                        if (!DBL_IS_EQUAL(field.vec_f[i], (float) rvars1.missval)) rsamp1.vec_d[i] += layer_weight;
                    }
                  else
                    {
                      for (size_t i = 0; i < gridsize; i++)
                        if (!DBL_IS_EQUAL(field.vec_d[i], rvars1.missval)) rsamp1.vec_d[i] += layer_weight;
                    }
                }

              if (lvarstd)
                {
                  if (IS_NOT_EQUAL(layer_weight, 1.0))
                    {
                      field2_sumqw(vars2[varID], field, layer_weight);
                      field2_sumw(rvars1, field, layer_weight);
                    }
                  else
                    {
                      field2_sumsumq(rvars1, vars2[varID], field);
                    }
                }
              else if (lrange)
                {
                  field2_maxmin(rvars1, vars2[varID], field);
                }
              else
                {
                  field2_function(rvars1, field, operfunc);
                }
            }
        }

      for (int varID = 0; varID < nvars; varID++)
        {
          const auto &rsamp1 = samp1[varID];
          auto &rvars1 = vars1[varID];

          if (rvars1.nsamp)
            {
              if (lmean)
                {
                  if (!rsamp1.empty())
                    field2_div(rvars1, rsamp1);
                  else
                    fieldc_div(rvars1, (double) rvars1.nsamp);
                }
              else if (lvarstd)
                {
                  if (!rsamp1.empty())
                    field2_stdvar_func(rvars1, vars2[varID], rsamp1, divisor);
                  else
                    fieldc_stdvar_func(rvars1, vars2[varID], rvars1.nsamp, divisor);
                }
              else if (lrange)
                {
                  field2_sub(rvars1, vars2[varID]);
                }

              cdo_def_record(streamID2, varID, 0);
              cdo_write_record(streamID2, rvars1);
              rvars1.nsamp = 0;
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  vlistDestroy(vlistID2);

  cdo_finish();

  return nullptr;
}
