/*
  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    ap2pl           Model air pressure level to pressure level interpolation
*/

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

static void
check_range_ps(int stepNum, const Field &psProg)
{
  auto mm = field_min_max(psProg);
  if (mm.min < MIN_PS || mm.max > MAX_PS)
    cdo_warning("Surface pressure out of range (min=%g max=%g) [timestep:%d]!", mm.min, mm.max, stepNum);
}

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;
      if (stdname == "height") isHeight = true;
    }
  return isHeight;
}

template <typename T>
static void
calc_half_press(size_t gridsize, size_t numFullLevels, const Varray<T> &fullPress, size_t numHalfLevels, Varray<T> &halfPress)
{
  for (size_t i = 0; i < gridsize; ++i) halfPress[i] = 0;
#ifdef _OPENMP
#pragma omp parallel for default(shared) schedule(static)
#endif
  for (size_t k = 1; k < numFullLevels; ++k)
    {
      auto fullPress_km1 = &fullPress[(k - 1) * gridsize];
      auto fullPress_k = &fullPress[k * gridsize];
      auto halfPress_k = &halfPress[k * gridsize];
      for (size_t i = 0; i < gridsize; ++i) halfPress_k[i] = 0.5 * (fullPress_km1[i] + fullPress_k[i]);
    }
  for (size_t i = 0; i < gridsize; ++i)
    halfPress[(numHalfLevels - 1) * gridsize + i] = fullPress[(numFullLevels - 1) * gridsize + i];
}

static void
calc_half_press(const Field3D &fullPress, Field3D &halfPress)
{
  if (fullPress.memType == MemType::Float)
    calc_half_press(fullPress.gridsize, fullPress.nlevels, fullPress.vec_f, halfPress.nlevels, halfPress.vec_f);
  else
    calc_half_press(fullPress.gridsize, fullPress.nlevels, fullPress.vec_d, halfPress.nlevels, halfPress.vec_d);
}

class ModuleVertintap
{
  enum
  {
    func_pl,
    func_hl
  };
  int apressID_FL = -1, apressID_HL = -1, dpressID = -1;
  int psID = -1;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID1;
  int taxisID2;

  size_t gridsize;

  int numPL;
  int nvars;
  int zaxisID_FL;
  int zaxisID_HL;
  int numFullLevels;
  int numHalfLevels;

  bool extrapolate;

  VarList varList1;
  VarList varList2;

  Varray<double> levels;

  std::vector<bool> processVars, interpVars;
  Varray2D<size_t> varnmiss;
  Field3DVector vardata1, vardata2;

  Varray<size_t> numMiss_FL, numMiss_HL;
  std::vector<int> vertIndex_FL, vertIndex_HL;
  Field psProg;
  Field3D fullPress, halfPress;

  CdoVar var3Dfull, var3Dhalf;

public:
  void
  init(void *process)
  {
    cdo_initialize(process);

    // clang-format off
                     cdo_operator_add("ap2pl",     func_pl, 0, "pressure levels in pascal");
    auto AP2PLX    = cdo_operator_add("ap2plx",    func_pl, 0, "pressure levels in pascal");
                     cdo_operator_add("ap2hl",     func_hl, 0, "height levels in meter");
    auto AP2HLX    = cdo_operator_add("ap2hlx",    func_hl, 0, "height levels in meter");
    // clang-format on

    auto operatorID = cdo_operator_id();
    auto useHeightLevel = (cdo_operator_f1(operatorID) == func_hl);

    extrapolate = (operatorID == AP2PLX || operatorID == AP2HLX);
    if (extrapolate == false) extrapolate = getenv_extrapolate();

    operator_input_arg(cdo_operator_enter(operatorID));

    if (cdo_operator_argc() == 1 && cdo_operator_argv(0) == "default")
      {
        if (useHeightLevel)
          levels = { 10, 50, 100, 500, 1000, 5000, 10000, 15000, 20000, 25000, 30000 };
        else
          levels = {
            100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000
          };
      }
    else
      {
        levels = cdo_argv_to_flt(cdo_get_oper_argv());
      }

    numPL = levels.size();

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

    auto zaxistype = useHeightLevel ? ZAXIS_HEIGHT : ZAXIS_PRESSURE;
    auto zaxisIDp = zaxisCreate(zaxistype, numPL);
    zaxisDefLevels(zaxisIDp, levels.data());

    varList_init(varList1, vlistID1);
    varListSetUniqueMemtype(varList1);
    auto memtype = varList1[0].memType;

    nvars = vlistNvars(vlistID1);
    for (int varID = 0; varID < nvars; ++varID)
      {
        auto stdname = string_to_lower(cdo::inq_key_string(vlistID1, varID, CDI_KEY_STDNAME));

        // clang-format off
        if      (stdname == var_stdname(surface_air_pressure))              psID = varID;
        else if (stdname == var_stdname(pressure_thickness))                dpressID = varID;
        else if (stdname == var_stdname(air_pressure) && apressID_FL == -1) apressID_FL = varID;
        else if (stdname == var_stdname(air_pressure) && apressID_HL == -1) apressID_HL = varID;
        // clang-format on
      }

    if (-1 != apressID_FL && -1 != apressID_HL)
      {
        if (varList1[apressID_FL].nlevels == (varList1[apressID_HL].nlevels + 1)) std::swap(apressID_FL, apressID_HL);
      }

    if (Options::cdoVerbose)
      {
        cdo_print("Found:");
        // clang-format off
        if (-1 != psID)        cdo_print("  %s -> %s", var_stdname(surface_air_pressure), varList1[psID].name);
        if (-1 != dpressID)    cdo_print("  %s -> %s", var_stdname(pressure_thickness), varList1[dpressID].name);
        if (-1 != apressID_FL) cdo_print("  %s (full) -> %s", var_stdname(air_pressure), varList1[apressID_FL].name);
        if (-1 != apressID_HL) cdo_print("  %s (half) -> %s", var_stdname(air_pressure), varList1[apressID_HL].name);
        // clang-format on
      }

    if (-1 == apressID_FL) cdo_abort("%s not found!", var_stdname(air_pressure));

    zaxisID_FL = (-1 == apressID_FL) ? -1 : varList1[apressID_FL].zaxisID;
    zaxisID_HL = (-1 == apressID_HL) ? -1 : varList1[apressID_HL].zaxisID;
    numFullLevels = (-1 == zaxisID_FL) ? 0 : varList1[apressID_FL].nlevels;
    numHalfLevels = (-1 == zaxisID_HL) ? numFullLevels + 1 : varList1[apressID_HL].nlevels;

    auto nzaxis = vlistNzaxis(vlistID1);
    for (int index = 0; index < nzaxis; ++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, zaxisIDp);
      }

    varList_init(varList2, vlistID2);
    varListSetMemtype(varList2, memtype);

    processVars = std::vector<bool>(nvars);
    interpVars = std::vector<bool>(nvars);
    varnmiss = Varray2D<size_t>(nvars);
    vardata1 = Field3DVector(nvars), vardata2 = Field3DVector(nvars);

    auto maxLevels = std::max(std::max(numFullLevels, numHalfLevels), numPL);

    if (!extrapolate) numMiss_FL.resize(numPL);
    if (!extrapolate) numMiss_HL.resize(numPL);

    vertIndex_FL.resize(gridsize * numPL);
    vertIndex_HL.resize(gridsize * numPL);

    var3Dfull.gridsize = gridsize;
    var3Dfull.nlevels = numFullLevels;
    var3Dfull.memType = memtype;
    fullPress.init(var3Dfull);

    var3Dhalf.gridsize = gridsize;
    var3Dhalf.nlevels = numHalfLevels;
    var3Dhalf.memType = memtype;
    halfPress.init(var3Dhalf);

    if (useHeightLevel)
      {
        Varray<double> phlev(numPL);
        height_to_pressure(levels.data(), phlev.data(), numPL);

        if (Options::cdoVerbose)
          for (int i = 0; i < numPL; ++i) cdo_print("level=%d   height=%g   pressure=%g", i + 1, levels[i], phlev[i]);

        levels = phlev;
      }

    for (int varID = 0; varID < nvars; ++varID)
      {
        const auto &var = varList1[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
               || (isHeightAxis && zaxisID_FL != -1 && (var.nlevels == numHalfLevels || var.nlevels == numFullLevels)));

        if (interpVars[varID])
          {
            varnmiss[varID].resize(maxLevels, 0);
            vardata2[varID].init(varList2[varID]);
          }
        else
          {
            if (isHeightAxis && zaxisID_FL != -1 && var.nlevels > 1)
              cdo_warning("Parameter %d has wrong number of levels, skipped! (name=%s nlevel=%d)", varID + 1, var.name,
                          var.nlevels);

            varnmiss[varID].resize(var.nlevels);
          }
      }

    if (zaxisID_FL != -1 && psID == -1)
      {
        if (dpressID != -1)
          cdo_warning("Surface pressure not found - set to vertical sum of %s!", var_stdname(pressure_thickness));
        else
          cdo_warning("Surface pressure not found - set to lower bound of %s!", var_stdname(air_pressure));
      }

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

    streamID2 = cdo_open_write(1);

    cdo_def_vlist(streamID2, vlistID2);
  }

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

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

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

        for (int recID = 0; recID < nrecs; ++recID)
          {
            int varID, levelID;
            cdo_inq_record(streamID1, &varID, &levelID);
            cdo_read_record(streamID1, vardata1[varID], levelID, &varnmiss[varID][levelID]);
            processVars[varID] = true;
          }

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

        if (zaxisID_FL != -1)
          {
            if (tsID == 1 && varList1[apressID_FL].timetype == TIME_CONSTANT)
              cdo_warning("%s does not vary in time!", var_stdname(air_pressure));

            if (psID != -1)
              {
                psProg.init(varList1[psID]);
                field_copy(vardata1[psID], psProg);
              }
            else if (dpressID != -1)
              {
                psProg.init(varList1[dpressID]);
                field_fill(psProg, 0);
                for (int k = 0; k < numFullLevels; ++k) field_add(psProg, vardata1[dpressID], k);
              }
            else
              {
                psProg.init(varList1[apressID_FL]);
                field_copy(vardata1[apressID_FL], numFullLevels - 1, psProg);
              }

            // check range of psProg
            check_range_ps(tsID + 1, psProg);

            field_copy(vardata1[apressID_FL], fullPress);

            if (-1 != zaxisID_HL)
              field_copy(vardata1[apressID_HL], halfPress);
            else
              calc_half_press(fullPress, halfPress);

            gen_vert_index(vertIndex_FL, levels, fullPress, gridsize);
            if (!extrapolate) gen_vert_index_mv(vertIndex_FL, levels, gridsize, psProg, numMiss_FL);

            gen_vert_index(vertIndex_HL, levels, halfPress, gridsize);
            if (!extrapolate) gen_vert_index_mv(vertIndex_HL, levels, gridsize, psProg, numMiss_HL);
          }

        for (int varID = 0; varID < nvars; ++varID)
          {
            if (processVars[varID])
              {
                const auto &var = varList1[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 (varnmiss[varID][levelID]) cdo_abort("Missing values unsupported for this operator!");
                      }

                    const auto &levels3D = (var.nlevels == numFullLevels) ? fullPress : halfPress;
                    const auto &vertIndex3D = (var.nlevels == numFullLevels) ? vertIndex_FL : vertIndex_HL;
                    vertical_interp_X(levels3D, vardata1[varID], vardata2[varID], vertIndex3D, levels, gridsize);

                    if (!extrapolate)
                      {
                        const auto &numMiss = (var.nlevels == numFullLevels) ? numMiss_FL : numMiss_HL;
                        varray_copy(numPL, numMiss, varnmiss[varID]);
                      }
                  }

                for (int levelID = 0; levelID < varList2[varID].nlevels; ++levelID)
                  {
                    cdo_def_record(streamID2, varID, levelID);
                    cdo_write_record(streamID2, interpVars[varID] ? vardata2[varID] : vardata1[varID], levelID,
                                     varnmiss[varID][levelID]);
                  }
              }
          }

        tsID++;
      }
  }

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

    cdo_finish();
  }
};

void *
Vertintap(void *process)
{
  ModuleVertintap verintap;
  verintap.init(process);
  verintap.run();
  verintap.close();
  return nullptr;
}
