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

  Copyright (C) 2003-2018 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:
   Ensstat3       ensrkhist_space  Ensemble ranked histogram averaged over time
   Ensstat3       ensrkhist_time   Ensemble ranked histogram averaged over space
   Ensstat3       ensroccurve      Ensamble Receiver Operating Characteristics
*/

#include <cdi.h>

#include "cdo_int.h"
#include "pstream_int.h"
#include "cdoOptions.h"
#include "util_files.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,
  FP,
  FN,
  TN
};
/* TP - True positive  ( event     forecast and     occured)  HIT */
/* FP - False positive ( event     forecast and not occured)  FALSE ALARM */
/* TN - True negative  ( event not forecast and not ocurred)  CORRECT REJECTION
 */
/* FN - False negative ( event not forecast and     ocurred)  MISSED */

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

double roc_curve_integrate(const double **roc, const int n);

void *
Ensstat3(void *process)
{
  int i, j;
  int nrecs = 0, nrecs0;
  size_t nmiss;
  int cum;
  int chksum;  // for check of histogram population
  int levelID, varID, binID = 0;
  int vlistID;
  int gridID, gridID2;
  int have_miss = 0;
  int streamID = 0, streamID2 = 0;
  int **array2 = NULL;
  int **ctg_tab = NULL, *hist = NULL;  // contingency table and histogram
  double missval;
  double *dat;                              // pointer to ensemble data for ROC
  double *uThresh = NULL, *lThresh = NULL;  // thresholds for histograms
  double **roc = NULL;                      // receiver operating characteristics table
  double val;
  int ival;

  struct ens_file_t
  {
    int streamID;
    int vlistID;
    double *array;
  };

  cdoInitialize(process);

  // clang-format off
  cdoOperatorAdd("ensroc",          func_roc,  0,          NULL);
  cdoOperatorAdd("ensrkhist_space", func_rank, space_data, NULL);
  cdoOperatorAdd("ensrkhist_time",  func_rank, time_data,  NULL);
  // clang-format on

  int operatorID = cdoOperatorID();
  int operfunc = cdoOperatorF1(operatorID);
  int datafunc = cdoOperatorF2(operatorID);

  int nbins = 0;
  if (operfunc == func_roc)
    {
      operatorInputArg("Number of eigen functions to write out");
      nbins = parameter2int(operatorArgv()[0]);
    }

  int nfiles = cdoStreamCnt() - 1;
  int nens = nfiles - 1;

  if (cdoVerbose) cdoPrint("Ensemble over %d files.", nfiles);

  const char *ofilename = cdoGetStreamName(nfiles).c_str();

  if (!cdoOverwriteMode && fileExists(ofilename) && !userFileOverwrite(ofilename))
    cdoAbort("Outputfile %s already exists!", ofilename);

  ens_file_t *ef = (ens_file_t *) Malloc(nfiles * sizeof(ens_file_t));

  /* *************************************************** */
  /* should each thread be allocating memory locally???? */
  /* ("first touch strategy")                            */
  /* --> #pragma omp parallel for ...                    */
  /* *************************************************** */
  Field *field = (Field *) Malloc(Threading::ompNumThreads * sizeof(Field));
  for (i = 0; i < Threading::ompNumThreads; i++)
    {
      field_init(&field[i]);
      field[i].size = nfiles;
      field[i].ptr = (double *) Malloc(nfiles * sizeof(double));
      field[i].weight = (double *) Malloc(nfiles * sizeof(double));
      for (int fileID = 0; fileID < nfiles; fileID++) field[i].weight[fileID] = 1;
    }

  for (int fileID = 0; fileID < nfiles; fileID++)
    {
      streamID = cdoStreamOpenRead(cdoStreamName(fileID));

      vlistID = cdoStreamInqVlist(streamID);

      ef[fileID].streamID = streamID;
      ef[fileID].vlistID = vlistID;
    }

  /* check for identical contents of all ensemble members */
  for (int fileID = 1; fileID < nfiles; fileID++) vlistCompare(ef[0].vlistID, ef[fileID].vlistID, CMP_ALL);

  int vlistID1 = ef[0].vlistID;
  int vlistID2 = vlistCreate();
  int nvars = vlistNvars(vlistID1);
  int *varID2 = (int *) Malloc(nvars * sizeof(int));

  double *levs = (double *) Calloc(nfiles, sizeof(double));
  int zaxisID2 = zaxisCreate(ZAXIS_GENERIC, nfiles);
  for (i = 0; i < nfiles; i++) levs[i] = i;
  zaxisDefLevels(zaxisID2, levs);
  zaxisDefName(zaxisID2, "histogram_binID");

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

  for (varID = 0; varID < nvars; varID++)
    {

      /* **************************************************************** */
      /* nfiles includes the observation, so there are nfiles-1 ensembles */
      /* and exactly nfiles bins, in which the observation could fall     */
      /* **************************************************************** */

      if (datafunc == TIME)
        {
          val = 0;
          gridID2 = gridCreate(GRID_LONLAT, 1);
          gridDefXsize(gridID2, 1);
          gridDefYsize(gridID2, 1);
          gridDefXvals(gridID2, &val);
          gridDefYvals(gridID2, &val);
        }
      else  // datafunc == SPACE
        gridID2 = vlistInqVarGrid(vlistID1, varID);

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

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

  for (varID = 0; varID < nvars; varID++)
    {
      if (zaxisInqSize(vlistInqVarZaxis(vlistID1, varID)) > 1)
        {
          cdoWarning("More than one level not supported when processing ranked histograms.");
          cdoWarning("Try to use `cdo splitlevel` to split the dataset into levels and apply");
          cdoWarning("the operator seperately to each level.");
          cdoAbort("Exit due to unsupported file structure");
        }
    }

  if (operfunc != func_roc)
    {
      streamID2 = cdoStreamOpenWrite(cdoStreamName(nfiles), cdoFiletype());
      pstreamDefVlist(streamID2, vlistID2);
    }

  size_t gridsize = vlistGridsizeMax(vlistID1);

  for (int fileID = 0; fileID < nfiles; fileID++) ef[fileID].array = (double *) Malloc(gridsize * sizeof(double));

  if (operfunc == func_rank && datafunc == SPACE)
    { /*need to memorize data for entire grid before writing          */
      array2 = (int **) Malloc((nfiles + 1) * sizeof(int *));
      for (binID = 0; binID < nfiles; binID++) array2[binID] = (int *) Calloc(gridsize, sizeof(int));
    }
  else if (operfunc == func_rank)
    { /* can process data separately for each timestep and only need */
      /* to cumulate values over the grid                            */
      array2 = (int **) Malloc((nfiles + 1) * sizeof(int *));
      for (binID = 0; binID < nfiles; binID++) array2[binID] = (int *) Calloc(1, sizeof(int));
    }

  if (operfunc == func_roc)
    {
      ctg_tab = (int **) Malloc((nbins + 1) * sizeof(int *));
      hist = (int *) Malloc(nbins * sizeof(int));
      uThresh = (double *) Malloc(nbins * sizeof(double));
      lThresh = (double *) Malloc(nbins * sizeof(double));
      roc = (double **) Malloc((nbins + 1) * sizeof(double *));

      for (i = 0; i < nbins; i++)
        {
          ctg_tab[i] = (int *) Calloc(4, sizeof(int));
          roc[i] = (double *) Calloc(2, sizeof(double));
          uThresh[i] = ((double) i + 1) / nbins;
          lThresh[i] = (double) i / nbins;
        }
      ctg_tab[nbins] = (int *) Calloc(4, sizeof(int));
      roc[nbins] = (double *) Calloc(2, sizeof(double));
    }

  int tsID = 0;
  do
    {
      nrecs0 = cdoStreamInqTimestep(ef[0].streamID, tsID);
      for (int fileID = 1; fileID < nfiles; fileID++)
        {
          streamID = ef[fileID].streamID;
          nrecs = cdoStreamInqTimestep(streamID, tsID);
          if (nrecs != nrecs0)
            {
              if (nrecs == 0)
                cdoAbort("Inconsistent ensemble file, too few time steps in %s!", cdoGetStreamName(fileID).c_str());
              else
                cdoAbort("Inconsistent ensemble file, number of records at time step %d of %s and %s differ!", tsID + 1,
                         cdoGetStreamName(0).c_str(), cdoGetStreamName(fileID).c_str());
            }
        }

      if (operfunc == func_rank && (datafunc == TIME || tsID == 0))
        {
          taxisCopyTimestep(taxisID2, taxisID1);
          if (nrecs0 > 0) pstreamDefTimestep(streamID2, tsID);
        }

      //      fprintf(stderr,"TIMESTEP %i varID %i rec %i\n",tsID,varID,recID);

      for (int recID = 0; recID < nrecs0; recID++)
        {
#ifdef _OPENMP
#pragma omp parallel for default(none) shared(ef, nfiles) private(streamID, nmiss) lastprivate(varID, levelID)
#endif
          for (int fileID = 0; fileID < nfiles; fileID++)
            {
              streamID = ef[fileID].streamID;
              pstreamInqRecord(streamID, &varID, &levelID);
              pstreamReadRecord(streamID, ef[fileID].array, &nmiss);
            }

          gridID = vlistInqVarGrid(vlistID1, varID);
          gridsize = gridInqSize(gridID);
          missval = vlistInqVarMissval(vlistID1, varID);

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

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

              field[ompthID].missval = missval;
              field[ompthID].nmiss = 0;
              have_miss = 0;
              for (int fileID = 0; fileID < nfiles; fileID++)
                {
                  field[ompthID].ptr[fileID] = ef[fileID].array[i];
                  if (DBL_IS_EQUAL(field[ompthID].ptr[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<nfiles; j++ )
//   fprintf(stderr,"%5.2g ",field[ompthID].ptr[j]);
#ifdef _OPENMP
#pragma omp critical
#endif
                      binID = (int) fldfun(field[ompthID], operfunc);
                      // 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 */
                      /* ********************************** */
                      dat = &field[ompthID].ptr[1];
                      ival = field[ompthID].ptr[0] > 0.5 ? 1 : 0;

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

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

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

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

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

                          cum += hist[0];
                          for (binID = 1; binID < nbins; binID++)
                            {
                              ctg_tab[binID][TP] += nens - cum;
                              ctg_tab[binID][FN] += cum;
                              cum += hist[binID];
                            }
                          ctg_tab[binID][TP] += nens - cum;
                          ctg_tab[binID][FN] += cum;
                        }
                      else if (ival == 0)
                        {
                          // all false positives in first bin
                          ctg_tab[0][FP] += nens;
                          cum += hist[0];
                          for (binID = 1; binID < nbins; binID++)
                            {
                              ctg_tab[binID][FP] += nens - cum;
                              ctg_tab[binID][TN] += cum;
                              cum += hist[binID];
                            }
                          ctg_tab[binID][FP] += nens - 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 < nfiles; binID++)
                {
                  val = (double) array2[binID][0];
                  //		fprintf(stderr,"%i ",(int)val);
                  pstreamDefRecord(streamID2, varID2[varID], binID);
                  pstreamWriteRecord(streamID2, &val, nmiss);
                }
              // fprintf(stderr,"\n");
            }
          else if (operfunc == func_roc)
            {
              if (DEBUG_ROC)
                {
                  fprintf(stderr, "#             :     TP     FP     FN     TN "
                                  "        TPR        FPR\n");

                  for (binID = 0; binID <= nbins; 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 < nbins ? 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", nbins);
                  fprintf(stderr, "#ROC CurveArea: %10.6f\n", roc_curve_integrate((const double **) roc, nbins));
                }  // if ( DEBUG_ROC )
            }      // else if (operfunc == func_roc )
        }          // for ( recID=0; recID<nrecs; recID++ )
      tsID++;
    }  // do [...]
  while (nrecs0 > 0);

  if (operfunc == func_rank)
    {
      double *tmpdoub;
      int osize = gridsize;

      if (datafunc == TIME) osize = 1;
      tmpdoub = (double *) Malloc(osize * sizeof(double));

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

          pstreamDefRecord(streamID2, varID2[varID], binID);
          pstreamWriteRecord(streamID2, tmpdoub, nmiss);
        }

      Free(tmpdoub);
    }
  else if (operfunc == func_roc)
    {
      fprintf(stdout, "#             :     TP     FP     FN     TN         TPR        FPR\n");

      for (i = 0; i <= nbins; 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 < nbins ? 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((const double **) roc, nbins));
    }

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

  if (operfunc != func_roc) pstreamClose(streamID2);

  for (int fileID = 0; fileID < nfiles; fileID++)
    if (ef[fileID].array) Free(ef[fileID].array);

  if (ef) Free(ef);
  if (array2)
    {
      for (binID = 0; binID < nfiles; binID++) Free(array2[binID]);
      Free(array2);
    }

  for (i = 0; i < Threading::ompNumThreads; i++)
    {
      if (field[i].ptr) Free(field[i].ptr);
      if (field[i].weight) Free(field[i].weight);
    }

  if (field) Free(field);

  cdoFinish();

  return 0;
}

double
roc_curve_integrate(const double **roc, const int n)
{
  double y1, y0, x1, x0, dx, dy;
  double step_area;
  double area = 0;

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

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

  return area - 0.5;
}
