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

  Author: Cedrick Ansorge

*/

/*
   This module contains the following operators:
   Ensstat3       ensrkhistspace   Ensemble ranked histogram averaged over time
   Ensstat3       ensrkhisttime    Ensemble ranked histogram averaged over space
   Ensstat3       ensroccurve      Ensamble Receiver Operating Characteristics
*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include "cdo_options.h"
#include "cdo_cdi_wrapper.h"
#include "util_files.h"
#include "cdo_omp.h"
#include "field_functions.h"

// Defines for rank histogram
enum TDATA_TYPE
{
  TIME,
  SPACE
};

#define time_data TIME
#define space_data SPACE

// Defines for Receiver Operating Characteristics (ROC)
#define DEBUG_ROC 0
enum CONTINGENCY_TYPE
{
  TP,  // TP - True positive  ( event     forecast and     occured)  HIT
  FP,  // FP - False positive ( event     forecast and not occured)  false ALARM
  FN,  // FN - False negative ( event not forecast and     ocurred)  MISSED
  TN   // TN - True negative  ( event not forecast and not ocurred)  CORRECT REJECTION
};

enum ROC_ENUM_TYPE
{
  TPR,  // TPR = True Positive Rate = TP / ( TP + FN )
  FPR   // FNR = False Negtive Rate = FN / ( FP + TN )
};

static double
roc_curve_integrate(const Varray2D<double> &roc, int n)
{
  double area = 0.0;

  for (int i = 1; i <= n; ++i)
    {
      auto x1 = roc[i][FPR];
      auto x0 = roc[i - 1][FPR];
      auto y1 = roc[i][TPR];
      auto y0 = roc[i - 1][TPR];
      auto dx = x1 - x0;
      auto dy = y1 - y0;

      auto step_area = -0.5 * dx * dy - dx * y0;
      area += step_area;
    }

  return area - 0.5;
}

class Ensstat3 : public Process
{
  enum
  {
    func_roc,
    func_rank
  };

public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Ensstat3",
    .operators = { { "ensrkhistspace", func_rank, space_data, Ensstat2Help },
                   { "ensrkhisttime", func_rank, time_data, Ensstat2Help },
                   { "ensroc", func_roc, 0, Ensstat2Help } },
    .aliases = {},
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { -1, 1, NoRestriction },
  };
  inline static RegisterEntry<Ensstat3> registration = RegisterEntry<Ensstat3>(module);

  int j;
  size_t numMissVals = 0;
  int chksum = 0;  // for check of histogram population
  int binID = 0;
  int have_miss = 0;
  CdoStreamID streamID2 = 0;

  struct ens_file_t
  {
    Varray<double> array;
    CdoStreamID streamID;
    int vlistID;
    VarList varList;
  };

  std::vector<ens_file_t> ef;
  int numFiles;
  int operfunc;
  int datafunc;

  int taxisID1;
  int taxisID2;

  int numBins = 0;
  int numEns;

  std::vector<int> varIDs2;

  size_t gridsizeMax;
  std::vector<std::vector<int>> array2;
  FieldVector field;

  std::vector<int> hist;
  Varray2D<double> roc;             // receiver operating characteristics table
  Varray2D<int> ctg_tab;            // contingency table and histogram
  Varray<double> uThresh, lThresh;  // thresholds for histograms

public:
  void
  init() override
  {
    auto operatorID = cdo_operator_id();
    operfunc = cdo_operator_f1(operatorID);
    datafunc = cdo_operator_f2(operatorID);

    if (operfunc == func_roc)
      {
        operator_input_arg("Number of eigen functions to write out");
        numBins = parameter_to_int(cdo_operator_argv(0));
      }

    numFiles = cdo_stream_cnt() - 1;
    numEns = numFiles - 1;

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

    std::string ofilename = cdo_get_stream_name(numFiles);

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

    ef = std::vector<ens_file_t>(numFiles);

    /* *************************************************** */
    /* should each thread be allocating memory locally???? */
    /* ("first touch strategy")                            */
    /* --> #pragma omp parallel for ...                    */
    /* *************************************************** */
    field = FieldVector(Threading::ompNumMaxThreads);
    for (int i = 0; i < Threading::ompNumMaxThreads; ++i)
      {
        field[i].resize(numFiles);
        field[i].weightv.resize(numFiles);
        for (int fileID = 0; fileID < numFiles; ++fileID) field[i].weightv[fileID] = 1;
      }

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

    // check for identical contents of all ensemble members
    for (int fileID = 1; fileID < numFiles; ++fileID) varList_compare(ef[0].varList, ef[fileID].varList);

    auto vlistID1 = ef[0].vlistID;
    const auto &varList1 = ef[0].varList;
    auto vlistID2 = vlistCreate();
    vlistDefNtsteps(vlistID2, varList1.numSteps());

    auto numVars = varList1.numVars();
    varIDs2 = std::vector<int>(numVars);

    auto zaxisID2 = zaxisCreate(ZAXIS_GENERIC, numFiles);
    {
      Varray<double> levs(numFiles, 0);
      for (int i = 0; i < numFiles; ++i) levs[i] = i;
      zaxisDefLevels(zaxisID2, levs.data());
      cdiDefKeyString(zaxisID2, CDI_GLOBAL, CDI_KEY_NAME, "histogram_binID");
    }

    int time_mode = (datafunc == TIME) ? TIME_VARYING : TIME_CONSTANT;

    for (int varID = 0; varID < numVars; ++varID)
      {
        int gridID2;
        /* ******************************************************************** */
        /* numFiles includes the observation, so there are numFiles-1 ensembles */
        /* and exactly numFiles bins, in which the observation could fall       */
        /* ******************************************************************** */
        if (datafunc == TIME)
          {
            double val = 0.0;
            gridID2 = gridCreate(GRID_LONLAT, 1);
            gridDefXsize(gridID2, 1);
            gridDefYsize(gridID2, 1);
            gridDefXvals(gridID2, &val);
            gridDefYvals(gridID2, &val);
          }
        else  // datafunc == SPACE
          gridID2 = varList1.vars[varID].gridID;

        varIDs2[varID] = vlistDefVar(vlistID2, gridID2, zaxisID2, time_mode);
      }

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

    for (int varID = 0; varID < numVars; ++varID)
      {
        if (varList1.vars[varID].nlevels > 1)
          {
            cdo_warning("More than one level not supported when processing ranked histograms.");
            cdo_warning("Try to use `cdo splitlevel` to split the dataset into levels and apply");
            cdo_warning("the operator seperately to each level.");
            cdo_abort("Exit due to unsupported file structure");
          }
      }

    if (operfunc != func_roc)
      {
        streamID2 = cdo_open_write(numFiles);
        cdo_def_vlist(streamID2, vlistID2);
      }

    gridsizeMax = cdo_vlist_gridsizemax(vlistID1);

    for (int fileID = 0; fileID < numFiles; ++fileID) ef[fileID].array.resize(gridsizeMax);

    array2 = std::vector<std::vector<int>>(numFiles + 1);
    if (operfunc == func_rank && datafunc == SPACE)
      {  // need to memorize data for entire grid before writing
        for (binID = 0; binID < numFiles; binID++) array2[binID].resize(gridsizeMax, 0);
      }
    else if (operfunc == func_rank)
      {  // can process data separately for each timestep and only need to cumulate values over the grid
        for (binID = 0; binID < numFiles; binID++) array2[binID].resize(1, 0);
      }

    if (operfunc == func_roc)
      {
        hist.resize(numBins);

        roc.resize(numBins + 1);
        ctg_tab.resize(numBins + 1);
        for (int i = 0; i <= numBins; ++i)
          {
            roc[i].resize(2, 0);
            ctg_tab[i].resize(4, 0);
          }

        uThresh.resize(numBins);
        lThresh.resize(numBins);
        for (int i = 0; i < numBins; ++i)
          {
            uThresh[i] = ((double) i + 1) / numBins;
            lThresh[i] = (double) i / numBins;
          }
      }
  }

  void
  run() override
  {
    const auto &varList1 = ef[0].varList;

    int varID = 0, levelID;
    int numFields0;
    int tsID = 0;
    do {
        numFields0 = cdo_stream_inq_timestep(ef[0].streamID, tsID);
        for (int fileID = 1; fileID < numFiles; ++fileID)
          {
            auto streamID = ef[fileID].streamID;
            int numFields = cdo_stream_inq_timestep(streamID, tsID);
            if (numFields != numFields0)
              {
                if (numFields == 0)
                  cdo_abort("Inconsistent ensemble file, too few time steps in %s!", cdo_get_stream_name(fileID));
                else
                  cdo_abort("Inconsistent ensemble file, number of fields at time step %d of %s and %s differ!", tsID + 1,
                            cdo_get_stream_name(0), cdo_get_stream_name(fileID));
              }
          }

        if (operfunc == func_rank && (datafunc == TIME || tsID == 0))
          {
            cdo_taxis_copy_timestep(taxisID2, taxisID1);
            if (numFields0 > 0) cdo_def_timestep(streamID2, tsID);
          }

        for (int fieldID = 0; fieldID < numFields0; ++fieldID)
          {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(ef, numFiles) private(numMissVals) lastprivate(varID, levelID)
#endif
            for (int fileID = 0; fileID < numFiles; ++fileID)
              {
                auto streamID = ef[fileID].streamID;
                cdo_inq_field(streamID, &varID, &levelID);
                cdo_read_field(streamID, ef[fileID].array.data(), &numMissVals);
              }

            gridsizeMax = varList1.vars[varID].gridsize;
            auto missval = varList1.vars[varID].missval;

            numMissVals = 0;
            if (datafunc == TIME && operfunc == func_rank)
              for (binID = 0; binID < numFiles; binID++) array2[binID][0] = 0;

#ifdef _OPENMP
#pragma omp parallel for default(shared) private(binID)
#endif
            for (size_t i = 0; i < gridsizeMax; ++i)
              {
                auto ompthID = cdo_omp_get_thread_num();

                field[ompthID].missval = missval;
                field[ompthID].numMissVals = 0;
                have_miss = 0;
                for (int fileID = 0; fileID < numFiles; ++fileID)
                  {
                    field[ompthID].vec_d[fileID] = ef[fileID].array[i];
                    if (fp_is_equal(field[ompthID].vec_d[fileID], missval))
                      {
                        have_miss = 1;
                        break;
                      }
                  }

                // need to ignore all data for a gridpoint if a single ensemble
                // has a missing value at that gridpoint.
                if (!have_miss)  // only process if no missing value in ensemble
                  {
                    switch (operfunc)
                      {
                      case (func_rank):
                        /* ****************/
                        /* RANK HISTOGRAM */
                        /* ************** */
                        // for ( j=0; j<numFiles; j++ )
                        //   fprintf(stderr,"%5.2g ",field[ompthID].vec_d[j]);
#ifdef _OPENMP
#pragma omp critical
#endif
                        binID = (int) field_rank(field[ompthID]);
                        // fprintf(stderr,"-->%i\n",binID);

                        if (datafunc == SPACE && !have_miss)
                          array2[binID][i]++;
                        else if (!have_miss)
                          array2[binID][0]++;

                        break;

                      case (func_roc):
                        /* ********************************** */
                        /* RECEIVER OPERATING CHARACTERISTICS */
                        /* ********************************** */
                        auto dat = &field[ompthID].vec_d[1];
                        auto ival = (field[ompthID].vec_d[0] > 0.5) ? 1 : 0;

                        for (binID = 0; binID < numBins; binID++) hist[binID] = 0;

                        for (j = 0; j < numEns; ++j)
                          for (binID = 0; binID < numBins; binID++)
                            if (dat[j] >= lThresh[binID] && dat[j] < uThresh[binID]) hist[binID]++;

                        chksum = 0;
                        for (binID = 0; binID < numBins; binID++) chksum += hist[binID];

                        if (chksum != numEns) exit(1);

                        int cum = 0;
                        if (ival == 1)
                          {
                            // all true positives in first bin
                            ctg_tab[0][TP] += numEns;

                            cum += hist[0];
                            for (binID = 1; binID < numBins; binID++)
                              {
                                ctg_tab[binID][TP] += numEns - cum;
                                ctg_tab[binID][FN] += cum;
                                cum += hist[binID];
                              }
                            ctg_tab[binID][TP] += numEns - cum;
                            ctg_tab[binID][FN] += cum;
                          }
                        else if (ival == 0)
                          {
                            // all false positives in first bin
                            ctg_tab[0][FP] += numEns;
                            cum += hist[0];
                            for (binID = 1; binID < numBins; binID++)
                              {
                                ctg_tab[binID][FP] += numEns - cum;
                                ctg_tab[binID][TN] += cum;
                                cum += hist[binID];
                              }
                            ctg_tab[binID][FP] += numEns - cum;
                            ctg_tab[binID][TN] += cum;
                          }
                        break;

                      }  // switch ( operfunc )
                  }  // if ( ! have_miss )
              }  // for ( i=0; i<gridsize; i++ )

            if (datafunc == TIME && operfunc == func_rank)
              {
                for (binID = 0; binID < numFiles; binID++)
                  {
                    double val = (double) array2[binID][0];
                    //		fprintf(stderr,"%i ",(int)val);
                    cdo_def_field(streamID2, varIDs2[varID], binID);
                    cdo_write_field(streamID2, &val, numMissVals);
                  }
                // fprintf(stderr,"\n");
              }
            else if (operfunc == func_roc)
              {
                if (DEBUG_ROC)
                  {
                    fprintf(stderr, "#             :     TP     FP     FN     TN         TPR        FPR\n");

                    for (binID = 0; binID <= numBins; binID++)
                      {
                        int p = ctg_tab[binID][TP] + ctg_tab[binID][FN];
                        int n = ctg_tab[binID][FP] + ctg_tab[binID][TN];
                        double tpr = ctg_tab[binID][TP] / (double) p;
                        double fpr = ctg_tab[binID][FP] / (double) n;
                        chksum += ctg_tab[binID][0] + ctg_tab[binID][1] + ctg_tab[binID][2] + ctg_tab[binID][3];

                        roc[binID][TPR] = tpr;
                        roc[binID][FPR] = fpr;

                        fprintf(stderr, "%3i %10.4g: %6i %6i %6i %6i: %10.4g %10.4g\n", binID,
                                (binID < numBins) ? lThresh[binID] : 1, ctg_tab[binID][0], ctg_tab[binID][1], ctg_tab[binID][2],
                                ctg_tab[binID][3], tpr, fpr);
                      }
                    fprintf(stderr, "nbins %10i\n", numBins);
                    fprintf(stderr, "#ROC CurveArea: %10.6f\n", roc_curve_integrate(roc, numBins));
                  }  // if ( DEBUG_ROC )
              }  // else if (operfunc == func_roc )
          }  // for ( fieldID=0; fieldID<numFields0; fieldID++ )
        tsID++;
      }  // do [...]
    while (numFields0 > 0);

    if (operfunc == func_rank)
      {
        int osize = (datafunc == TIME) ? 1 : gridsizeMax;
        Varray<double> tmpdoub(osize);

        for (binID = 0; binID < numFiles; binID++)
          {
            for (int i = 0; i < osize; ++i) tmpdoub[i] = (double) array2[binID][i];

            cdo_def_field(streamID2, varIDs2[varID], binID);
            cdo_write_field(streamID2, tmpdoub.data(), numMissVals);
          }
      }
    else if (operfunc == func_roc)
      {
        fprintf(stdout, "#             :     TP     FP     FN     TN         TPR        FPR\n");

        for (int i = 0; i <= numBins; ++i)
          {
            int p = ctg_tab[i][TP] + ctg_tab[i][FN];
            int n = ctg_tab[i][FP] + ctg_tab[i][TN];
            double tpr = ctg_tab[i][TP] / (double) p;
            double fpr = ctg_tab[i][FP] / (double) n;
            // chksum += ctg_tab[i][0] + ctg_tab[i][1] + ctg_tab[i][2] + ctg_tab[i][3];

            roc[i][TPR] = tpr;
            roc[i][FPR] = fpr;

            int sum = ctg_tab[i][TP] + ctg_tab[i][TN] + ctg_tab[i][FP] + ctg_tab[i][FN];

            fprintf(stdout, "%3i %10.4g: %6i %6i %6i %6i (%6i): %10.4g %10.4g\n", i, (i < numBins) ? lThresh[i] : 1, ctg_tab[i][0],
                    ctg_tab[i][1], ctg_tab[i][2], ctg_tab[i][3], sum, tpr, fpr);
          }

        fprintf(stdout, "#ROC CurveArea: %10.6f\n", roc_curve_integrate(roc, numBins));
      }
  }
  void
  close() override
  {
    for (int fileID = 0; fileID < numFiles; ++fileID) cdo_stream_close(ef[fileID].streamID);

    if (operfunc != func_roc) cdo_stream_close(streamID2);
  }
};
