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

      Math       abs             Absolute value
      Math       sqr             Square
      Math       sqrt            Square root
      Math       exp             Exponential
      Math       ln              Natural logarithm
      Math       log10           Base 10 logarithm
      Math       sin             Sine
      Math       cos             Cosine
      Math       tan             Tangent
      Math       asin            Arc sine
      Math       acos            Arc cosine
      Math       atan            Arc tangent
      Math       pow             Power
      Math       reci            Reciprocal
*/

#include <cstdlib>
#include <cdi.h>

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

static void
check_out_of_range(size_t &numMissVals, const size_t len, double missval, Varray<double> &v, double rmin, double rmax)
{
  if (numMissVals)
    {
      for (size_t i = 0; i < len; ++i)
        if (fp_is_not_equal(v[i], missval) && (v[i] < rmin || v[i] > rmax))
          {
            v[i] = missval;
            numMissVals++;
          }
    }
  else
    {
      for (size_t i = 0; i < len; ++i)
        if (v[i] < rmin || v[i] > rmax)
          {
            v[i] = missval;
            numMissVals++;
          }
    }
}

static void
check_lower_range(size_t &numMissVals, const size_t len, double missval, Varray<double> &v, double rmin)
{
  if (numMissVals)
    {
      for (size_t i = 0; i < len; ++i)
        if (fp_is_not_equal(v[i], missval) && v[i] < rmin)
          {
            v[i] = missval;
            numMissVals++;
          }
    }
  else
    {
      for (size_t i = 0; i < len; ++i)
        if (v[i] < rmin)
          {
            v[i] = missval;
            numMissVals++;
          }
    }
}

template <typename T, class UnaryOperation>
void
math_varray_transform(const size_t numMissVals, const size_t len, double missval1, const Varray<T> &array1, Varray<T> &array2,
                      UnaryOperation unary_op)
{
  if (numMissVals)
    for (size_t i = 0; i < len; ++i) array2[i] = fp_is_equal(array1[i], missval1) ? missval1 : unary_op(array1[i]);
  else
    for (size_t i = 0; i < len; ++i) array2[i] = unary_op(array1[i]);
}

static void
math_varray_sqr_cplx(const size_t len, const Varray<double> &array1, Varray<double> &array2)
{
  for (size_t i = 0; i < len; ++i)
    {
      array2[i * 2] = array1[i * 2] * array1[i * 2] + array1[i * 2 + 1] * array1[i * 2 + 1];
      array2[i * 2 + 1] = 0.0;
    }
}

static void
math_varray_sqrt_cplx(const size_t len, double missval1, const Varray<double> &array1, Varray<double> &array2)
{
  auto is_EQ = fp_is_equal;
  auto missval2 = missval1;
  auto rsqrt2 = 1.0 / std::sqrt(2.0);
  for (size_t i = 0; i < len; ++i)
    {
      double abs = SQRTM(ADDM(MULM(array1[2 * i], array1[2 * i]), MULM(array1[2 * i + 1], array1[2 * i + 1])));
      array2[i * 2] = MULM(rsqrt2, SQRTM(ADDM(array1[i * 2], abs)));
      array2[i * 2 + 1] = MULM(rsqrt2, DIVM(array1[2 * i + 1], SQRTM(ADDM(array1[2 * i], abs))));
    }
}

static void
math_varray_conj_cplx(const size_t len, const Varray<double> &array1, Varray<double> &array2)
{
  for (size_t i = 0; i < len; ++i)
    {
      array2[i * 2] = array1[i * 2];
      array2[i * 2 + 1] = -array1[i * 2 + 1];
    }
}

static void
math_varray_abs_cplx(const size_t len, double missval1, const Varray<double> &array1, Varray<double> &array2)
{
  auto is_EQ = fp_is_equal;
  auto missval2 = missval1;
  for (size_t i = 0; i < len; ++i)
    {
      array2[i] = SQRTM(ADDM(MULM(array1[2 * i], array1[2 * i]), MULM(array1[2 * i + 1], array1[2 * i + 1])));
    }
}

static void
math_varray_arg_cplx(const size_t len, double missval1, const Varray<double> &array1, Varray<double> &array2)
{
  for (size_t i = 0; i < len; ++i)
    {
      array2[i] = (fp_is_equal(array1[2 * i], missval1) || fp_is_equal(array1[2 * i + 1], missval1))
                      ? missval1
                      : std::atan2(array1[2 * i + 1], array1[2 * i]);
    }
}

enum struct Oper
{
  Abs,
  Int,
  Nint,
  Sqr,
  Sqrt,
  Exp,
  Ln,
  Log10,
  Sin,
  Cos,
  Tan,
  Asin,
  Acos,
  Atan,
  Pow,
  Rand,
  Reci,
  Not,
  Conj,
  Re,
  Im,
  Arg
};

class Math : public Process
{
public:
  using Process::Process;
  inline static CdoModule module = {
    .name = "Math",
    .operators = { { "abs", (int) Oper::Abs, 0, MathHelp },   { "int", (int) Oper::Int, 0, MathHelp },
                   { "nint", (int) Oper::Nint, 0, MathHelp }, { "sqr", (int) Oper::Sqr, 0, MathHelp },
                   { "sqrt", (int) Oper::Sqrt, 0, MathHelp }, { "exp", (int) Oper::Exp, 0, MathHelp },
                   { "ln", (int) Oper::Ln, 0, MathHelp },     { "log10", (int) Oper::Log10, 0, MathHelp },
                   { "sin", (int) Oper::Sin, 0, MathHelp },   { "cos", (int) Oper::Cos, 0, MathHelp },
                   { "tan", (int) Oper::Tan, 0, MathHelp },   { "asin", (int) Oper::Asin, 0, MathHelp },
                   { "acos", (int) Oper::Acos, 0, MathHelp }, { "atan", (int) Oper::Atan, 0, MathHelp },
                   { "pow", (int) Oper::Pow, 0, MathHelp },   { "rand", (int) Oper::Rand, 0, MathHelp },
                   { "reci", (int) Oper::Reci, 0, MathHelp }, { "not", (int) Oper::Not, 0, MathHelp },
                   { "conj", (int) Oper::Conj, 0, MathHelp }, { "re", (int) Oper::Re, 0, MathHelp },
                   { "im", (int) Oper::Im, 0, MathHelp },     { "arg", (int) Oper::Arg, 0, MathHelp } },
    .aliases = { { "log", "ln" } },
    .mode = EXPOSED,     // Module mode: 0:intern 1:extern
    .number = CDI_BOTH,  // Allowed number type
    .constraints = { 1, 1, NoRestriction },
  };
  inline static RegisterEntry<Math> registration = RegisterEntry<Math>(module);

private:
  Oper operfunc;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID1;
  int taxisID2;
  int vlistID1;
  int vlistID2;

  VarList varList1;

  double rc = 0.0;

public:
  void
  init() override
  {
    auto operatorID = cdo_operator_id();
    operfunc = (Oper) cdo_operator_f1(operatorID);

    if (operfunc == Oper::Pow)
      {
        operator_input_arg("value");
        rc = parameter_to_double(cdo_operator_argv(0));
      }
    else { operator_check_argc(0); }

    if (operfunc == Oper::Rand) std::srand(Options::Random_Seed);

    streamID1 = cdo_open_read(0);

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

    varList1 = VarList(vlistID1);

    if (operfunc == Oper::Re || operfunc == Oper::Im || operfunc == Oper::Abs || operfunc == Oper::Arg)
      {
        auto numVars = varList1.numVars();
        for (int varID = 0; varID < numVars; ++varID)
          {
            const auto &var1 = varList1.vars[varID];
            if (var1.dataType == CDI_DATATYPE_CPX32) vlistDefVarDatatype(vlistID2, varID, CDI_DATATYPE_FLT32);
            if (var1.dataType == CDI_DATATYPE_CPX64) vlistDefVarDatatype(vlistID2, varID, CDI_DATATYPE_FLT64);
          }
      }

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

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

  void
  run() override
  {
    auto gridsizemax = vlistGridsizeMax(vlistID1);
    if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;

    auto array1 = Varray<double>(gridsizemax);
    auto array2 = Varray<double>(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);

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

            const auto &var = varList1.vars[varID];
            auto is_EQ = fp_is_equal;
            auto missval1 = var.missval;
            auto n = var.gridsize;
            auto number = var.nwpv;

            if (number == CDI_REAL)
              {
                // clang-format off
              switch (operfunc)
                {
                case Oper::Abs:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_abs); break;
                case Oper::Int:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_int); break;
                case Oper::Nint:  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_nint); break;
                case Oper::Sqr:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_sqr); break;
                case Oper::Sqrt:
                  for (size_t i = 0; i < n; ++i) array2[i] = SQRTM(array1[i]);
                  break;
                case Oper::Exp:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_exp); break;
                case Oper::Ln:    check_lower_range(numMissVals, n, missval1, array1, -1);
                                  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_log); break;
                case Oper::Log10: check_lower_range(numMissVals, n, missval1, array1, -1);
                                  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_log10); break;
                case Oper::Sin:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_sin); break;
                case Oper::Cos:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_cos); break;
                case Oper::Tan:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_tan); break;
                case Oper::Asin:  check_out_of_range(numMissVals, n, missval1, array1, -1, 1);
                                  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_asin); break;
                case Oper::Acos:  check_out_of_range(numMissVals, n, missval1, array1, -1, 1);
                                  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_acos); break;
                case Oper::Atan:  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_atan); break;
                case Oper::Pow:
                  for (size_t i = 0; i < n; ++i) array2[i] = fp_is_equal(array1[i], missval1) ? missval1 : std::pow(array1[i], rc);
                  break;
                case Oper::Rand:
                  for (size_t i = 0; i < n; ++i) array2[i] = fp_is_equal(array1[i], missval1) ? missval1 : ((double) std::rand()) / ((double) RAND_MAX);
                  break;
                case Oper::Reci:  math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_reci); break;
                case Oper::Not:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_not); break;
                case Oper::Re:
                case Oper::Arg:   math_varray_transform(numMissVals, n, missval1, array1, array2, unary_op_nop); break;
                default: cdo_abort("Operator not implemented for real data!"); break;
                }
                // clang-format on

                numMissVals = varray_num_mv(n, array2, missval1);
              }
            else
              {
                // clang-format off
              switch (operfunc)
                {
                case Oper::Sqr:   math_varray_sqr_cplx(n, array1, array2); break;
                case Oper::Sqrt:  math_varray_sqrt_cplx(n, missval1, array1, array2); break;
                case Oper::Conj:  math_varray_conj_cplx(n, array1, array2); break;
                case Oper::Re:    for (size_t i = 0; i < n; ++i) array2[i] = array1[i * 2]; break;
                case Oper::Im:    for (size_t i = 0; i < n; ++i) array2[i] = array1[i * 2 + 1]; break;
                case Oper::Abs:   math_varray_abs_cplx(n, missval1, array1, array2); break;
                case Oper::Arg:   math_varray_arg_cplx(n, missval1, array1, array2); break;
                default: cdo_abort("Fields with complex numbers are not supported by this operator!"); break;
                }
                // clang-format on

                numMissVals = 0;
              }

            cdo_def_field(streamID2, varID, levelID);
            cdo_write_field(streamID2, array2.data(), numMissVals);
          }

        tsID++;
      }
  }

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

    vlistDestroy(vlistID2);
  }
};
