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

      Change     chcode          Change code number
      Change     chtabnum        Change GRIB1 parameter table number
      Change     chparam         Change parameter identifier
      Change     chname          Change variable or coordinate name
      Change     chlevel         Change level
      Change     chlevelc        Change level of one code
      Change     chlevelv        Change level of one variable
      Change     chltype         Change GRIB level type
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "param_conversion.h"

static void
change_code(const VarList &varList1, int vlistID2, int nch, const std::vector<int> &chints)
{
  for (const auto &var1 : varList1.vars)
    {
      for (int i = 0; i < nch; i += 2)
        if (var1.code == chints[i]) vlistDefVarCode(vlistID2, var1.ID, chints[i + 1]);
    }
}

static void
change_tabnum(const VarList &varList1, int vlistID2, int nch, const std::vector<int> &chints)
{
  for (const auto &var1 : varList1.vars)
    {
      auto tabnum = tableInqNum(vlistInqVarTable(vlistID2, var1.ID));
      for (int i = 0; i < nch; i += 2)
        if (tabnum == chints[i])
          {
            auto tableID = tableDef(-1, chints[i + 1], nullptr);
            vlistDefVarTable(vlistID2, var1.ID, tableID);
          }
    }
}

static void
change_param(const VarList &varList1, int vlistID2, int nch, const std::vector<const char *> &chnames)
{
  for (const auto &var1 : varList1.vars)
    {
      if (Options::cdoVerbose)
        {
          int pnum, pcat, pdis;
          cdiDecodeParam(var1.param, &pnum, &pcat, &pdis);
          cdo_print("pnum, pcat, pdis: %d.%d.%d", pnum, pcat, pdis);
        }
      for (int i = 0; i < nch; i += 2)
        if (var1.param == string_to_param(chnames[i])) vlistDefVarParam(vlistID2, var1.ID, string_to_param(chnames[i + 1]));
    }
}

static void
change_name(const VarList &varList1, int vlistID2, int nch, const std::vector<const char *> &chnames)
{
  auto npairs = nch / 2;
  std::vector<std::pair<const char *, const char *>> vpairs(npairs);
  for (int i = 0; i < npairs; ++i) vpairs[i].first = chnames[i * 2];
  for (int i = 0; i < npairs; ++i) vpairs[i].second = chnames[i * 2 + 1];

  std::vector<bool> namefound(npairs, false);
  for (const auto &var1 : varList1.vars)
    {
      for (int i = 0; i < npairs; ++i)
        if (var1.name == vpairs[i].first)
          {
            namefound[i] = true;
            cdiDefKeyString(vlistID2, var1.ID, CDI_KEY_NAME, vpairs[i].second);
            break;
          }
    }

  auto searchForGridName = false;
  for (int i = 0; i < npairs; ++i)
    if (!namefound[i])
      {
        searchForGridName = true;
        break;
      }

  if (searchForGridName)
    {
      auto ngrids = vlistNumGrids(vlistID2);
      for (int index = 0; index < ngrids; ++index)
        {
          int gridID2 = -1;
          auto gridID1 = vlistGrid(vlistID2, index);
          auto xname = cdo::inq_key_string(gridID1, CDI_XAXIS, CDI_KEY_NAME);
          auto yname = cdo::inq_key_string(gridID1, CDI_YAXIS, CDI_KEY_NAME);
          auto xfound = false, yfound = false;
          for (int i = 0; i < npairs; ++i)
            {
              if (!namefound[i])
                {
                  if (xname == vpairs[i].first)
                    {
                      xfound = true;
                      namefound[i] = true;
                      if (gridID2 == -1) gridID2 = gridDuplicate(gridID1);
                      cdiDefKeyString(gridID2, CDI_XAXIS, CDI_KEY_NAME, vpairs[i].second);
                    }
                }
              if (!namefound[i])
                {
                  if (yname == vpairs[i].first)
                    {
                      yfound = true;
                      namefound[i] = true;
                      if (gridID2 == -1) gridID2 = gridDuplicate(gridID1);
                      cdiDefKeyString(gridID2, CDI_YAXIS, CDI_KEY_NAME, vpairs[i].second);
                    }
                }

              if (xfound && yfound) break;
            }

          if (gridID2 != -1) vlistChangeGrid(vlistID2, gridID1, gridID2);
        }
    }

  auto searchForZaxisName = false;
  for (int i = 0; i < npairs; ++i)
    if (!namefound[i])
      {
        searchForZaxisName = true;
        break;
      }

  if (searchForZaxisName)
    {
      auto nzaxis = vlistNumZaxis(vlistID2);
      for (int index = 0; index < nzaxis; ++index)
        {
          auto zaxisID1 = vlistZaxis(vlistID2, index);
          auto varname = cdo::inq_key_string(zaxisID1, CDI_GLOBAL, CDI_KEY_NAME);
          for (int i = 0; i < npairs; ++i)
            {
              if (!namefound[i])
                {
                  if (varname == vpairs[i].first)
                    {
                      namefound[i] = true;
                      auto zaxisID2 = zaxisDuplicate(zaxisID1);
                      cdiDefKeyString(zaxisID2, CDI_GLOBAL, CDI_KEY_NAME, vpairs[i].second);
                      vlistChangeZaxis(vlistID2, zaxisID1, zaxisID2);
                      break;
                    }
                }
            }
        }
    }

  for (int i = 0; i < npairs; ++i)
    if (!namefound[i]) cdo_warning("Variable name %s not found!", vpairs[i].first);
}

static void
change_unit(const VarList &varList1, int vlistID2, int nch, const std::vector<const char *> &chnames)
{
  for (const auto &var1 : varList1.vars)
    {
      for (int i = 0; i < nch; i += 2)
        if (var1.units == chnames[i]) cdiDefKeyString(vlistID2, var1.ID, CDI_KEY_UNITS, chnames[i + 1]);
    }
}

static void
change_level(int vlistID2, int nch, const std::vector<double> &chlevels)
{
  auto nzaxis = vlistNumZaxis(vlistID2);
  for (int index = 0; index < nzaxis; ++index)
    {
      auto zaxisID1 = vlistZaxis(vlistID2, index);
      if (zaxisInqLevels(zaxisID1, nullptr))
        {
          auto nlevs = zaxisInqSize(zaxisID1);
          Varray<double> levels(nlevs);
          zaxisInqLevels(zaxisID1, &levels[0]);

          int nfound = 0;
          for (int i = 0; i < nch; i += 2)
            for (int k = 0; k < nlevs; ++k)
              if (std::fabs(levels[k] - chlevels[i]) < 0.0001) nfound++;

          if (nfound)
            {
              Varray<double> newlevels = levels;
              auto zaxisID2 = zaxisDuplicate(zaxisID1);
              for (int i = 0; i < nch; i += 2)
                for (int k = 0; k < nlevs; ++k)
                  if (std::fabs(levels[k] - chlevels[i]) < 0.001) newlevels[k] = chlevels[i + 1];

              zaxisDefLevels(zaxisID2, &newlevels[0]);
              vlistChangeZaxis(vlistID2, zaxisID1, zaxisID2);
            }
        }
    }
}

static void
change_varLevel(int varID, int vlistID2, const std::vector<double> &chlevels)
{
  auto zaxisID1 = vlistInqVarZaxis(vlistID2, varID);
  if (zaxisInqLevels(zaxisID1, nullptr))
    {
      auto nlevs = zaxisInqSize(zaxisID1);
      Varray<double> levels(nlevs);
      zaxisInqLevels(zaxisID1, &levels[0]);

      int nfound = 0;
      for (int k = 0; k < nlevs; ++k)
        if (std::fabs(levels[k] - chlevels[0]) < 0.0001) nfound++;

      if (nfound)
        {
          auto zaxisID2 = zaxisDuplicate(zaxisID1);
          for (int k = 0; k < nlevs; ++k)
            if (std::fabs(levels[k] - chlevels[0]) < 0.001) levels[k] = chlevels[1];

          zaxisDefLevels(zaxisID2, &levels[0]);
          vlistChangeVarZaxis(vlistID2, varID, zaxisID2);
        }
      else
        cdo_abort("Level %g not found!", chlevels[0]);
    }
}

static void
change_levelByCode(int chcode, const VarList &varList1, int vlistID2, const std::vector<double> &chlevels)
{
  for (const auto &var1 : varList1.vars)
    {
      if (var1.code == chcode)
        {
          change_varLevel(var1.ID, vlistID2, chlevels);
          return;
        }
    }

  cdo_abort("Code %d not found!", chcode);
}

static void
change_levelByName(const char *chname, const VarList &varList1, int vlistID2, const std::vector<double> &chlevels)
{
  for (const auto &var1 : varList1.vars)
    {
      if (var1.name == chname)
        {
          change_varLevel(var1.ID, vlistID2, chlevels);
          return;
        }
    }

  cdo_abort("Variable name %s not found!", chname);
}

static void
change_ltype(int vlistID2, int nch, const std::vector<int> &chltypes)
{
  auto nzaxis = vlistNumZaxis(vlistID2);
  for (int index = 0; index < nzaxis; ++index)
    {
      auto zaxisID1 = vlistZaxis(vlistID2, index);
      auto zaxisID2 = zaxisDuplicate(zaxisID1);
      int ltype = 0;
      cdiInqKeyInt(zaxisID1, CDI_GLOBAL, CDI_KEY_TYPEOFFIRSTFIXEDSURFACE, &ltype);

      for (int i = 0; i < nch; i += 2)
        {
          auto ltype1 = chltypes[i];
          auto ltype2 = chltypes[i + 1];
          if (ltype1 == ltype)
            {
              zaxisChangeType(zaxisID2, ZAXIS_GENERIC);
              cdiDefKeyInt(zaxisID2, CDI_GLOBAL, CDI_KEY_TYPEOFFIRSTFIXEDSURFACE, ltype2);
              vlistChangeZaxis(vlistID2, zaxisID1, zaxisID2);
            }
        }
    }
}

class Change : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Change",
    .operators = { { "chcode", 0, 0, "pairs of old and new code numbers", ChangeHelp },
                   { "chtabnum", 0, 0, "pairs of old and new GRIB1 table numbers", ChangeHelp },
                   { "chparam", 0, 0, "pairs of old and new parameter identifiers", ChangeHelp },
                   { "chname", 0, 0, "pairs of old and new variable names", ChangeHelp },
                   { "chunit", 0, 0, "pairs of old and new variable units", ChangeHelp },
                   { "chlevel", 0, 0, "pairs of old and new levels", ChangeHelp },
                   { "chlevelc", 0, 0, "code number, old and new level", ChangeHelp },
                   { "chlevelv", 0, 0, "variable name, old and new level", ChangeHelp },
                   { "chltype", 0, 0, "pairs of old and new level type", ChangeHelp } },
    .aliases = { { "chvar", "chname" } },
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Change> registration = RegisterEntry<Change>(module);

  int CHCODE, CHTABNUM, CHPARAM, CHNAME, CHUNIT, CHLEVEL, CHLEVELC, CHLEVELV, CHLTYPE;

  CdoStreamID streamID1;
  int taxisID1;

  CdoStreamID streamID2;
  int taxisID2;

  VarList varList1;

public:
  void
  init() override
  {
    CHCODE = module.get_id("chcode");
    CHTABNUM = module.get_id("chtabnum");
    CHPARAM = module.get_id("chparam");
    CHNAME = module.get_id("chname");
    CHUNIT = module.get_id("chunit");
    CHLEVEL = module.get_id("chlevel");
    CHLEVELC = module.get_id("chlevelc");
    CHLEVELV = module.get_id("chlevelv");
    CHLTYPE = module.get_id("chltype");

    auto operatorID = cdo_operator_id();

    operator_input_arg(cdo_operator_enter(operatorID));

    auto nch = cdo_operator_argc();

    const char *chname = nullptr;
    int chcode = 0;
    std::vector<const char *> chnames;
    std::vector<int> chints, chltypes;
    std::vector<double> chlevels;

    if (operatorID == CHCODE || operatorID == CHTABNUM)
      {
        if (nch % 2) cdo_abort("Odd number of input arguments!");
        chints.resize(nch);
        for (int i = 0; i < nch; ++i) chints[i] = parameter_to_int(cdo_operator_argv(i));
      }
    else if (operatorID == CHPARAM || operatorID == CHNAME || operatorID == CHUNIT)
      {
        if (nch % 2) cdo_abort("Odd number of input arguments!");
        chnames.resize(nch);
        for (int i = 0; i < nch; ++i) chnames[i] = &cdo_operator_argv(i)[0];
      }
    else if (operatorID == CHLEVEL)
      {
        if (nch % 2) cdo_abort("Odd number of input arguments!");
        chlevels.resize(nch);
        for (int i = 0; i < nch; ++i) chlevels[i] = parameter_to_double(cdo_operator_argv(i));
      }
    else if (operatorID == CHLEVELC)
      {
        operator_check_argc(3);

        chcode = parameter_to_int(cdo_operator_argv(0));
        chlevels.resize(2);
        chlevels[0] = parameter_to_double(cdo_operator_argv(1));
        chlevels[1] = parameter_to_double(cdo_operator_argv(2));
      }
    else if (operatorID == CHLEVELV)
      {
        operator_check_argc(3);

        chname = cdo_operator_argv(0).c_str();
        chlevels.resize(2);
        chlevels[0] = parameter_to_double(cdo_operator_argv(1));
        chlevels[1] = parameter_to_double(cdo_operator_argv(2));
      }
    else if (operatorID == CHLTYPE)
      {
        if (nch % 2) cdo_abort("Odd number of input arguments!");
        chltypes.resize(nch);
        for (int i = 0; i < nch; ++i) chltypes[i] = parameter_to_int(cdo_operator_argv(i));
      }

    streamID1 = cdo_open_read(0);

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

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

    varList1 = VarList(vlistID1);

    // clang-format off
    if      (operatorID == CHCODE)   change_code(varList1, vlistID2, nch, chints);
    else if (operatorID == CHTABNUM) change_tabnum(varList1, vlistID2, nch, chints);
    else if (operatorID == CHPARAM)  change_param(varList1, vlistID2, nch, chnames);
    else if (operatorID == CHNAME)   change_name(varList1, vlistID2, nch, chnames);
    else if (operatorID == CHUNIT)   change_unit(varList1, vlistID2, nch, chnames);
    else if (operatorID == CHLEVEL)  change_level(vlistID2, nch, chlevels);
    else if (operatorID == CHLEVELC) change_levelByCode(chcode, varList1, vlistID2, chlevels);
    else if (operatorID == CHLEVELV) change_levelByName(chname, varList1, vlistID2, chlevels);
    else if (operatorID == CHLTYPE)  change_ltype(vlistID2, nch, chltypes);
    // clang-format on

    streamID2 = cdo_open_write(1);
    cdo_def_vlist(streamID2, vlistID2);
  }

  void
  run() override
  {
    Field field;

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

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

            field.init(varList1.vars[varID]);
            cdo_read_field(streamID1, field);
            cdo_write_field(streamID2, field);
          }

        tsID++;
      }
  }

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