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

      Splittime  splithour       Split hours
      Splittime  splitday        Split days
      Splittime  splitmon        Split months
      Splittime  splitseas       Split seasons
*/

#include <time.h>
#include <cdi.h>

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

constexpr int MaxStreams = 32;

struct tm
datetime_to_tm(CdiDateTime vDateTime)
{
  int year, month, day, hour, minute, second, ms;
  cdiDate_decode(vDateTime.date, &year, &month, &day);
  cdiTime_decode(vDateTime.time, &hour, &minute, &second, &ms);

  struct tm stime;
  memset(&stime, 0, sizeof(struct tm));

  stime.tm_sec = second;
  stime.tm_min = minute;
  stime.tm_hour = hour;
  stime.tm_mday = day;
  stime.tm_mon = month - 1;
  stime.tm_year = year - 1900;

  return stime;
}

class Splittime : public Process
{
  enum
  {
    func_time,
    func_date
  };

public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Splittime",
    .operators = { { "splithour", func_time, 10000, SplittimeHelp },
                   { "splitday", func_date, 1, SplittimeHelp },
                   { "splitmon", func_date, 100, SplittimeHelp },
                   { "splitseas", func_date, 100, SplittimeHelp } },
    .aliases = {},
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_BOTH,  // Allowed number type
    .constraints = { 1, OBASE, OnlyFirst },
  };
  inline static RegisterEntry<Splittime> registration = RegisterEntry<Splittime>(module);

  int SPLITMON, SPLITSEAS;
  CdoStreamID streamID1;
  CdoStreamID streamID2;
  CdoStreamID streamIDs[MaxStreams];
  int tsIDs[MaxStreams];
  int index = 0;
  int taxisID1;
  int taxisID2;

  int vlistID1;
  int vlistID2;

  int operfunc;
  int operintval;
  int operatorID;

  std::string fileSuffix;
  const char *format = nullptr;

  VarList varList1;
  FieldVector2D vars;
  Varray<double> array;

  bool haveConstVars;
  bool dataIsUnchanged;

public:
  void
  init() override
  {
    dataIsUnchanged = data_is_unchanged();

    SPLITMON = module.get_id("splitmon");
    SPLITSEAS = module.get_id("splitseas");

    operatorID = cdo_operator_id();
    operfunc = cdo_operator_f1(operatorID);
    operintval = cdo_operator_f2(operatorID);

    if (operatorID == SPLITMON && cdo_operator_argc() == 1)
      format = cdo_operator_argv(0).c_str();
    else
      operator_check_argc(0);

    for (int i = 0; i < MaxStreams; ++i) streamIDs[i] = CDO_STREAM_UNDEF;
    for (int i = 0; i < MaxStreams; ++i) tsIDs[i] = 0;

    streamID1 = cdo_open_read(0);

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

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

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

    varList1 = VarList(vlistID1);

    //  if (! dataIsUnchanged)
    {
      auto gridsizemax = vlistGridsizeMax(vlistID1);
      if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;
      array.resize(gridsizemax);
    }

    haveConstVars = (varList1.numConstVars() > 0);

    if (haveConstVars)
      {
        int numVars = varList1.numVars();
        vars.resize(numVars);

        for (int varID = 0; varID < numVars; ++varID)
          {
            const auto &var = varList1.vars[varID];
            if (var.isConstant)
              {
                vars[varID].resize(var.nlevels);

                for (int levelID = 0; levelID < var.nlevels; ++levelID)
                  {
                    vars[varID][levelID].grid = var.gridID;
                    vars[varID][levelID].resize(var.gridsize);
                  }
              }
          }
      }
  }

  void
  run() override
  {
    auto seasonNames = get_season_name();

    int tsID = 0;
    while (true)
      {
        auto numFields = cdo_stream_inq_timestep(streamID1, tsID);
        if (numFields == 0) break;

        cdo_taxis_copy_timestep(taxisID2, taxisID1);
        auto vDateTime = taxisInqVdatetime(taxisID1);

        if (operfunc == func_date)
          {
            index = (cdiDate_get(vDateTime.date) / operintval) % 100;
            if (index < 0) index = -index;

            if (operatorID == SPLITSEAS) index = month_to_season(index);
          }
        else if (operfunc == func_time) { index = (cdiTime_get(vDateTime.time) / operintval) % 100; }

        if (index < 0 || index >= MaxStreams) cdo_abort("Index out of range!");

        streamID2 = streamIDs[index];
        if (streamID2 == CDO_STREAM_UNDEF)
          {
            auto fileName = cdo_get_obase();
            if (operatorID == SPLITSEAS)
              {
                fileName += string_format("%3s", seasonNames[index]);
                if (fileSuffix.size() > 0) fileName += fileSuffix;
              }
            else
              {
                char oformat[32];
                std::strcpy(oformat, "%02d");

                if (operatorID == SPLITMON && format)
                  {
                    char sbuf[32];
                    auto stime = datetime_to_tm(vDateTime);
                    auto slen = strftime(sbuf, sizeof(sbuf), format, &stime);
                    if (slen) std::strcpy(oformat, sbuf);
                  }

                fileName += string_format(oformat, index);
                if (fileSuffix.size() > 0) fileName += fileSuffix;
              }

            if (Options::cdoVerbose) cdo_print("create file %s", fileName);

            streamID2 = cdo_open_write(fileName.c_str());
            cdo_def_vlist(streamID2, vlistID2);
            streamIDs[index] = streamID2;
          }

        cdo_def_timestep(streamID2, tsIDs[index]);

        if (tsID > 0 && tsIDs[index] == 0 && haveConstVars)
          {
            int numVars = varList1.numVars();
            for (int varID = 0; varID < numVars; ++varID)
              {
                const auto &var = varList1.vars[varID];
                if (var.isConstant)
                  {
                    for (int levelID = 0; levelID < var.nlevels; ++levelID)
                      {
                        cdo_def_field(streamID2, varID, levelID);
                        auto numMissVals = vars[varID][levelID].numMissVals;
                        cdo_write_field(streamID2, vars[varID][levelID].vec_d.data(), numMissVals);
                      }
                  }
              }
          }

        for (int fieldID = 0; fieldID < numFields; ++fieldID)
          {
            auto [varID, levelID] = cdo_inq_field(streamID1);
            cdo_def_field(streamID2, varID, levelID);

            if (dataIsUnchanged && !(tsID == 0 && haveConstVars)) { cdo_copy_field(streamID2, streamID1); }
            else
              {
                size_t numMissVals;
                cdo_read_field(streamID1, array.data(), &numMissVals);
                cdo_write_field(streamID2, array.data(), numMissVals);

                if (tsID == 0 && haveConstVars)
                  {
                    const auto &var = varList1.vars[varID];
                    if (var.isConstant)
                      {
                        varray_copy(var.gridsize, array, vars[varID][levelID].vec_d);
                        vars[varID][levelID].numMissVals = numMissVals;
                      }
                  }
              }
          }

        tsIDs[index]++;
        tsID++;
      }
  }

  void
  close() override
  {
    cdo_stream_close(streamID1);

    for (index = 0; index < MaxStreams; ++index)
      {
        if (streamIDs[index] != CDO_STREAM_UNDEF) cdo_stream_close(streamIDs[index]);
      }

    vlistDestroy(vlistID2);
  }
};
