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

  Copyright (C) 2003-2020 Uwe Schulzweida, <uwe.schulzweida AT mpimet.mpg.de>
  See COPYING file for copying and redistribution conditions.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; version 2 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
*/

/*
   This module contains the following operators:

      Vargen     const           Create a constant field
      Vargen     random          Field with random values
      Vargen     stdatm          Field values for pressure and temperature for the standard atmosphere
*/

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

#include "cdo_options.h"
#include "process_int.h"
#include "cdo_cdi_wrapper.h"
#include "cdo_zaxis.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "griddes.h"
#include "constants.h"
#include "stdnametable.h"
#include "timebase.h"
#include "param_conversion.h"

static constexpr double etopo_scale = 3.0;
static constexpr double etopo_offset = 11000.0;
static constexpr unsigned short etopo[] = {
#include "etopo.dat"
};

static constexpr double temp_scale = 500.0;
static constexpr double temp_offset = -220.0;
static constexpr unsigned short temp[] = {
#include "temp.dat"
};

static constexpr double mask_scale = 1.0;
static constexpr double mask_offset = 0.0;
static constexpr unsigned short mask[] = {
#include "mask.dat"
};

// Some Constants for creating temperatur and pressure for the standard atmosphere
constexpr double T_ZERO = 213.0;
constexpr double T_DELTA = 75.0;
constexpr double SCALEHEIGHT = 10000.0;  // [m]
constexpr double P_ZERO = 1013.25;       // surface pressure [hPa]
constexpr double CC_R = 287.05;          // specific gas constant for air
constexpr double TMP4PRESSURE = (C_EARTH_GRAV * SCALEHEIGHT) / (CC_R * T_ZERO);

static double
std_atm_temperatur(const double height)
{
  // Compute the temperatur for the given height (in meters) according to the solution of the hydrostatic atmosphere
  return (T_ZERO + T_DELTA * std::exp((-1) * (height / SCALEHEIGHT)));
}

static double
std_atm_pressure(const double height)
{
  // Compute the pressure for the given height (in meters) according to the solution of the hydrostatic atmosphere
  return (P_ZERO
          * std::exp((-1) * TMP4PRESSURE * std::log((std::exp(height / SCALEHEIGHT) * T_ZERO + T_DELTA) / (T_ZERO + T_DELTA))));
}

static void
conv_generic_grid(const int gridID, const size_t gridsize, Varray<double> &xvals2D, Varray<double> &yvals2D)
{
  const auto xsize = gridInqXsize(gridID);
  const auto ysize = gridInqYsize(gridID);

  assert(gridsize == xsize * ysize);

  Varray<double> xcoord(xsize), ycoord(ysize);
  gridInqXvals(gridID, &xcoord[0]);
  gridInqYvals(gridID, &ycoord[0]);

  const auto xrange = varrayRange(xsize, xcoord);
  const auto yrange = varrayRange(ysize, ycoord);

  for (size_t j = 0; j < ysize; ++j)
    for (size_t i = 0; i < xsize; ++i)
      {
        xvals2D[j * xsize + i] = xcoord[i] * M_PI / xrange;
        yvals2D[j * xsize + i] = ycoord[j] * M_PI / yrange;
      }
}

static void
remap_nn_reg2d_reg2d(size_t nx, size_t ny, const double *restrict data, const int gridID, double *restrict array)
{
  if (gridInqType(gridID) != GRID_LONLAT) cdoAbort("Internal error, wrong grid type!");

  auto nxvals = gridInqXsize(gridID);
  auto nyvals = gridInqYsize(gridID);
  Varray<double> xvals(nxvals), yvals(nyvals);

  gridInqXvals(gridID, &xvals[0]);
  gridInqYvals(gridID, &yvals[0]);

  // Convert lat/lon units if required
  cdo_grid_to_degree(gridID, CDI_XAXIS, nxvals, &xvals[0], "grid center lon");
  cdo_grid_to_degree(gridID, CDI_YAXIS, nyvals, &yvals[0], "grid center lat");

#ifdef _OPENMP
#pragma omp parallel for default(none) shared(nx, ny, nxvals, nyvals, xvals, yvals, data, array)
#endif
  for (size_t j = 0; j < nyvals; j++)
    {
      const auto yval = yvals[j];
      for (size_t i = 0; i < nxvals; i++)
        {
          auto xval = xvals[i];
          if (xval >= 180) xval -= 360;
          if (xval < -180) xval += 360;
          size_t ii = (xval + 180) * 2;
          size_t jj = (yval + 90) * 2;
          if (ii >= nx) ii = nx - 1;
          if (jj >= ny) jj = ny - 1;
          array[j * nxvals + i] = data[jj * nx + ii];
        }
    }
}

static void
remap_nn_reg2d_nonreg2d(const size_t nx, const size_t ny, const double *restrict data, const int gridID, double *restrict array)
{
  auto gridID2 = gridID;
  const auto gridsize = gridInqSize(gridID2);
  Varray<double> xvals(gridsize), yvals(gridsize);

  if (gridInqType(gridID2) == GRID_GME || gridInqType(gridID2) == GRID_GAUSSIAN_REDUCED) gridID2 = gridToUnstructured(gridID2, 0);

  if (gridInqType(gridID2) != GRID_UNSTRUCTURED && gridInqType(gridID2) != GRID_CURVILINEAR)
    gridID2 = gridToCurvilinear(gridID2, 0);

  gridInqXvals(gridID2, &xvals[0]);
  gridInqYvals(gridID2, &yvals[0]);

  // Convert lat/lon units if required
  cdo_grid_to_degree(gridID2, CDI_XAXIS, gridsize, &xvals[0], "grid center lon");
  cdo_grid_to_degree(gridID2, CDI_YAXIS, gridsize, &yvals[0], "grid center lat");

  for (size_t i = 0; i < gridsize; i++)
    {
      auto xval = xvals[i];
      auto yval = yvals[i];
      if (xval >= 180) xval -= 360;
      if (xval < -180) xval += 360;
      size_t ii = (xval + 180) * 2;
      size_t jj = (yval + 90) * 2;
      if (ii >= nx) ii = nx - 1;
      if (jj >= ny) jj = ny - 1;
      array[i] = data[jj * nx + ii];
    }

  if (gridID != gridID2) gridDestroy(gridID2);
}

static void
remap_nn_reg2d(const size_t nx, const size_t ny, const double *restrict data, const int gridID, double *restrict array)
{
  if (gridInqType(gridID) == GRID_LONLAT)
    remap_nn_reg2d_reg2d(nx, ny, data, gridID, array);
  else
    remap_nn_reg2d_nonreg2d(nx, ny, data, gridID, array);
}

static int
randomInit(const int operatorID)
{
  unsigned int seed = Options::Random_Seed;
  operatorInputArg(cdoOperatorEnter(operatorID));
  if (operatorArgc() < 1) cdoAbort("Too few arguments!");
  if (operatorArgc() > 2) cdoAbort("Too many arguments!");
  const auto gridID = cdoDefineGrid(cdoOperatorArgv(0));
  if (operatorArgc() == 2)
    {
      const auto idum = parameter2int(cdoOperatorArgv(1));
      if (idum >= 0 && idum < 0x7FFFFFFF) seed = idum;
    }
  std::srand(seed);
  return gridID;
}

static void
randomCompute(const size_t gridsize, Varray<double> &array)
{
  for (size_t i = 0; i < gridsize; ++i) array[i] = ((double) std::rand()) / ((double) RAND_MAX);
}

static void
sincosCompute(const size_t gridsize, Varray<double> &array, const Varray<double> &xvals, const Varray<double> &yvals)
{
  for (size_t i = 0; i < gridsize; ++i) array[i] = std::cos(1.0 * xvals[i]) * std::sin(2.0 * yvals[i]);
}

static void
coshillCompute(const size_t gridsize, Varray<double> &array, const Varray<double> &xvals, const Varray<double> &yvals)
{
  for (size_t i = 0; i < gridsize; ++i) array[i] = 2 - std::cos(std::acos(std::cos(xvals[i]) * std::cos(yvals[i])) / 1.2);
}

static void
testfieldCompute(const size_t gridsize, Varray<double> &array, const Varray<double> &xvals, const Varray<double> &yvals)
{
  double xyz[3];
  for (size_t i = 0; i < gridsize; ++i)
    {
      gcLLtoXYZ(xvals[i], yvals[i], xyz);
      const auto x = xyz[0];
      const auto y = xyz[1];
      const auto z = xyz[2];
      array[i] = 1.0 + std::pow(x, 8.0) + std::exp(2.0 * y * y * y) + std::exp(2.0 * x * x) + 10.0 * x * y * z;
    }
}

static void
unpackData(const size_t datasize, Varray<double> &data, const double scale, const double offset, const unsigned short *zdata)
{
  for (size_t i = 0; i < datasize; ++i) data[i] = zdata[i] / scale - offset;
}

static int
definePointGrid()
{
  const auto gridID = gridCreate(GRID_LONLAT, 1);
  gridDefXsize(gridID, 1);
  gridDefYsize(gridID, 1);
  const double value = 0.0;
  gridDefXvals(gridID, &value);
  gridDefYvals(gridID, &value);

  return gridID;
}

static int
defineZaxis(const bool lstdatm, const int nlevels, double *levels)
{
  int zaxisID = -1;

  if (lstdatm)
    {
      zaxisID = zaxisCreate(ZAXIS_HEIGHT, nlevels);
      zaxisDefLevels(zaxisID, levels);
      cdiDefKeyString(zaxisID, CDI_GLOBAL, CDI_KEY_NAME, "level");
      cdiDefKeyString(zaxisID, CDI_GLOBAL, CDI_KEY_LONGNAME, "Level");
      cdiDefKeyString(zaxisID, CDI_GLOBAL, CDI_KEY_UNITS, "m");
    }
  else
    {
      zaxisID = zaxisFromName("surface");
    }

  return zaxisID;
}

static void
definePressureAttributes(const int vlistID, const int varID)
{
  cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, "P");
  vlistDefVarParam(vlistID, varID, cdiEncodeParam(1, 255, 255));
  cdiDefKeyString(vlistID, varID, CDI_KEY_STDNAME, "air_pressure");
  cdiDefKeyString(vlistID, varID, CDI_KEY_LONGNAME, "pressure");
  cdiDefKeyString(vlistID, varID, CDI_KEY_UNITS, "hPa");
}

static void
defineTemperatureAttributes(const int vlistID, const int varID)
{
  cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, "T");
  vlistDefVarParam(vlistID, varID, cdiEncodeParam(130, 128, 255));
  cdiDefKeyString(vlistID, varID, CDI_KEY_STDNAME, var_stdname(air_temperature));
  cdiDefKeyString(vlistID, varID, CDI_KEY_LONGNAME, "temperature");
  cdiDefKeyString(vlistID, varID, CDI_KEY_UNITS, "K");
}

void *
Vargen(void *process)
{
  constexpr size_t nlat = 360;
  constexpr size_t nlon = 720;
  double lon[nlon], lat[nlat];
  int ntimesteps, nlevels = 1;
  int gridID = -1, gridIDdata = -1;
  double rstart = 0.0, rstop = 0.0, rinc = 0.0;
  double rconst = 0.0;
  std::vector<double> levels;

  cdoInitialize(process);

  // clang-format off
  const auto RANDOM    = cdoOperatorAdd("random",    0, 0, "grid description file or name, <seed>");
  const auto SINCOS    = cdoOperatorAdd("sincos",    0, 0, "grid description file or name");
  const auto COSHILL   = cdoOperatorAdd("coshill",   0, 0, "grid description file or name");
  const auto TESTFIELD = cdoOperatorAdd("testfield", 0, 0, "grid description file or name");
  const auto CONST     = cdoOperatorAdd("const",     0, 0, "constant value, grid description file or name");
  const auto SEQ       = cdoOperatorAdd("seq",       0, 0, "start, end, <increment>");
  const auto TOPO      = cdoOperatorAdd("topo",      0, 0, nullptr);
  const auto TEMP      = cdoOperatorAdd("temp",      0, 0, nullptr);
  const auto MASK      = cdoOperatorAdd("mask",      0, 0, nullptr);
  const auto STDATM    = cdoOperatorAdd("stdatm",    0, 0, "height levels [m]");
  // clang-format on

  const auto operatorID = cdoOperatorID();

  if (operatorID == RANDOM)
    {
      gridID = randomInit(operatorID);
    }
  else if (operatorID == SINCOS || operatorID == COSHILL || operatorID == TESTFIELD)
    {
      operatorInputArg(cdoOperatorEnter(operatorID));
      operatorCheckArgc(1);
      gridID = cdoDefineGrid(cdoOperatorArgv(0));
    }
  else if (operatorID == CONST)
    {
      operatorInputArg(cdoOperatorEnter(operatorID));
      operatorCheckArgc(2);
      rconst = parameter2double(cdoOperatorArgv(0));
      gridID = cdoDefineGrid(cdoOperatorArgv(1));
    }
  else if (operatorID == TOPO || operatorID == TEMP || operatorID == MASK)
    {
      gridIDdata = gridCreate(GRID_LONLAT, nlon * nlat);
      gridDefXsize(gridIDdata, nlon);
      gridDefYsize(gridIDdata, nlat);

      for (size_t i = 0; i < nlon; i++) lon[i] = -179.75 + i * 0.5;
      for (size_t i = 0; i < nlat; i++) lat[i] = -89.75 + i * 0.5;

      gridDefXvals(gridIDdata, lon);
      gridDefYvals(gridIDdata, lat);

      gridID = gridIDdata;

      if (operatorArgc() == 1) gridID = cdoDefineGrid(cdoOperatorArgv(0));
      if (operatorArgc() > 1) cdoAbort("Too many arguments!");
    }
  else if (operatorID == SEQ)
    {
      operatorInputArg(cdoOperatorEnter(operatorID));
      if (operatorArgc() < 2) cdoAbort("Too few arguments!");
      if (operatorArgc() > 3) cdoAbort("Too many arguments!");

      rstart = parameter2double(cdoOperatorArgv(0));
      rstop = parameter2double(cdoOperatorArgv(1));
      rinc = (operatorArgc() == 3) ? parameter2double(cdoOperatorArgv(2)) : 1;
      if (DBL_IS_EQUAL(rinc, 0.0)) cdoAbort("Increment is zero!");

      gridID = definePointGrid();
    }
  else if (operatorID == STDATM)
    {
      operatorInputArg(cdoOperatorEnter(operatorID));
      levels = cdoArgvToFlt(cdoGetOperArgv());
      nlevels = levels.size();

      if (Options::cdoVerbose)
        for (int i = 0; i < nlevels; ++i) printf("levels %d: %g\n", i, levels[i]);

      gridID = definePointGrid();
    }

  const auto zaxisID = defineZaxis(operatorID == STDATM, nlevels, levels.data());

  const auto vlistID = vlistCreate();

  const auto timetype = (operatorID == SEQ) ? TIME_VARYING : TIME_CONSTANT;

  auto varID = vlistDefVar(vlistID, gridID, zaxisID, timetype);
  /*
     For the standard atmosphere two output variables are generated: pressure
     and temperatur. The first (varID) is pressure, second (varID2) is
     temperatur. Add an additional variable for the standard atmosphere.
   */
  const int varID2 = (operatorID == STDATM) ? vlistDefVar(vlistID, gridID, zaxisID, TIME_CONSTANT) : -1;

  if (operatorID == MASK) vlistDefVarDatatype(vlistID, varID, CDI_DATATYPE_INT8);

  if (operatorID == STDATM)
    {
      definePressureAttributes(vlistID, varID);
      defineTemperatureAttributes(vlistID, varID2);
    }
  else
    {
      cdiDefKeyString(vlistID, varID, CDI_KEY_NAME, cdoOperatorName(operatorID));
      if (operatorID == TOPO) cdiDefKeyString(vlistID, varID, CDI_KEY_UNITS, "m");
      if (operatorID == TEMP) cdiDefKeyString(vlistID, varID, CDI_KEY_UNITS, "K");
    }

  const auto taxisID = cdoTaxisCreate(TAXIS_RELATIVE);
  vlistDefTaxis(vlistID, taxisID);

  if (operatorID == RANDOM || operatorID == SINCOS || operatorID == COSHILL || operatorID == TESTFIELD || operatorID == CONST
      || operatorID == TOPO || operatorID == TEMP || operatorID == MASK || operatorID == STDATM)
    vlistDefNtsteps(vlistID, 1);

  const auto streamID = cdoOpenWrite(0);

  cdoDefVlist(streamID, vlistID);

  const auto gridsize = gridInqSize(gridID);
  Varray<double> array(gridsize);

  if (operatorID == SEQ)
    ntimesteps = 1.001 + ((rstop - rstart) / rinc);
  else
    {
      vlistDefNtsteps(vlistID, 0);
      ntimesteps = 1;
    }

  const auto julday = date_to_julday(CALENDAR_PROLEPTIC, 10101);

  const auto nvars = vlistNvars(vlistID);

  for (int tsID = 0; tsID < ntimesteps; tsID++)
    {
      const double rval = rstart + rinc * tsID;
      const auto vdate = julday_to_date(CALENDAR_PROLEPTIC, julday + tsID);
      const int vtime = 0;
      taxisDefVdate(taxisID, vdate);
      taxisDefVtime(taxisID, vtime);
      cdoDefTimestep(streamID, tsID);

      for (varID = 0; varID < nvars; varID++)
        {
          nlevels = zaxisInqSize(vlistInqVarZaxis(vlistID, varID));
          for (int levelID = 0; levelID < nlevels; levelID++)
            {
              cdoDefRecord(streamID, varID, levelID);

              if (operatorID == RANDOM)
                {
                  randomCompute(gridsize, array);
                }
              else if (operatorID == SINCOS || operatorID == COSHILL || operatorID == TESTFIELD)
                {
                  Varray<double> xvals(gridsize), yvals(gridsize);

                  if (grid_is_distance_generic(gridID))
                    {
                      conv_generic_grid(gridID, gridsize, xvals, yvals);
                    }
                  else
                    {
                      if (gridInqType(gridID) == GRID_GME) gridID = gridToUnstructured(gridID, 0);

                      if (gridInqType(gridID) != GRID_UNSTRUCTURED && gridInqType(gridID) != GRID_CURVILINEAR)
                        gridID = gridToCurvilinear(gridID, 0);

                      gridInqXvals(gridID, xvals.data());
                      gridInqYvals(gridID, yvals.data());

                      // Convert lat/lon units if required
                      cdo_grid_to_radian(gridID, CDI_XAXIS, gridsize, xvals.data(), "grid center lon");
                      cdo_grid_to_radian(gridID, CDI_YAXIS, gridsize, yvals.data(), "grid center lat");
                    }

                  if (operatorID == SINCOS)
                    {
                      sincosCompute(gridsize, array, xvals, yvals);
                    }
                  else if (operatorID == COSHILL)
                    {
                      coshillCompute(gridsize, array, xvals, yvals);
                    }
                  else if (operatorID == TESTFIELD)
                    {
                      testfieldCompute(gridsize, array, xvals, yvals);
                    }
                }
              else if (operatorID == CONST)
                {
                  for (size_t i = 0; i < gridsize; i++) array[i] = rconst;
                }
              else if (operatorID == TOPO || operatorID == TEMP || operatorID == MASK)
                {
                  const auto datasize = gridInqSize(gridIDdata);
                  Varray<double> data(datasize);

                  // clang-format off
                  if      (operatorID == TOPO) unpackData(datasize, data, etopo_scale, etopo_offset, etopo);
                  else if (operatorID == TEMP) unpackData(datasize, data, temp_scale, temp_offset, temp);
                  else if (operatorID == MASK) unpackData(datasize, data, mask_scale, mask_offset, mask);
                  // clang-format on

                  if (gridID != gridIDdata && gridIDdata != -1)
                    {
                      remap_nn_reg2d(nlon, nlat, data.data(), gridID, array.data());
                    }
                  else
                    {
                      for (size_t i = 0; i < gridsize; ++i) array[i] = data[i];
                    }
                }
              else if (operatorID == SEQ)
                {
                  array[0] = rval;
                }
              else if (operatorID == STDATM)
                {
                  array[0] = (varID == varID2) ? std_atm_temperatur(levels[levelID]) : std_atm_pressure(levels[levelID]);
                }

              cdoWriteRecord(streamID, array.data(), 0);
            }
        }
    }

  cdoStreamClose(streamID);

  vlistDestroy(vlistID);

  cdoFinish();

  return nullptr;
}
