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

      Ensstat    ensrange        Ensemble range
      Ensstat    ensmin          Ensemble minimum
      Ensstat    ensmax          Ensemble maximum
      Ensstat    enssum          Ensemble sum
      Ensstat    ensmean         Ensemble mean
      Ensstat    ensavg          Ensemble average
      Ensstat    ensstd          Ensemble standard deviation
      Ensstat    ensstd1         Ensemble standard deviation
      Ensstat    ensvar          Ensemble variance
      Ensstat    ensvar1         Ensemble variance
      Ensstat    enspctl         Ensemble percentiles
*/

#include <atomic>

#include <cdi.h>

#include "cdo_rlimit.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "param_conversion.h"
#include "cdo_task.h"
#include "cdo_options.h"
#include "util_files.h"
#include "cimdOmp.h"
#include "field_functions.h"

struct ens_file_t
{
  CdoStreamID streamID;
  int vlistID;
  size_t nmiss[2];
  double missval[2];
  Varray<double> array[2];
};

struct ensstat_arg_t
{
  int t;
  int varID[2];
  int levelID[2];
  int vlistID1;
  CdoStreamID streamID2;
  int nfiles;
  ens_file_t *efData;
  double *array2Data;
  double *count2Data;
  Field *fieldsData;
  int operfunc;
  double pn;
  bool lpctl;
  bool withCountData;
  int nvars;
};

static void *
ensstat_func(void *ensarg)
{
  if (Options::CDO_task) cdo_omp_set_num_threads(Threading::ompNumThreads);

  auto arg = (ensstat_arg_t *) ensarg;
  auto t = arg->t;
  auto nfiles = arg->nfiles;
  auto ef = arg->efData;
  auto fields = arg->fieldsData;
  auto array2 = arg->array2Data;
  auto count2 = arg->count2Data;
  auto withCountData = arg->withCountData;

  auto hasMissvals = false;
  for (int fileID = 0; fileID < nfiles; ++fileID)
    if (ef[fileID].nmiss[t] > 0) hasMissvals = true;

  const auto gridID = vlistInqVarGrid(arg->vlistID1, arg->varID[t]);
  auto gridsize = gridInqSize(gridID);
  auto missval = vlistInqVarMissval(arg->vlistID1, arg->varID[t]);

  std::atomic<size_t> atomicNumMiss{0};
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t i = 0; i < gridsize; ++i)
    {
      const auto ompthID = cdo_omp_get_thread_num();

      auto &field = fields[ompthID];
      field.missval = missval;
      field.nmiss = 0;
      for (int fileID = 0; fileID < nfiles; ++fileID)
        {
          field.vec_d[fileID] = ef[fileID].array[t][i];
          if (hasMissvals && DBL_IS_EQUAL(field.vec_d[fileID], ef[fileID].missval[t]))
            {
              field.vec_d[fileID] = missval;
              field.nmiss++;
            }
        }

      array2[i] = arg->lpctl ? field_pctl(field, arg->pn) : field_function(field, arg->operfunc);

      if (DBL_IS_EQUAL(array2[i], field.missval)) atomicNumMiss++;

      if (withCountData) count2[i] = nfiles - field.nmiss;
    }

  size_t nmiss = atomicNumMiss;

  cdo_def_record(arg->streamID2, arg->varID[t], arg->levelID[t]);
  cdo_write_record(arg->streamID2, array2, nmiss);

  if (withCountData)
    {
      cdo_def_record(arg->streamID2, arg->varID[t] + arg->nvars, arg->levelID[t]);
      cdo_write_record(arg->streamID2, count2, 0);
    }

  return nullptr;
}

static void
addOperators(void)
{
  // clang-format off
  cdo_operator_add("ensrange",  FieldFunc_Range,  0, nullptr);
  cdo_operator_add("ensmin",    FieldFunc_Min,    0, nullptr);
  cdo_operator_add("ensmax",    FieldFunc_Max,    0, nullptr);
  cdo_operator_add("enssum",    FieldFunc_Sum,    0, nullptr);
  cdo_operator_add("ensmean",   FieldFunc_Mean,   0, nullptr);
  cdo_operator_add("ensavg",    FieldFunc_Avg,    0, nullptr);
  cdo_operator_add("ensstd",    FieldFunc_Std,    0, nullptr);
  cdo_operator_add("ensstd1",   FieldFunc_Std1,   0, nullptr);
  cdo_operator_add("ensvar",    FieldFunc_Var,    0, nullptr);
  cdo_operator_add("ensvar1",   FieldFunc_Var1,   0, nullptr);
  cdo_operator_add("ensskew",   FieldFunc_Skew,   0, nullptr);
  cdo_operator_add("enskurt",   FieldFunc_Kurt,   0, nullptr);
  cdo_operator_add("ensmedian", FieldFunc_Median, 0, nullptr);
  cdo_operator_add("enspctl",   FieldFunc_Pctl,   0, nullptr);
  // clang-format on
}

void *
Ensstat(void *process)
{
  cdo::Task *task = Options::CDO_task ? new cdo::Task : nullptr;
  int nrecs0;

  cdo_initialize(process);

  addOperators();

  const auto operatorID = cdo_operator_id();
  const auto operfunc = cdo_operator_f1(operatorID);

  const auto lpctl = (operfunc == FieldFunc_Pctl);

  auto argc = cdo_operator_argc();
  const auto nargc = argc;

  double pn = 0.0;
  if (operfunc == FieldFunc_Pctl)
    {
      operator_input_arg("percentile number");
      pn = parameter_to_double(cdo_operator_argv(0));
      argc--;
    }

  auto withCountData = false;
  if (argc == 1)
    {
      if (cdo_operator_argv(nargc - 1) == "count")
        withCountData = true;
      else
        cdo_abort("Unknown parameter: >%s<", cdo_operator_argv(nargc - 1));
    }

  auto nfiles = cdo_stream_cnt() - 1;

  if (Options::cdoVerbose) cdo_print("Ensemble over %d files.", nfiles);

  cdo::set_numfiles(nfiles + 8);

  const auto ofilename = cdo_get_stream_name(nfiles);

  if (!Options::cdoOverwriteMode && FileUtils::file_exists(ofilename) && !FileUtils::user_file_overwrite(ofilename))
    cdo_abort("Outputfile %s already exists!", ofilename);

  std::vector<ens_file_t> ef(nfiles);

  FieldVector fields(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; ++i) fields[i].resize(nfiles);

  for (int fileID = 0; fileID < nfiles; ++fileID)
    {
      ef[fileID].streamID = cdo_open_read(fileID);
      ef[fileID].vlistID = cdo_stream_inq_vlist(ef[fileID].streamID);
    }

  // check that the contents is always the same
  for (int fileID = 1; fileID < nfiles; ++fileID) vlist_compare(ef[0].vlistID, ef[fileID].vlistID, CMP_ALL);

  const auto vlistID1 = ef[0].vlistID;
  const auto vlistID2 = vlistDuplicate(vlistID1);
  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

  const auto gridsizemax = vlistGridsizeMax(vlistID1);

  for (int fileID = 0; fileID < nfiles; ++fileID)
    {
      ef[fileID].array[0].resize(gridsizemax);
      if (Options::CDO_task) ef[fileID].array[1].resize(gridsizemax);
    }

  Varray<double> array2(gridsizemax);

  const auto nvars = vlistNvars(vlistID2);
  Varray<double> count2;
  if (withCountData)
    {
      count2.resize(gridsizemax);
      for (int varID = 0; varID < nvars; ++varID)
        {
          char name[CDI_MAX_NAME];
          vlistInqVarName(vlistID2, varID, name);
          strcat(name, "_count");
          const auto gridID = vlistInqVarGrid(vlistID2, varID);
          const auto zaxisID = vlistInqVarZaxis(vlistID2, varID);
          const auto timetype = vlistInqVarTimetype(vlistID2, varID);
          const auto cvarID = vlistDefVar(vlistID2, gridID, zaxisID, timetype);
          cdiDefKeyString(vlistID2, cvarID, CDI_KEY_NAME, name);
          vlistDefVarDatatype(vlistID2, cvarID, CDI_DATATYPE_INT16);
          if (cvarID != (varID + nvars)) cdo_abort("Internal error, varIDs do not match!");
        }
    }

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

  ensstat_arg_t ensstat_arg;
  ensstat_arg.vlistID1 = vlistID1;
  ensstat_arg.streamID2 = streamID2;
  ensstat_arg.nfiles = nfiles;
  ensstat_arg.array2Data = array2.data();
  ensstat_arg.count2Data = count2.data();
  ensstat_arg.fieldsData = fields.data();
  ensstat_arg.operfunc = operfunc;
  ensstat_arg.pn = pn;
  ensstat_arg.lpctl = lpctl;
  ensstat_arg.withCountData = withCountData;
  ensstat_arg.nvars = nvars;
  ensstat_arg.t = 0;

  auto printWarning = false;
  auto printError = false;
  int t = 0;
  int tsID = 0;
  do
    {
      nrecs0 = cdo_stream_inq_timestep(ef[0].streamID, tsID);
      for (int fileID = 1; fileID < nfiles; ++fileID)
        {
          const auto streamID = ef[fileID].streamID;
          const auto nrecs = cdo_stream_inq_timestep(streamID, tsID);
          if (nrecs != nrecs0)
            {
              if (nrecs == 0)
                {
                  printWarning = true;
                  cdo_warning("Inconsistent ensemble file, too few time steps in %s!", cdo_get_stream_name(fileID));
                }
              else if (nrecs0 == 0)
                {
                  printWarning = true;
                  cdo_warning("Inconsistent ensemble file, too few time steps in %s!", cdo_get_stream_name(0));
                }
              else
                {
                  printError = true;
                  cdo_warning("Inconsistent ensemble file, number of records at time step %d of %s and %s differ!", tsID + 1,
                              cdo_get_stream_name(0), cdo_get_stream_name(fileID));
                }
              goto CLEANUP;
            }
        }

      if (nrecs0 > 0)
        {
          cdo_taxis_copy_timestep(taxisID2, taxisID1);
          cdo_def_timestep(streamID2, tsID);
        }

      for (int recID = 0; recID < nrecs0; ++recID)
        {
          int varID = -1, levelID = -1;

          for (int fileID = 0; fileID < nfiles; ++fileID)
            {
              cdo_inq_record(ef[fileID].streamID, &varID, &levelID);
              ef[fileID].missval[t] = vlistInqVarMissval(ef[fileID].vlistID, varID);
            }
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
          for (int fileID = 0; fileID < nfiles; ++fileID)
            {
              cdo_read_record(ef[fileID].streamID, ef[fileID].array[t].data(), &ef[fileID].nmiss[t]);
            }

          ensstat_arg.efData = ef.data();
          ensstat_arg.varID[t] = varID;
          ensstat_arg.levelID[t] = levelID;
          if (Options::CDO_task)
            {
              task->start(ensstat_func, &ensstat_arg);
              task->wait();
              // t = !t;
            }
          else
            {
              ensstat_func(&ensstat_arg);
            }
        }

      tsID++;
    }
  while (nrecs0 > 0);

CLEANUP:

  if (printWarning) cdo_warning("Inconsistent ensemble, processed only the first %d timesteps!", tsID);
  if (printError) cdo_abort("Inconsistent ensemble, processed only the first %d timesteps!", tsID);

  for (int fileID = 0; fileID < nfiles; ++fileID) cdo_stream_close(ef[fileID].streamID);

  cdo_stream_close(streamID2);

  if (task) delete task;

  cdo_finish();

  return nullptr;
}
