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

  Eofcoeff             eofcoeff             process eof coefficients
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "util_files.h"
#include "util_string.h"

// NO MISSING VALUE SUPPORT ADDED SO FAR

class ModuleEofcoeff
{
  double missval1 = -999, missval2;
  int varID, levelID;
  int nrecs;
  size_t nmiss;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID2;
  int taxisID3;

  int vlistID1;
  int vlistID2;
  int vlistID3;

  int gridID1;
  int eofID;
  size_t gridsize;
  int neof;
  int nvars;
  int nlevs;

  std::vector<CdoStreamID> streamIDs;

  FieldVector3D eof;
  std::string fileSuffix;

public:
  void
  init(void *process)
  {

    cdo_initialize(process);

    if (process_self().m_ID != 0) cdo_abort("This operator can't be combined with other operators!");

    cdo_operator_add("eofcoeff", 0, 0, nullptr);

    streamID1 = cdo_open_read(0);
    streamID2 = cdo_open_read(1);

    vlistID1 = cdo_stream_inq_vlist(streamID1);
    vlistID2 = cdo_stream_inq_vlist(streamID2);
    vlistID3 = vlistDuplicate(vlistID2);

    // taxisID1 = vlistInqTaxis(vlistID1);
    taxisID2 = vlistInqTaxis(vlistID2);
    taxisID3 = taxisDuplicate(taxisID2);

    gridID1 = vlistInqVarGrid(vlistID1, 0);
    int gridID2 = vlistInqVarGrid(vlistID2, 0);

    gridsize = vlistGridsizeMax(vlistID1);
    if (gridsize != vlistGridsizeMax(vlistID2))
      cdo_abort("Gridsize of input files does not match! %zu and %zu", gridsize, vlistGridsizeMax(vlistID2));

    if (vlistNgrids(vlistID2) > 1 || vlistNgrids(vlistID1) > 1) cdo_abort("Too many different grids in input!");

    nvars = vlistNvars(vlistID1) == vlistNvars(vlistID2) ? vlistNvars(vlistID1) : -1;
    nlevs = zaxisInqSize(vlistInqVarZaxis(vlistID1, 0));

    if (gridID1 != gridID2) cdo_compare_grids(gridID1, gridID2);

    fileSuffix = FileUtils::gen_suffix(cdo_inq_filetype(streamID1), vlistID1, cdo_get_stream_name(0));

    eof = FieldVector3D(nvars);
    for (varID = 0; varID < nvars; ++varID) eof[varID].resize(nlevs);

    eofID = 0;
  }
  void
  run()
  {
    while (1)
      {
        nrecs = cdo_stream_inq_timestep(streamID1, eofID);
        if (nrecs == 0) break;

        for (int recID = 0; recID < nrecs; ++recID)
          {
            cdo_inq_record(streamID1, &varID, &levelID);
            missval1 = vlistInqVarMissval(vlistID1, varID);
            eof[varID][levelID].resize((eofID == 0) ? 1 : eofID + 1);
            eof[varID][levelID][eofID].grid = gridID1;
            eof[varID][levelID][eofID].missval = missval1;
            eof[varID][levelID][eofID].resize(gridsize);
            varray_fill(gridsize, eof[varID][levelID][eofID].vec_d, missval1);

            if (varID >= nvars) cdo_abort("Internal error - varID >= nvars");
            if (levelID >= nlevs) cdo_abort("Internal error - levelID >= nlevs");

            cdo_read_record(streamID1, eof[varID][levelID][eofID].vec_d.data(), &nmiss);
            eof[varID][levelID][eofID].nmiss = nmiss;
          }
        eofID++;
      }

    neof = eofID;

    if (Options::cdoVerbose) cdo_print("%s contains %i eof's", cdo_get_stream_name(0), neof);
    // Create 1x1 Grid for output
    auto gridID3 = gridCreate(GRID_LONLAT, 1);
    gridDefXsize(gridID3, 1);
    gridDefYsize(gridID3, 1);
    double xvals = 0.0, yvals = 0.0;
    gridDefXvals(gridID3, &xvals);
    gridDefYvals(gridID3, &yvals);

    // Create var-list and time-axis for output

    auto ngrids = vlistNgrids(vlistID3);
    for (int i = 0; i < ngrids; ++i) vlistChangeGridIndex(vlistID3, i, gridID3);

    vlistDefTaxis(vlistID3, taxisID3);
    for (varID = 0; varID < nvars; ++varID) vlistDefVarTimetype(vlistID3, varID, TIME_VARYING);

    // open streams for eofcoeff output
    streamIDs = std::vector<CdoStreamID>(neof);
    for (eofID = 0; eofID < neof; eofID++)
      {
        auto fileName = cdo_get_obase() + string_format("%5.5i", eofID);
        if (fileSuffix.size() > 0) fileName += fileSuffix;

        streamIDs[eofID] = cdo_open_write(fileName.c_str());

        if (Options::cdoVerbose)
          cdo_print("opened %s ('w')  as stream%i for %i. eof", fileName, streamIDs[eofID]->get_id(), eofID + 1);

        cdo_def_vlist(streamIDs[eofID], vlistID3);
      }

    // ALLOCATE temporary fields for data read and write
    Field in, out;
    in.resize(gridsize);
    in.grid = gridID1;
    out.missval = missval1;
    out.nmiss = 0;
    out.resize(1);

    int tsID = 0;
    while (1)
      {
        nrecs = cdo_stream_inq_timestep(streamID2, tsID);
        if (nrecs == 0) break;

        cdo_taxis_copy_timestep(taxisID3, taxisID2);
        /*for ( eofID=0; eofID<neof; eofID++)
          {
            fprintf(stderr, "defining ts %i\n", tsID);
            cdo_def_timestep(streamIDs[eofID],tsID);
          }
        */
        for (int recID = 0; recID < nrecs; ++recID)
          {
            cdo_inq_record(streamID2, &varID, &levelID);
            missval2 = vlistInqVarMissval(vlistID2, varID);
            cdo_read_record(streamID2, in.vec_d.data(), &in.nmiss);

            for (eofID = 0; eofID < neof; eofID++)
              {
                if (recID == 0) cdo_def_timestep(streamIDs[eofID], tsID);
                // if ( recID == 0 ) fprintf(stderr, "ts%i rec%i eof%i\n", tsID, recID, eofID);
                out.vec_d[0] = 0;
                out.grid = gridID3;
                out.missval = missval2;
                for (size_t i = 0; i < gridsize; ++i)
                  {
                    if (!DBL_IS_EQUAL(in.vec_d[i], missval2) && !DBL_IS_EQUAL(eof[varID][levelID][eofID].vec_d[i], missval1))
                      {
                        // out.vec_d[0] += w[i]*in.vec_d[i]*eof[varID][levelID][eofID].vec_d[i];
                        out.vec_d[0] += in.vec_d[i] * eof[varID][levelID][eofID].vec_d[i];
                      }
                  }
                if (!DBL_IS_EQUAL(out.vec_d[0], 0.))
                  nmiss = 0;
                else
                  {
                    nmiss = 1;
                    out.vec_d[0] = missval2;
                  }

                cdo_def_record(streamIDs[eofID], varID, levelID);
                // fprintf(stderr, "%d %d %d %d %d %g\n", streamIDs[eofID],tsID, recID, varID, levelID,*out.vec_d.data());
                cdo_write_record(streamIDs[eofID], out.vec_d.data(), nmiss);
              }
            if (varID >= nvars) cdo_abort("Internal error - varID >= nvars");
            if (levelID >= nlevs) cdo_abort("Internal error - levelID >= nlevs");
          }

        tsID++;
      }
  }
  void
  close()
  {
    for (eofID = 0; eofID < neof; eofID++) cdo_stream_close(streamIDs[eofID]);

    cdo_stream_close(streamID2);
    cdo_stream_close(streamID1);

    cdo_finish();
  }
};

void *
Eofcoeff(void *process)
{
  ModuleEofcoeff eofcoeff;

  eofcoeff.init(process);
  eofcoeff.run();
  eofcoeff.close();

  return nullptr;
}
