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

      Setpartab  setpartab       Set parameter table
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "table.h"
#include "param_conversion.h"
#include "cdo_cmor.h"
#include "pmlist.h"
#include "convert_units.h"
#include "util_files.h"
#include "cdi_lockedIO.h"
#include "parse_literals.h"

enum pt_mode_t
{
  CODE_NUMBER,
  PARAMETER_ID,
  VARIABLE_NAME,
  STANDARD_NAME
};

void cdo_define_var_units(CmorVar &cmorVar, int vlistID2, int varID, const std::string &units);

void
mapvar(int vlistID, int varID, const KeyValues &kv, CmorVar &cmorVar, bool &hasValidMin, bool &hasValidMax, int ptab,
       bool isnPtmodeName)
{
  const auto key = string_to_lower(kv.key);
  const auto &value = kv.values[0];
  auto lv1 = (kv.nvalues == 1);

  // printf("key=%s  value=>%s<\n", key.c_str(), value.c_str());

  // clang-format off
  if (cmorVar.name == "cdocmor" )
    {
      if ((key == "cn") || (key == "cmor_name"))
        {
          auto name = cdo::inq_var_name(vlistID, varID);
          if (name[0] != 0) cdiDefAttTxt(vlistID, varID, "original_name", (int) name.size(), name.c_str());
          cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, parameter_to_word(value.c_str()));
        }
      else if (key == "character_axis")
        cdiDefAttTxt(vlistID, varID, "character_axis", (int) value.size(), value.c_str());
      else if (key == "z_axis")
        cdiDefAttTxt(vlistID, varID, "z_axis", (int) value.size(), value.c_str());
      else if (key == "variable_comment")
        cdiDefAttTxt(vlistID, varID, "variable_comment", (int) value.size(), value.c_str());
      else if (key == "positive")
        {
          if (!isspace(value[0])) cdiDefAttTxt(vlistID, varID, "positive", (int) value.size(), value.c_str());
        }
      else
        {
          if (Options::cdoVerbose) cdo_print("In applying the mapping table:\n          Key: '%s' is ignored.", key);
        }
    }
  if (key == "standard_name")
    {
      if (not lv1) cdo_abort("%s can only have one string value!", key);
      cdiDefKeyString(vlistID, varID, CDI_KEY_STDNAME, value.c_str());
    }
  else if (key == "long_name")
    {
      if (not lv1) cdo_abort("%s can only have one string value!", key);
      cdiDefKeyString(vlistID, varID, CDI_KEY_LONGNAME, value.c_str());
    }
  else if (key == "units")
    {
      if (not lv1) cdo_abort("%s can only have one string value!", key);
      cdo_define_var_units(cmorVar, vlistID, varID, value);
    }
  else if (key == "filterspec")
    {
      if (not lv1) cdo_abort("%s can only have one string value!", key);
      cdiDefKeyString(vlistID, varID, CDI_KEY_FILTERSPEC, value.c_str());
    }
  else if (key == "name")
    {
      if (isnPtmodeName)
        {
          if (not lv1) cdo_abort("%s can only have one string value!", key);
          cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, parameter_to_word(value.c_str()));
        }
    }
  else if (key == "out_name")
    {
      if (not lv1) cdo_abort("%s can only have one string value!", key);
      auto outname = parameter_to_word(value);
      if (cmorVar.name != outname)
        {
          cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, outname.c_str());
          cdiDefAttTxt(vlistID, varID, "original_name", (int) cmorVar.name.size(), cmorVar.name.c_str());
        }
    }
  else if (lv1 && key == "param")
    vlistDefVarParam(vlistID, varID, string_to_param(parameter_to_word(value)));
  else if (lv1 && key == "out_param")
    vlistDefVarParam(vlistID, varID, string_to_param(parameter_to_word(value)));
  else if (lv1 && key == "code")
    vlistDefVarParam(vlistID, varID, cdiEncodeParam(parameter_to_int(value), ptab, 255));
  else if (lv1 && key == "out_code")
    vlistDefVarParam(vlistID, varID, cdiEncodeParam(parameter_to_int(value), ptab, 255));
  else if (lv1 && key == "uvRelativeToGrid")
    cdiDefKeyInt(vlistID, varID, CDI_KEY_UVRELATIVETOGRID, parameter_to_bool(value));
  else if (lv1 && key == "comment")
    cdiDefAttTxt(vlistID, varID, key.c_str(), (int) value.size(), value.c_str());
  else if (lv1 && key == "chunktype")
    ;
  else if (lv1 && key == "cell_methods")
    cdiDefAttTxt(vlistID, varID, key.c_str(), (int) value.size(), value.c_str());
  else if (lv1 && key == "cell_measures")
    cdiDefAttTxt(vlistID, varID, key.c_str(), (int) value.size(), value.c_str());
  else if (lv1 && key == "delete")
    cmorVar.remove = parameter_to_bool(value);
  else if (lv1 && key == "convert")
    cmorVar.convert = parameter_to_bool(value);
  else if (lv1 && key == "factor")
    {
      cmorVar.applyFactor = true;
      cmorVar.factor = parameter_to_double(value);
      if (Options::cdoVerbose) cdo_print("%s - scale factor %g", cmorVar.name, cmorVar.factor);
    }
  else if (lv1 && (key == "missval" || key == "missing_value"))
    {
      auto missval = parameter_to_double(value);
      auto missvalOld = vlistInqVarMissval(vlistID, varID);
      if (!fp_is_equal(missval, missvalOld))
        {
          if (Options::cdoVerbose) cdo_print("%s - change missval from %g to %g", cmorVar.name, missvalOld, missval);
          cmorVar.changeMissval = true;
          cmorVar.missvalOld = missvalOld;
          vlistDefVarMissval(vlistID, varID, missval);
        }
    }
  else if (lv1 && key == "valid_min")
    {
      hasValidMin = true;
      cmorVar.valid_min = parameter_to_double(value);
    }
  else if (lv1 && key == "valid_max")
    {
      hasValidMax = true;
      cmorVar.valid_max = parameter_to_double(value);
    }
  else if (lv1 && key == "ok_min_mean_abs")
    {
      cmorVar.check_min_mean_abs = true;
      cmorVar.ok_min_mean_abs = parameter_to_double(value);
    }
  else if (lv1 && key == "ok_max_mean_abs")
    {
      cmorVar.check_max_mean_abs = true;
      cmorVar.ok_max_mean_abs = parameter_to_double(value);
    }
  else if (lv1 && (key == "datatype" || key == "type"))
    {
      auto datatype = cdo::str_to_datatype(parameter_to_word(value));
      if (datatype != -1) vlistDefVarDatatype(vlistID, varID, datatype);
    }
  else if (lv1 && key == "dimensions") {}
  else
    {
      const auto &values = kv.values;
      const auto &rvalue = kv.values[0];
      auto nvalues = kv.nvalues;
      if (nvalues == 1 && rvalue.empty()) nvalues = 0;

      auto dtype = literals_find_datatype(nvalues, values);

      if (dtype == CDI_DATATYPE_INT8 || dtype == CDI_DATATYPE_INT16 || dtype == CDI_DATATYPE_INT32)
        {
          std::vector<int> ivals(nvalues);
          for (int i = 0; i < nvalues; ++i) ivals[i] = literal_to_int(values[i]);
          cdiDefAttInt(vlistID, varID, key.c_str(), dtype, nvalues, ivals.data());
        }
      else if (dtype == CDI_DATATYPE_FLT32 || dtype == CDI_DATATYPE_FLT64)
        {
          std::vector<double> dvals(nvalues);
          for (int i = 0; i < nvalues; ++i) dvals[i] = literal_to_double(values[i]);
          cdiDefAttFlt(vlistID, varID, key.c_str(), dtype, nvalues, dvals.data());
        }
      else { cdiDefAttTxt(vlistID, varID, key.c_str(), (int) rvalue.size(), rvalue.c_str()); }
    }
  // clang-format on
}

static void
search_global_missval(PMList &pmlist, bool &hasMissvals, double &missval)
{
  const std::vector<std::string> hentry = { "Header" };
  auto kvlist = pmlist.getKVListVentry(hentry);
  if (kvlist)
    {
      auto kv = kvlist->search("missing_value");
      if (kv && kv->nvalues > 0)
        {
          hasMissvals = true;
          missval = parameter_to_double(kv->values[0]);
        }
    }
}

static void
apply_parameterList(pt_mode_t ptmode, PMList &pmlist, int vlistID2, std::vector<CmorVar> &cmorVars)
{
  // search for global missing value
  auto hasMissvals = false;
  double missval = 0.0;
  search_global_missval(pmlist, hasMissvals, missval);

  const std::vector<std::string> ventry = { "variable_entry", "parameter" };
  char valstr[CDI_MAX_NAME];
  char paramstr[32];
  int codenum = 0;

  VarList varList2(vlistID2);

  int numVarsFound = 0;
  int numVars = cmorVars.size();
  for (int varID = 0; varID < numVars; ++varID)
    {
      auto &cmorVar = cmorVars[varID];
      cmorVar.name = cdo::inq_var_name(vlistID2, varID);

      const auto &var2 = varList2.vars[varID];
      if (hasMissvals)
        {
          if (fp_is_not_equal(missval, var2.missval))
            {
              cmorVar.changeMissval = true;
              cmorVar.missvalOld = var2.missval;
              vlistDefVarMissval(vlistID2, varID, missval);
            }
        }

      const KVList *kvlist = nullptr;
      if (ptmode == CODE_NUMBER)
        {
          codenum = var2.code;
          std::snprintf(valstr, sizeof(valstr), "%d", codenum);
          kvlist = pmlist.searchKVListVentry("code", valstr, ventry);
          if (kvlist)
            {
              auto tableID = vlistInqVarTable(vlistID2, varID);
              auto tabnum = tableInqNum(tableID);
              int levtype = 0;
              cdiInqKeyInt(var2.zaxisID, CDI_GLOBAL, CDI_KEY_TYPEOFFIRSTFIXEDSURFACE, &levtype);
              auto table = tabnum;
              auto ltype = levtype;
              {
                auto kv = kvlist->search("table");
                if (kv && kv->nvalues == 1) table = parameter_to_int(kv->values[0]);
              }
              {
                auto kv = kvlist->search("ltype");
                if (kv && kv->nvalues == 1) ltype = parameter_to_int(kv->values[0]);
              }
              if (!(tabnum == table && levtype == ltype)) kvlist = nullptr;
            }
        }
      else if (ptmode == PARAMETER_ID)
        {
          auto param = var2.param;
          cdiParamToString(param, paramstr, sizeof(paramstr));
          std::snprintf(valstr, sizeof(valstr), "%s", paramstr);
          kvlist = pmlist.searchKVListVentry("param", valstr, ventry);
          if (kvlist)
            {
              int levtype = 0;
              cdiInqKeyInt(var2.zaxisID, CDI_GLOBAL, CDI_KEY_TYPEOFFIRSTFIXEDSURFACE, &levtype);
              auto kv = kvlist->search("ltype");
              auto ltype = (kv && kv->nvalues == 1) ? parameter_to_int(kv->values[0]) : levtype;
              if (levtype != ltype) kvlist = nullptr;
            }
        }
      else if (ptmode == VARIABLE_NAME) { kvlist = pmlist.searchKVListVentry("name", cmorVar.name, ventry); }

      if (kvlist)
        {
          numVarsFound++;
          int pnum, ptab, pdum;
          cdiDecodeParam(var2.param, &pnum, &ptab, &pdum);

          auto hasValidMin = false;
          auto hasValidMax = false;
          for (const auto &kv : *kvlist)
            {
              mapvar(vlistID2, varID, kv, cmorVar, hasValidMin, hasValidMax, ptab, (ptmode != VARIABLE_NAME));
            }
          if (hasValidMin && hasValidMax) cmorVar.checkValid = true;
        }
      else if (Options::cdoVerbose)
        {
          // clang-format off
          if      (ptmode == CODE_NUMBER)   cdo_print("Code number %d not found in parameter table!", codenum);
          else if (ptmode == PARAMETER_ID)  cdo_print("Parameter ID %s not found in parameter table!", paramstr);
          else if (ptmode == VARIABLE_NAME) cdo_print("Variable %s not found in parameter table!", cmorVar.name);
          // clang-format on
        }
    }

  if (numVarsFound == 0)
    {
      // clang-format off
      if      (ptmode == CODE_NUMBER)   cdo_warning("None of the input variables has a code number that matches the entries in the parameter table!");
      else if (ptmode == PARAMETER_ID)  cdo_warning("None of the input variables has a parameter ID that matches the entries in the parameter table!");
      else if (ptmode == VARIABLE_NAME) cdo_warning("None of the input variables has a name that matches the entries in the parameter table!");
      // clang-format on
    }
}

static void
apply_codetable(int vlistID2, int tableID)
{
  char name[CDI_MAX_NAME], longname[CDI_MAX_NAME], units[CDI_MAX_NAME];
  auto numVars = vlistNvars(vlistID2);
  for (int varID = 0; varID < numVars; ++varID)
    {
      auto param = vlistInqVarParam(vlistID2, varID);
      int pdis, pcat, pnum;
      cdiDecodeParam(param, &pnum, &pcat, &pdis);
      if (pdis == 255)
        {
          auto code = pnum;
          int ltype = 0;
          cdiInqKeyInt(vlistInqVarZaxis(vlistID2, varID), CDI_GLOBAL, CDI_KEY_TYPEOFFIRSTFIXEDSURFACE, &ltype);
          name[0] = 0;
          longname[0] = 0;
          units[0] = 0;
          tableInqEntry(tableID, code, ltype, name, longname, units);
          if (name[0])
            {
              cdiDefKeyString(vlistID2, varID, CDI_KEY_NAME, name);
              if (longname[0]) cdiDefKeyString(vlistID2, varID, CDI_KEY_LONGNAME, longname);
              if (units[0]) cdiDefKeyString(vlistID2, varID, CDI_KEY_UNITS, units);
            }
        }
      vlistDefVarTable(vlistID2, varID, tableID);
    }
}

class Setpartab : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Setpartab",
    .operators = { { "setcodetab", 0, 0, "parameter code table name", SetHelp },
                   { "setpartabc", 0, 0, "parameter table name", SetpartabHelp },
                   { "setpartabp", 0, 0, "parameter table name", SetpartabHelp },
                   { "setpartabn", 0, 0, "parameter table name", SetpartabHelp } },
    .aliases = { { "setpartab", "setcodetab" }, { "setpartabv", "setpartabn" } },
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Setpartab> registration = RegisterEntry<Setpartab>(module);

  int tableID = -1;
  int tableformat = 0;
  bool deleteVars = false;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID1;
  int taxisID2;

  int vlistID1;
  int vlistID2;

  std::vector<CmorVar> cmorVars;

  VarList varList2;

public:
  void
  init() override
  {
    auto SETCODETAB = module.get_id("setcodetab");
    auto SETPARTABC = module.get_id("setpartabc");
    auto SETPARTABP = module.get_id("setpartabp");
    auto SETPARTABN = module.get_id("setpartabn");

    auto operatorID = cdo_operator_id();

    operator_input_arg(cdo_operator_enter(operatorID));

    if (cdo_operator_argc() < 1) cdo_abort("Too few arguments!");

    auto convertData = false;
    if (cdo_operator_argc() == 2)
      {
        if (cdo_operator_argv(1) == "convert")
          convertData = true;
        else
          cdo_abort("Unknown parameter: >%s<", cdo_operator_argv(1));
      }

    if (cdo_operator_argc() > 2) cdo_abort("Too many arguments!");

    pt_mode_t ptmode = CODE_NUMBER;
    // clang-format off
    if      (operatorID == SETCODETAB) ptmode = CODE_NUMBER;
    else if (operatorID == SETPARTABC) ptmode = CODE_NUMBER;
    else if (operatorID == SETPARTABP) ptmode = PARAMETER_ID;
    else if (operatorID == SETPARTABN) ptmode = VARIABLE_NAME;
    // clang-format on

    if (ptmode == CODE_NUMBER)
      {
        const auto &partab = cdo_operator_argv(0);
        FILE *fp = FileUtils::file_exists(partab) ? std::fopen(partab.c_str(), "r") : nullptr;
        if (fp != nullptr)
          {
            fseek(fp, 0L, SEEK_END);
            auto fsize = (size_t) ftell(fp);
            std::vector<char> parbuf(fsize + 1);
            fseek(fp, 0L, SEEK_SET);
            fread(parbuf.data(), fsize, 1, fp);
            parbuf[fsize] = 0;
            fseek(fp, 0L, SEEK_SET);

            if (atoi(parbuf.data()) == 0) tableformat = 1;

            std::fclose(fp);
          }

        if (tableformat == 0) tableID = cdo::define_table(partab);
      }
    else if (ptmode == PARAMETER_ID) { tableformat = 1; }
    else if (ptmode == VARIABLE_NAME) { tableformat = 1; }

    if (Options::cdoVerbose) cdo_print("Table format version %d", tableformat);

    streamID1 = cdo_open_read(0);

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

    auto numVars = vlistNvars(vlistID2);
    cmorVars.resize(numVars);

    if (convertData)
      for (auto &var : cmorVars) var.convert = true;

    if (tableformat == 0)
      {
        // for (int varID = 0; varID < numVars; ++varID) vlistDefVarTable(vlistID2, varID, tableID);
        apply_codetable(vlistID2, tableID);
      }
    else
      {
        {
          auto filename = cdo_operator_argv(0).c_str();
          auto fp = std::fopen(filename, "r");
          if (fp == nullptr) cdo_abort("Open failed on: %s\n", filename);

          PMList pmlist;
          pmlist.read_namelist(fp, filename);
          std::fclose(fp);

          apply_parameterList(ptmode, pmlist, vlistID2, cmorVars);
        }

        for (auto &var : cmorVars)
          if (var.remove)
            {
              deleteVars = true;
              break;
            }

        if (deleteVars)
          {
            vlistClearFlag(vlistID1);
            vlistClearFlag(vlistID2);

            for (int varID = 0; varID < numVars; ++varID)
              {
                auto zaxisID = vlistInqVarZaxis(vlistID2, varID);
                auto numLevels = zaxisInqSize(zaxisID);
                for (int levelID = 0; levelID < numLevels; levelID++)
                  {
                    vlistDefFlag(vlistID1, varID, levelID, true);
                    vlistDefFlag(vlistID2, varID, levelID, true);
                    if (cmorVars[varID].remove)
                      {
                        vlistDefFlag(vlistID1, varID, levelID, false);
                        vlistDefFlag(vlistID2, varID, levelID, false);
                      }
                  }
              }

            auto vlistIDx = vlistCreate();
            cdo_vlist_copy_flag(vlistIDx, vlistID2);

            vlistDestroy(vlistID2);

            vlistID2 = vlistIDx;
            if (vlistNvars(vlistID2) == 0) cdo_abort("No variable selected!");
          }

        for (auto &var : cmorVars)
          {
            if (!var.convert) var.changeUnits = false;
            if (var.changeUnits)
              cdo::convert_units(&var.ut_converter, &var.changeUnits, (char *) &var.units, (char *) &var.unitsOld, var.name);
          }
      }

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

    // vlistPrint(vlistID2);
    streamID2 = cdo_open_write(1);
    cdo_def_vlist(streamID2, vlistID2);

    varList2 = VarList(vlistID2);
  }

  void
  run() override
  {
    auto gridsizemax = vlistGridsizeMax(vlistID1);
    if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;
    Varray<double> array(gridsizemax);

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

        cdo_taxis_copy_timestep(taxisID2, taxisID1);
        cdo_def_timestep(streamID2, tsID);

        cmor_check_init(cmorVars);

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

            auto &cmorVar = cmorVars[varID];
            auto varID2 = varID;
            auto levelID2 = levelID;

            if (deleteVars)
              {
                if (cmorVar.remove) continue;

                if (vlistInqFlag(vlistID1, varID, levelID) == true)
                  {
                    varID2 = vlistFindVar(vlistID2, varID);
                    levelID2 = vlistFindLevel(vlistID2, varID, levelID);
                  }
              }

            cdo_def_field(streamID2, varID2, levelID2);

            size_t numMissVals;
            cdo_read_field(streamID1, array.data(), &numMissVals);

            const auto &var2 = varList2.vars[varID2];
            auto missval = var2.missval;
            auto gridsize = var2.nwpv * var2.gridsize;

            if (numMissVals && cmorVar.changeMissval)
              {
                for (size_t i = 0; i < gridsize; ++i)
                  {
                    if (fp_is_equal(array[i], cmorVar.missvalOld)) array[i] = missval;
                  }
              }

            if (cmorVar.applyFactor)
              {
                for (size_t i = 0; i < gridsize; ++i)
                  {
                    if (fp_is_not_equal(array[i], missval)) array[i] *= cmorVar.factor;
                  }
              }

#ifdef HAVE_UDUNITS2
            if (cmorVar.changeUnits)
              {
                int nerr = 0;
                for (size_t i = 0; i < gridsize; ++i)
                  {
                    if (fp_is_not_equal(array[i], missval))
                      {
                        array[i] = cv_convert_double((const cv_converter *) cmorVar.ut_converter, array[i]);
                        if (ut_get_status() != UT_SUCCESS) nerr++;
                      }
                  }
                if (nerr)
                  {
                    cdo_warning("Udunits: Error converting units from [%s] to [%s], parameter: %s", cmorVar.unitsOld, cmorVar.units,
                                cmorVar.name);
                    cmorVar.changeUnits = false;
                  }
              }
#endif

            cdo_write_field(streamID2, array.data(), numMissVals);

            cmor_check_prep(cmorVar, gridsize, missval, array.data());
          }

        cmor_check_eval(vlistID2, cmorVars);

        tsID++;
      }
  }

  void
  close() override
  {
    cdo_stream_close(streamID2);
    cdo_stream_close(streamID1);

#ifdef HAVE_UDUNITS2
    for (auto &var : cmorVars)
      if (var.changeUnits) cdo::convert_free(var.ut_converter);

    cdo::convert_destroy();
#endif
  }
};
