/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Uwe Schulzweida

*/

#include <cdi.h>

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

template <typename T>
static void
isosurface_kernel(double isoval, size_t nmiss, const Varray<double> &levels, int nlevels, size_t gridsize, T missval,
                  const Varray<const T *> &data3D, Varray<T> &data2D)
{
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t i = 0; i < gridsize; ++i)
    {
      data2D[i] = missval;

      for (int k = 0; k < (nlevels - 1); ++k)
        {
          const double val1 = data3D[k][i];
          const double val2 = data3D[k + 1][i];

          if (nmiss)
            {
              const auto hasMissvals1 = DBL_IS_EQUAL(val1, missval);
              const auto hasMissvals2 = DBL_IS_EQUAL(val2, missval);
              if (hasMissvals1 && hasMissvals2) continue;
              if (hasMissvals1 && IS_EQUAL(isoval, val2)) data2D[i] = levels[k + 1];
              if (hasMissvals2 && IS_EQUAL(isoval, val1)) data2D[i] = levels[k];
              if (hasMissvals1 || hasMissvals2) continue;
            }

          if ((isoval >= val1 && isoval <= val2) || (isoval >= val2 && isoval <= val1))
            {
              data2D[i] = IS_EQUAL(val1, val2) ? levels[k] : intlin(isoval, levels[k], val1, levels[k + 1], val2);
              break;
            }
        }
    }
}

static void
isosurface(double isoval, int nlevels, const Varray<double> &levels, const FieldVector &field3D, Field &field2D)
{
  const auto gridsize = gridInqSize(field3D[0].grid);
  const auto missval = field3D[0].missval;

  auto nmiss = field3D[0].nmiss;
  for (int k = 1; k < nlevels; ++k) nmiss += field3D[k].nmiss;

  if (field3D[0].memType == MemType::Float)
    {
      Varray<const float *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_f.data();
      isosurface_kernel(isoval, nmiss, levels, nlevels, gridsize, (float) missval, data3D, field2D.vec_f);
    }
  else
    {
      Varray<const double *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_d.data();
      isosurface_kernel(isoval, nmiss, levels, nlevels, gridsize, missval, data3D, field2D.vec_d);
    }

  field_num_mv(field2D);
}

template <typename T>
static void
layer_value_min_kernel(int nlevels, size_t gridsize, T missval, const Varray<const T *> &data3D, Varray<T> &data2D)
{
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t i = 0; i < gridsize; ++i)
    {
      data2D[i] = missval;

      for (int k = 0; k < nlevels; ++k)
        {
          const auto val = data3D[k][i];
          if (!DBL_IS_EQUAL(val, missval))
            {
              data2D[i] = val;
              break;
            }
        }
    }
}

static void
layer_value_min(int nlevels, const FieldVector &field3D, Field &field2D)
{
  const auto gridsize = gridInqSize(field3D[0].grid);
  const auto missval = field3D[0].missval;

  if (field3D[0].memType == MemType::Float)
    {
      Varray<const float *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_f.data();
      layer_value_min_kernel(nlevels, gridsize, (float) missval, data3D, field2D.vec_f);
    }
  else
    {
      Varray<const double *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_d.data();
      layer_value_min_kernel(nlevels, gridsize, missval, data3D, field2D.vec_d);
    }

  field_num_mv(field2D);
}

template <typename T>
static void
layer_value_max_kernel(int nlevels, size_t gridsize, T missval, const Varray<const T *> &data3D, Varray<T> &data2D)
{
#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t i = 0; i < gridsize; ++i)
    {
      data2D[i] = missval;

      for (int k = nlevels - 1; k >= 0; --k)
        {
          const auto val = data3D[k][i];
          if (!DBL_IS_EQUAL(val, missval))
            {
              data2D[i] = val;
              break;
            }
        }
    }
}

static void
layer_value_max(int nlevels, const FieldVector &field3D, Field &field2D)
{
  const auto gridsize = gridInqSize(field3D[0].grid);
  const auto missval = field3D[0].missval;

  if (field3D[0].memType == MemType::Float)
    {
      Varray<const float *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_f.data();
      layer_value_max_kernel(nlevels, gridsize, (float) missval, data3D, field2D.vec_f);
    }
  else
    {
      Varray<const double *> data3D(nlevels);
      for (int k = 0; k < nlevels; ++k) data3D[k] = field3D[k].vec_d.data();
      layer_value_max_kernel(nlevels, gridsize, missval, data3D, field2D.vec_d);
    }

  field_num_mv(field2D);
}

static bool
positiveIsDown(int zaxisID)
{
  return zaxisInqPositive(zaxisID) == 2;
}

void *
Selsurface(void *process)
{
  cdo_initialize(process);

  // clang-format off
  const auto ISOSURFACE  = cdo_operator_add("isosurface",  0,  0, nullptr);
  const auto BOTTOMVALUE = cdo_operator_add("bottomvalue", 0,  0, nullptr);
  const auto TOPVALUE    = cdo_operator_add("topvalue",    0,  0, nullptr);
  // clang-format on

  const auto operatorID = cdo_operator_id();

  double isoval = 0.0;
  if (operatorID == ISOSURFACE)
    {
      operator_input_arg("isoval");
      operator_check_argc(1);
      isoval = parameter_to_double(cdo_operator_argv(0));
    }

  if (Options::cdoVerbose) cdo_print("Isoval: %g", isoval);

  const auto streamID1 = cdo_open_read(0);

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

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

  int zaxisID1 = -1;
  const auto nzaxis = vlistNzaxis(vlistID1);
  for (int i = 0; i < nzaxis; i++)
    {
      const auto zaxisID = vlistZaxis(vlistID1, i);
      const auto nlevels = zaxisInqSize(zaxisID);
      if (zaxisInqType(zaxisID) != ZAXIS_HYBRID && zaxisInqType(zaxisID) != ZAXIS_HYBRID_HALF)
        if (nlevels > 1)
          {
            zaxisID1 = zaxisID;
            break;
          }
    }
  if (zaxisID1 == -1) cdo_abort("No processable variable found!");

  const auto nlevels = zaxisInqSize(zaxisID1);
  Varray<double> levels(nlevels);
  cdo_zaxis_inq_levels(zaxisID1, levels.data());

  const auto lpositive = !(levels[0] < 0.0 && levels[nlevels - 1] < 0.0);
  const auto lreverse = (levels[0] > levels[nlevels - 1]);
  auto bottom_value_func = lreverse ? layer_value_max : layer_value_min;
  auto top_value_func = lreverse ? layer_value_min : layer_value_max;
  if (lpositive && positiveIsDown(zaxisID1)) std::swap(bottom_value_func, top_value_func);

  const auto zaxisIDsfc = zaxis_from_name("surface");
  for (int i = 0; i < nzaxis; i++)
    if (zaxisID1 == vlistZaxis(vlistID1, i)) vlistChangeZaxisIndex(vlistID2, i, zaxisIDsfc);

  const auto streamID2 = cdo_open_write(1);
  cdo_def_vlist(streamID2, vlistID2);

  VarList varList1;
  varListInit(varList1, vlistID1);

  const auto nvars = vlistNvars(vlistID1);
  std::vector<bool> isVar3D(nvars), foundVar(nvars);

  Field field2;
  FieldVector2D vars1;
  fields_from_vlist(vlistID1, vars1);

  for (int varID = 0; varID < nvars; varID++) isVar3D[varID] = (varList1[varID].zaxisID == zaxisID1);

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

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

      for (int varID = 0; varID < nvars; varID++) foundVar[varID] = false;

      for (int recID = 0; recID < nrecs; recID++)
        {
          int varID, levelID;
          cdo_inq_record(streamID1, &varID, &levelID);
          auto &field1 = vars1[varID][levelID];
          field1.init(varList1[varID]);
          cdo_read_record(streamID1, field1);
          foundVar[varID] = true;
        }

      for (int varID = 0; varID < nvars; varID++)
        {
          if (foundVar[varID])
            {
              if (isVar3D[varID])
                {
                  field2.init(varList1[varID]);
                  // clang-format off
                  if      (operatorID == ISOSURFACE)  isosurface(isoval, nlevels, levels, vars1[varID], field2);
                  else if (operatorID == BOTTOMVALUE) bottom_value_func(nlevels, vars1[varID], field2);
                  else if (operatorID == TOPVALUE)    top_value_func(nlevels, vars1[varID], field2);
                  // clang-format on

                  cdo_def_record(streamID2, varID, 0);
                  cdo_write_record(streamID2, field2);
                }
              else
                {
                  for (int levelID = 0; levelID < varList1[varID].nlevels; levelID++)
                    {
                      cdo_def_record(streamID2, varID, levelID);
                      cdo_write_record(streamID2, vars1[varID][levelID]);
                    }
                }
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  vlistDestroy(vlistID2);

  cdo_finish();

  return nullptr;
}
