/*
  This file is part of CDO. CDO is a collection of Operators to
  manipulate and analyse Climate model Data.

  Copyright (C) 2003-2021 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:

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

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

  const auto arg = (ensstat_arg_t *) ensarg;
  const auto t = arg->t;
  const auto nfiles = arg->nfiles;
  auto ef = arg->efData;
  auto fields = arg->fieldsData;
  auto array2 = arg->array2Data;
  auto count2 = arg->count2Data;
  const 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]);
  const auto gridsize = gridInqSize(gridID);
  const auto missval = vlistInqVarMissval(arg->vlistID1, arg->varID[t]);

  size_t nmiss = 0;
#ifdef HAVE_OPENMP4
#pragma omp parallel for default(shared) reduction(+ : nmiss)
#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)) nmiss++;

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

  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).c_str());
    }

  const 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 lwarning = false;
  auto lerror = 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)
                {
                  lwarning = true;
                  cdo_warning("Inconsistent ensemble file, too few time steps in %s!", cdo_get_stream_name(fileID));
                }
              else if (nrecs0 == 0)
                {
                  lwarning = true;
                  cdo_warning("Inconsistent ensemble file, too few time steps in %s!", cdo_get_stream_name(0));
                }
              else
                {
                  lerror = 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);
            }
          //#pragma omp parallel for default(none) shared(ef, t)
          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 (lwarning) cdo_warning("Inconsistent ensemble, processed only the first %d timesteps!", tsID);
  if (lerror) 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;
}
