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

      Merge      mergetime       Merge datasets sorted by date and time
*/

#include <cdi.h>

#include "cdo_options.h"
#include "cdo_rlimit.h"
#include "process_int.h"
#include "util_files.h"
#include "printinfo.h"
#include "param_conversion.h"
#include "pmlist.h"

struct StreamInfo
{
  CdoStreamID streamID;
  CdiDateTime vDateTime{};
  int vlistID;
  int taxisID;
  int tsID;
  int numFields;
  VarList varList;
  std::map<int, int> mapOfVarIDs;
};

bool
getenv_skip_same_time()
{
  auto envString = getenv_string("SKIP_SAME_TIME");
  if (envString.size())
    {
      auto ival = std::stoi(envString);
      if (ival == 1)
        {
          if (Options::cdoVerbose) cdo_print("Set SKIP_SAME_TIME to %d", ival);
          return true;
        }
    }

  return false;
}

static void
open_all_files(int numFiles, std::vector<StreamInfo> &streamInfoList)
{
  for (int fileID = 0; fileID < numFiles; ++fileID)
    {
      if (Options::cdoVerbose) cdo_print("process: %s", cdo_get_stream_name(fileID));

      auto &si = streamInfoList[fileID];
      si.streamID = cdo_open_read(fileID);
      si.vlistID = cdo_stream_inq_vlist(si.streamID);
      si.taxisID = vlistInqTaxis(si.vlistID);
      si.varList = VarList(si.vlistID);
    }
}

static void
read_first_timestep(int numFiles, std::vector<StreamInfo> &streamInfoList)
{
  for (int fileID = 0; fileID < numFiles; ++fileID)
    {
      auto &si = streamInfoList[fileID];
      si.tsID = 0;
      si.numFields = cdo_stream_inq_timestep(si.streamID, si.tsID);
      if (si.numFields == 0)
        {
          cdo_stream_close(si.streamID);
          si.streamID = CDO_STREAM_UNDEF;
        }
      else { si.vDateTime = taxisInqVdatetime(si.taxisID); }
    }
}

static void
get_parameter(bool &skipSameTime, MapFlag &mapFlag)
{
  auto pargc = cdo_operator_argc();
  if (pargc)
    {
      const auto &pargv = cdo_get_oper_argv();

      KVList kvlist;
      kvlist.name = cdo_module_name();
      if (kvlist.parse_arguments(pargv) != 0) cdo_abort("Parse error!");
      if (Options::cdoVerbose) kvlist.print();

      for (const auto &kv : kvlist)
        {
          const auto &key = kv.key;
          if (kv.nvalues > 1) cdo_abort("Too many values for parameter key >%s<!", key);
          if (kv.nvalues < 1) cdo_abort("Missing value for parameter key >%s<!", key);
          const auto &value = kv.values[0];

          // clang-format off
          if      (key == "skip_same_time")  skipSameTime = parameter_to_bool(value);
          else if (key == "names")
            {
              if      (value == "union")     mapFlag = MapFlag::Left;
              else if (value == "intersect") mapFlag = MapFlag::Intersect;
              else cdo_abort("Invalid value for key >%s< (names=<union/intersect>)", key, value);
            }
          else cdo_abort("Invalid parameter key >%s<!", key);
          // clang-format on
        }
    }
}

class Mergetime : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Mergetime",
    .operators = { { "mergetime", MergeHelp } },
    .aliases = {},
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { -1, 1, NoRestriction },
  };
  inline static RegisterEntry<Mergetime> registration = RegisterEntry<Mergetime>(module);

  int tsID2 = 0;
  int vlistID2 = CDI_UNDEFID;
  int taxisID2 = CDI_UNDEFID;
  CdiDateTime lastDateTime{};
  Field field;

  CdoStreamID streamID2;

  int numFiles{ 0 };
  bool dataIsUnchanged{ false };
  bool skipSameTime{ false };
  std::vector<StreamInfo> streamInfoList;
  int vlistFileIDmin{ -1 };
  int vlistFileIDmax{ -1 };

  MapFlag mapFlag{ MapFlag::Undefined };

public:
  void
  init() override
  {
    skipSameTime = getenv_skip_same_time();

    get_parameter(skipSameTime, mapFlag);

    dataIsUnchanged = data_is_unchanged();

    numFiles = cdo_stream_cnt() - 1;
    streamInfoList.resize(numFiles);

    cdo::set_numfiles(numFiles + 8);

    open_all_files(numFiles, streamInfoList);

    // check that the contents is always the same
    if (mapFlag == MapFlag::Undefined)
      {
        for (int fileID = 1; fileID < numFiles; ++fileID)
          {
            auto &si = streamInfoList[fileID];
            varList_compare(streamInfoList[0].varList, si.varList);
            for (const auto &var : si.varList.vars) si.mapOfVarIDs[var.ID] = var.ID;
          }
      }
    else
      {
        vlistFileIDmin = 0;
        vlistFileIDmax = 0;
        for (int fileID = 1; fileID < numFiles; ++fileID)
          {
            auto numVars = streamInfoList[fileID].varList.numVars();
            if (numVars < streamInfoList[vlistFileIDmin].varList.numVars()) vlistFileIDmin = fileID;
            if (numVars > streamInfoList[vlistFileIDmax].varList.numVars()) vlistFileIDmax = fileID;
          }

        const auto &varList2 = streamInfoList[(mapFlag == MapFlag::Intersect) ? vlistFileIDmin : vlistFileIDmax].varList;
        for (int fileID = 0; fileID < numFiles; ++fileID)
          {
            auto &si = streamInfoList[fileID];
            varList_map(si.varList, varList2, mapFlag, si.mapOfVarIDs);
          }
      }

    // read the first time step
    read_first_timestep(numFiles, streamInfoList);

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

    streamID2 = cdo_open_write(numFiles);
  }

  void
  fill_missing_fields(const StreamInfo &si1, StreamInfo &si2)
  {
    auto maxVars = si1.varList.numVars();
    auto numVars = si2.varList.numVars();
    if (numVars < maxVars)
      {
        std::vector<short> missingIDs(maxVars, 0);
        for (int varID = 0; varID < numVars; ++varID) { missingIDs[si2.mapOfVarIDs[varID]] = 1; }
        for (int varID = 0; varID < maxVars; ++varID)
          {
            if (missingIDs[varID] != 1)
              {
                const auto &var1 = si1.varList.vars[varID];
                field.init(var1);
                field_fill(field, var1.missval);
                field.numMissVals = field.gridsize;
                for (int levelID = 0; levelID < var1.nlevels; levelID++)
                  {
                    cdo_def_field(streamID2, varID, levelID);
                    cdo_write_field(streamID2, field);
                  }
              }
          }
      }
  }

  void
  run() override
  {
    while (true)
      {
        auto processTimestep = true;

        int nextFileID = -1;
        CdiDateTime vDateTime{};
        for (int fileID = 0; fileID < numFiles; ++fileID)
          {
            if (streamInfoList[fileID].streamID != CDO_STREAM_UNDEF)
              {
                auto vdate = cdiDate_get(streamInfoList[fileID].vDateTime.date);
                auto vtime = cdiTime_get(streamInfoList[fileID].vDateTime.time);
                if (nextFileID == -1 || vdate < cdiDate_get(vDateTime.date)
                    || (vdate == cdiDate_get(vDateTime.date) && vtime < cdiTime_get(vDateTime.time)))
                  {
                    nextFileID = fileID;
                    vDateTime = streamInfoList[fileID].vDateTime;
                  }
              }
          }

        auto fileID = nextFileID;
        if (Options::cdoVerbose) cdo_print("nextstep = %d  vDateTime = %s", fileID, datetime_to_string(vDateTime));
        if (fileID == -1) break;

        auto &si = streamInfoList[fileID];

        if (skipSameTime && cdiDateTime_isEQ(vDateTime, lastDateTime))
          {
            cdo_print("Timestep %4d in stream %d (%s) already exists, skipped!", si.tsID + 1, si.streamID->get_id(),
                      datetime_to_string(vDateTime));
            processTimestep = false;
          }

        if (processTimestep)
          {
            if (tsID2 == 0)
              {
                auto vlistID1 = (mapFlag == MapFlag::Undefined)
                                    ? si.vlistID
                                    : ((mapFlag == MapFlag::Intersect) ? streamInfoList[vlistFileIDmin].vlistID
                                                                       : streamInfoList[vlistFileIDmax].vlistID);
                vlistID2 = vlistDuplicate(vlistID1);
                auto taxisID1 = vlistInqTaxis(si.vlistID);
                taxisID2 = taxisDuplicate(taxisID1);
                vlistDefTaxis(vlistID2, taxisID2);
                cdo_def_vlist(streamID2, vlistID2);
              }

            lastDateTime = vDateTime;

            cdo_taxis_copy_timestep(taxisID2, si.taxisID);
            cdo_def_timestep(streamID2, tsID2);

            for (int fieldID = 0; fieldID < si.numFields; ++fieldID)
              {
                auto [varID, levelID] = cdo_inq_field(si.streamID);

                const auto &var = si.varList.vars[varID];
                if (tsID2 > 0 && si.tsID == 0 && var.isConstant) continue;

                auto varID2 = varID;
                if (mapFlag != MapFlag::Undefined)
                  {
                    auto it = si.mapOfVarIDs.find(varID);
                    if (it == si.mapOfVarIDs.end()) continue;
                    varID2 = it->second;
                  }

                cdo_def_field(streamID2, varID2, levelID);

                if (dataIsUnchanged) { cdo_copy_field(streamID2, si.streamID); }
                else
                  {
                    field.init(var);
                    cdo_read_field(si.streamID, field);
                    cdo_write_field(streamID2, field);
                  }
              }

            if (mapFlag == MapFlag::Left) fill_missing_fields(streamInfoList[vlistFileIDmax], si);

            tsID2++;
          }

        si.numFields = cdo_stream_inq_timestep(si.streamID, ++si.tsID);
        if (si.numFields == 0)
          {
            cdo_stream_close(si.streamID);
            si.streamID = CDO_STREAM_UNDEF;
          }
        else { si.vDateTime = taxisInqVdatetime(si.taxisID); }
      }
  }

  void
  close() override
  {
    cdo_stream_close(streamID2);
    if (vlistID2 != CDI_UNDEFID) vlistDestroy(vlistID2);
    if (taxisID2 != CDI_UNDEFID) taxisDestroy(taxisID2);
  }
};
