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

      Vertint    gh2hl           Model geometric height level to height level interpolation
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "field_vinterp.h"
#include "stdnametable.h"
#include "util_string.h"
#include "cdo_zaxis.h"
#include "param_conversion.h"
#include "vertint_util.h"

static bool
is_height_axis(int zaxisID)
{
  auto isHeight = false;
  if (zaxisInqType(zaxisID) == ZAXIS_REFERENCE)
    {
      auto units = cdo::inq_key_string(zaxisID, CDI_GLOBAL, CDI_KEY_UNITS);
      auto stdname = cdo::inq_key_string(zaxisID, CDI_GLOBAL, CDI_KEY_STDNAME);
      if (stdname == "height" && units.empty()) isHeight = true;
    }
  return isHeight;
}

static int
create_zaxis_height(Varray<double> &heightLevels)
{
  int zaxisID = CDI_UNDEFID;
  const auto &arg1 = cdo_operator_argv(0);
  if (cdo_operator_argc() == 1 && !std::isdigit(arg1[0]))
    {
      auto zfilename = arg1.c_str();
      auto zfp = std::fopen(zfilename, "r");
      if (zfp)
        {
          zaxisID = zaxis_from_file(zfp, zfilename);
          std::fclose(zfp);
          if (zaxisID == CDI_UNDEFID) cdo_abort("Invalid zaxis description file %s!", zfilename);
          auto nlevels = zaxisInqSize(zaxisID);
          heightLevels.resize(nlevels);
          zaxisInqLevels(zaxisID, heightLevels.data());
        }
      else if (arg1 == "default")
        heightLevels = { 10, 50, 100, 500, 1000, 5000, 10000, 15000, 20000, 25000, 30000 };
      else
        cdo_abort("Open failed on %s", zfilename);
    }
  else { heightLevels = cdo_argv_to_fltarr(cdo_get_oper_argv()); }

  if (zaxisID == CDI_UNDEFID)
    {
      zaxisID = zaxisCreate(ZAXIS_HEIGHT, heightLevels.size());
      zaxisDefLevels(zaxisID, heightLevels.data());
    }

  return zaxisID;
}

class Vertintgh : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Vertintgh",
    // clang-format off
    .operators = { { "gh2hl", 0, 0, "height levels in meter", VertintghHelp },
                   { "gh2hlx", 0, 0, "height levels in meter", VertintghHelp } },
    // clang-format on
    .aliases = {},
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_REAL,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Vertintgh> registration = RegisterEntry<Vertintgh>(module);

  int GH2HLX;

  CdoStreamID streamID1;
  int taxisID1;

  CdoStreamID streamID2;
  int taxisID2;

  size_t gridsize;
  int numVars;
  int numFullLevels;
  int numHalfLevels;

  bool extrapolate;

  std::vector<bool> processVars;
  std::vector<bool> interpVars;

  int heightID_FL = -1, heightID_HL = -1;

  VarList varList1;
  VarList varList2;
  Varray2D<size_t> varnumMissVals;
  Field3DVector vardata1, vardata2;

  Varray<double> heightLevels;

  Varray<size_t> numMiss_FL, numMiss_HL;
  std::vector<int> vertIndex_FL;
  std::vector<int> vertIndex_HL;

public:
  void
  init() override
  {
    GH2HLX = module.get_id("gh2hlx");

    auto operatorID = cdo_operator_id();

    extrapolate = (operatorID == GH2HLX);
    if (extrapolate == false) extrapolate = getenv_extrapolate();

    operator_input_arg(cdo_operator_enter(operatorID));

    auto zaxisID2 = create_zaxis_height(heightLevels);

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

    gridsize = vlist_check_gridsize(vlistID1);

    varList1 = VarList(vlistID1);
    varList_set_unique_memtype(varList1);
    auto memtype = varList1.vars[0].memType;

    auto stdnameHeight_FL = var_stdname(geometric_height_at_full_level_center);
    auto stdnameHeight_HL = var_stdname(geometric_height_at_half_level_center);

    numVars = varList1.numVars();
    for (int varID = 0; varID < numVars; ++varID)
      {
        auto stdname = string_to_lower(varList1.vars[varID].stdname);
        if (stdname == stdnameHeight_FL) heightID_FL = varID;
        if (stdname == stdnameHeight_HL) heightID_HL = varID;
      }

    if (-1 == heightID_FL && -1 == heightID_HL)
      {
        for (int varID = 0; varID < numVars; ++varID)
          {
            const auto &var = varList1.vars[varID];
            auto stdname = string_to_lower(var.stdname);
            if (stdname == "height" && var.units == "m")
              {
                // clang-format off
                if      (var.longname == "geometric height at full level center") heightID_FL = varID;
                else if (var.longname == "geometric height at half level center") heightID_HL = varID;
                // clang-format on
              }
          }
      }

    if (Options::cdoVerbose)
      {
        cdo_print("Found:");
        if (-1 != heightID_FL) cdo_print("  %s -> %s", stdnameHeight_FL, varList1.vars[heightID_FL].name);
        if (-1 != heightID_HL) cdo_print("  %s -> %s", stdnameHeight_HL, varList1.vars[heightID_HL].name);
      }

    if (-1 == heightID_FL && -1 == heightID_HL) cdo_abort("%s not found!", stdnameHeight_FL);

    auto zaxisID_FL = (-1 == heightID_FL) ? -1 : varList1.vars[heightID_FL].zaxisID;
    auto zaxisID_HL = (-1 == heightID_HL) ? -1 : varList1.vars[heightID_HL].zaxisID;
    numFullLevels = (-1 == zaxisID_FL) ? 0 : varList1.vars[heightID_FL].nlevels;
    numHalfLevels = (-1 == zaxisID_HL) ? 0 : varList1.vars[heightID_HL].nlevels;

    auto numZaxes = vlistNumZaxis(vlistID1);
    for (int index = 0; index < numZaxes; ++index)
      {
        auto zaxisID = vlistZaxis(vlistID1, index);
        auto nlevels = zaxisInqSize(zaxisID);
        if (zaxisID == zaxisID_FL || zaxisID == zaxisID_HL
            || (is_height_axis(zaxisID) && (nlevels == numHalfLevels || nlevels == numFullLevels)))
          vlistChangeZaxis(vlistID2, zaxisID, zaxisID2);
      }

    varList2 = VarList(vlistID2);
    varList_set_memtype(varList2, memtype);

    if (!extrapolate) numMiss_FL.resize(heightLevels.size());
    if (!extrapolate) numMiss_HL.resize(heightLevels.size());

    if (-1 != heightID_FL) vertIndex_FL.resize(gridsize * heightLevels.size());
    if (-1 != heightID_HL) vertIndex_HL.resize(gridsize * heightLevels.size());

    processVars = std::vector<bool>(numVars);
    interpVars = std::vector<bool>(numVars);
    varnumMissVals = Varray2D<size_t>(numVars);
    vardata1 = Field3DVector(numVars);
    vardata2 = Field3DVector(numVars);

    auto maxLevels = std::max(std::max(numFullLevels, numHalfLevels), (int) heightLevels.size());

    for (int varID = 0; varID < numVars; ++varID)
      {
        const auto &var = varList1.vars[varID];
        auto isHeightAxis = is_height_axis(var.zaxisID);

        if (gridInqType(var.gridID) == GRID_SPECTRAL) cdo_abort("Spectral data unsupported!");

        vardata1[varID].init(var);

        interpVars[varID] = (var.zaxisID == zaxisID_FL || var.zaxisID == zaxisID_HL
                             || (isHeightAxis && (var.nlevels == numHalfLevels || var.nlevels == numFullLevels)));

        if (interpVars[varID])
          {
            varnumMissVals[varID].resize(maxLevels, 0);
            vardata2[varID].init(varList2.vars[varID]);
          }
        else
          {
            if (isHeightAxis && var.nlevels > 1)
              {
                if (-1 == heightID_FL && -1 != heightID_HL && var.nlevels == (numHalfLevels - 1))
                  cdo_abort("%s not found (needed for %s)!", stdnameHeight_FL, var.name);
                else if (-1 != heightID_FL && -1 == heightID_HL && var.nlevels == (numFullLevels + 1))
                  cdo_abort("%s not found (needed for %s)!", stdnameHeight_HL, var.name);
                else
                  cdo_warning("Parameter %d has wrong number of levels, skipped! (name=%s nlevel=%d)", varID + 1, var.name,
                              var.nlevels);
              }
            varnumMissVals[varID].resize(var.nlevels);
          }
      }

    for (int varID = 0; varID < numVars; ++varID)
      {
        const auto &var = varList1.vars[varID];
        if (interpVars[varID] && var.isConstant) vlistDefVarTimetype(vlistID2, varID, TIME_VARYING);
      }

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

  void
  run() override
  {
    Field heightBottom;

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

        for (int varID = 0; varID < numVars; ++varID)
          {
            processVars[varID] = false;
            const auto &var = varList1.vars[varID];
            for (int levelID = 0; levelID < var.nlevels; ++levelID) varnumMissVals[varID][levelID] = 0;
          }

        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_read_field(streamID1, vardata1[varID], levelID, &varnumMissVals[varID][levelID]);
            processVars[varID] = true;
          }

        for (int varID = 0; varID < numVars; ++varID)
          if (interpVars[varID]) processVars[varID] = true;

        auto lreverse = true;
        if (-1 != heightID_FL && (tsID == 0 || !varList1.vars[heightID_FL].isConstant))
          {
            gen_vert_index(vertIndex_FL, heightLevels, vardata1[heightID_FL], gridsize, lreverse);
            if (!extrapolate)
              {
                heightBottom.init(varList1.vars[heightID_FL]);
                field_copy(vardata1[heightID_FL], numFullLevels - 1, heightBottom);
                gen_vert_index_mv(vertIndex_FL, heightLevels, gridsize, heightBottom, numMiss_FL, lreverse);
              }
          }

        if (-1 != heightID_HL && (tsID == 0 || !varList1.vars[heightID_HL].isConstant))
          {
            gen_vert_index(vertIndex_HL, heightLevels, vardata1[heightID_HL], gridsize, lreverse);
            if (!extrapolate)
              {
                heightBottom.init(varList1.vars[heightID_HL]);
                field_copy(vardata1[heightID_HL], numHalfLevels - 1, heightBottom);
                gen_vert_index_mv(vertIndex_HL, heightLevels, gridsize, heightBottom, numMiss_HL, lreverse);
              }
          }

        for (int varID = 0; varID < numVars; ++varID)
          {
            if (processVars[varID])
              {
                const auto &var = varList1.vars[varID];
                if (tsID > 0 && !interpVars[varID] && var.isConstant) continue;

                if (interpVars[varID])
                  {
                    if (var.nlevels != numFullLevels && var.nlevels != numHalfLevels)
                      cdo_abort("Number of generalized height level differ from full/half level (param=%s)!", var.name);

                    for (int levelID = 0; levelID < var.nlevels; ++levelID)
                      {
                        if (varnumMissVals[varID][levelID]) cdo_abort("Missing values unsupported for this operator!");
                      }

                    const auto &height3D = (var.nlevels == numFullLevels) ? vardata1[heightID_FL] : vardata1[heightID_HL];
                    const auto &vertIndex3D = (var.nlevels == numFullLevels) ? vertIndex_FL : vertIndex_HL;
                    vertical_interp_X(height3D, vardata1[varID], vardata2[varID], vertIndex3D, heightLevels, gridsize);

                    if (!extrapolate)
                      {
                        auto numMiss = (var.nlevels == numFullLevels) ? numMiss_FL : numMiss_HL;
                        varray_copy(heightLevels.size(), numMiss, varnumMissVals[varID]);
                      }
                  }

                for (int levelID = 0; levelID < varList2.vars[varID].nlevels; ++levelID)
                  {
                    cdo_def_field(streamID2, varID, levelID);
                    cdo_write_field(streamID2, interpVars[varID] ? vardata2[varID] : vardata1[varID], levelID,
                                     varnumMissVals[varID][levelID]);
                  }
              }
          }

        tsID++;
      }
  }

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