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

*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>

static int
gentpngrid(int gridID1)
{
  const auto nlon1 = gridInqXsize(gridID1);
  const auto nlat1 = gridInqYsize(gridID1);

  const auto nlon2 = nlon1;
  const auto nlat2 = nlat1 + 2;

  const auto gridtype = gridInqType(gridID1);

  const auto gridID2 = gridCreate(gridtype, nlon2 * nlat2);
  gridDefXsize(gridID2, nlon2);
  gridDefYsize(gridID2, nlat2);

  int datatype = CDI_UNDEFID;
  cdiInqKeyInt(gridID1, CDI_GLOBAL, CDI_KEY_DATATYPE, &datatype);
  cdiDefKeyInt(gridID2, CDI_GLOBAL, CDI_KEY_DATATYPE, datatype);

  grid_copy_keys(gridID1, gridID2);

  char xunits[CDI_MAX_NAME], yunits[CDI_MAX_NAME];
  int length = CDI_MAX_NAME;
  cdiInqKeyString(gridID1, CDI_XAXIS, CDI_KEY_UNITS, xunits, &length);
  length = CDI_MAX_NAME;
  cdiInqKeyString(gridID1, CDI_YAXIS, CDI_KEY_UNITS, yunits, &length);

  if (gridHasCoordinates(gridID1))
    {
      if (gridtype == GRID_CURVILINEAR)
        {
          Varray<double> xvals1(nlon1 * nlat1), yvals1(nlon1 * nlat1);
          Varray<double> xvals2(nlon2 * nlat2), yvals2(nlon2 * nlat2);

          gridInqXvals(gridID1, xvals1.data());
          gridInqYvals(gridID1, yvals1.data());

          for (size_t ilat = 0; ilat < nlat1; ilat++)
            {
              for (size_t ilon = 0; ilon < nlon1; ilon++)
                {
                  xvals2[(ilat + 2) * nlon1 + ilon] = xvals1[ilat * nlon1 + ilon];
                  yvals2[(ilat + 2) * nlon1 + ilon] = yvals1[ilat * nlon1 + ilon];
                }
            }

          for (size_t ilon = 0; ilon < nlon1; ilon++)
            {
              const auto ilonr = nlon1 - ilon - 1;
              xvals2[1 * nlon1 + ilon] = xvals2[2 * nlon1 + ilonr]; // syncronise line 2 with line 3
              xvals2[0 * nlon1 + ilon] = xvals2[3 * nlon1 + ilonr]; // syncronise line 1 with line 4
              yvals2[1 * nlon1 + ilon] = yvals2[2 * nlon1 + ilonr]; // syncronise line 2 with line 3
              yvals2[0 * nlon1 + ilon] = yvals2[3 * nlon1 + ilonr]; // syncronise line 1 with line 4
            }

          gridDefXvals(gridID2, xvals2.data());
          gridDefYvals(gridID2, yvals2.data());
        }
    }

  if (gridHasBounds(gridID1))
    {
      if (gridtype == GRID_CURVILINEAR)
        {
          Varray<double> xbounds1(4 * nlon1 * nlat1), ybounds1(4 * nlon1 * nlat1);
          Varray<double> xbounds2(4 * nlon2 * nlat2), ybounds2(4 * nlon2 * nlat2);

          gridInqXbounds(gridID1, xbounds1.data());
          gridInqYbounds(gridID1, ybounds1.data());

          if (gridtype == GRID_CURVILINEAR)
            {
              gridDefNvertex(gridID2, 4);

              const auto nlon4 = 4 * nlon1;
              for (size_t ilat = 0; ilat < nlat1; ilat++)
                {
                  for (size_t ilon = 0; ilon < nlon4; ilon++)
                    {
                      xbounds2[(ilat + 2) * nlon4 + ilon] = xbounds1[ilat * nlon4 + ilon];
                      ybounds2[(ilat + 2) * nlon4 + ilon] = ybounds1[ilat * nlon4 + ilon];
                    }
                }

              for (size_t ilon = 0; ilon < nlon1; ilon++)
                {
                  const auto ilonr = nlon1 - ilon - 1;
                  for (size_t k = 0; k < 4; ++k)
                    {
                      const auto kr = 3 - k;
                      xbounds2[1 * nlon4 + 4 * ilon + k] = xbounds2[2 * nlon4 + 4 * ilonr + kr];
                      xbounds2[0 * nlon4 + 4 * ilon + k] = xbounds2[3 * nlon4 + 4 * ilonr + kr];
                      ybounds2[1 * nlon4 + 4 * ilon + k] = ybounds2[2 * nlon4 + 4 * ilonr + kr];
                      ybounds2[0 * nlon4 + 4 * ilon + k] = ybounds2[3 * nlon4 + 4 * ilonr + kr];
                    }
                }
              /*
              for (size_t ilon = 0; ilon < 4*nlon1; ilon++)
                {
                  const auto ilonr = nlon4 - ilon - 1;
                  xbounds2[1*nlon4 + ilon] = xbounds2[2*nlon4 + ilonr]; xbounds2[0*nlon4 + ilon] = xbounds2[3*nlon4 + ilonr];
                  ybounds2[1*nlon4 + ilon] = ybounds2[2*nlon4 + ilonr]; ybounds2[0*nlon4 + ilon] = ybounds2[3*nlon4 + ilonr];
                }
              */
            }

          gridDefXbounds(gridID2, xbounds2.data());
          gridDefYbounds(gridID2, ybounds2.data());
        }
    }

  return gridID2;
}

static int
gengrid(int gridID1, long lhalo, long rhalo)
{
  long i, ilat, ilon;
  double cpi2 = M_PI * 2;

  const long nlon1 = gridInqXsize(gridID1);
  const long nlat1 = gridInqYsize(gridID1);

  const long nlon2 = nlon1 + lhalo + rhalo;
  const long nlat2 = nlat1;

  long nmin = 0;
  long nmax = nlon1;
  if (lhalo < 0) nmin = -lhalo;
  if (rhalo < 0) nmax += rhalo;
  // printf("nlon1=%d, nlon2=%d, lhalo=%d, rhalo=%d nmin=%d, nmax=%d\n", nlon1, nlon2, lhalo, rhalo, nmin, nmax);

  const auto gridtype = gridInqType(gridID1);

  const auto gridID2 = gridCreate(gridtype, nlon2 * nlat2);
  gridDefXsize(gridID2, nlon2);
  gridDefYsize(gridID2, nlat2);

  int datatype = CDI_UNDEFID;
  cdiInqKeyInt(gridID1, CDI_GLOBAL, CDI_KEY_DATATYPE, &datatype);
  cdiDefKeyInt(gridID2, CDI_GLOBAL, CDI_KEY_DATATYPE, datatype);

  grid_copy_keys(gridID1, gridID2);

  char xunits[CDI_MAX_NAME], yunits[CDI_MAX_NAME];
  int length = CDI_MAX_NAME;
  cdiInqKeyString(gridID1, CDI_XAXIS, CDI_KEY_UNITS, xunits, &length);
  length = CDI_MAX_NAME;
  cdiInqKeyString(gridID1, CDI_YAXIS, CDI_KEY_UNITS, yunits, &length);

  if (memcmp(xunits, "degree", 6) == 0) cpi2 *= RAD2DEG;

  if (gridHasCoordinates(gridID1))
    {
      Varray<double> xvals1, yvals1, xvals2, yvals2;
      if (gridtype == GRID_CURVILINEAR)
        {
          xvals1.resize(nlon1 * nlat1);
          yvals1.resize(nlon1 * nlat1);
          xvals2.resize(nlon2 * nlat2);
          yvals2.resize(nlon2 * nlat2);
        }
      else
        {
          xvals1.resize(nlon1);
          yvals1.resize(nlat1);
          xvals2.resize(nlon2);
          yvals2.resize(nlat2);
        }

      double *pxvals2 = xvals2.data();
      double *pyvals2 = yvals2.data();

      gridInqXvals(gridID1, xvals1.data());
      gridInqYvals(gridID1, yvals1.data());

      if (gridtype == GRID_CURVILINEAR)
        {
          for (ilat = 0; ilat < nlat2; ilat++)
            {
              for (ilon = nlon1 - lhalo; ilon < nlon1; ilon++)
                {
                  *pxvals2++ = xvals1[ilat * nlon1 + ilon];
                  *pyvals2++ = yvals1[ilat * nlon1 + ilon];
                }

              for (ilon = nmin; ilon < nmax; ilon++)
                {
                  *pxvals2++ = xvals1[ilat * nlon1 + ilon];
                  *pyvals2++ = yvals1[ilat * nlon1 + ilon];
                }

              for (ilon = 0; ilon < rhalo; ilon++)
                {
                  *pxvals2++ = xvals1[ilat * nlon1 + ilon];
                  *pyvals2++ = yvals1[ilat * nlon1 + ilon];
                }
            }
        }
      else
        {
          for (i = nlon1 - lhalo; i < nlon1; i++) *pxvals2++ = xvals1[i] - cpi2;
          for (i = nmin; i < nmax; i++) *pxvals2++ = xvals1[i];
          for (i = 0; i < rhalo; i++) *pxvals2++ = xvals1[i] + cpi2;

          for (i = 0; i < nlat1; i++) yvals2[i] = yvals1[i];
        }
      /*
      for (i = 0; i < nlat2; i++) printf("lat : %d %g\n", i+1, yvals2[i]);
      for (i = 0; i < nlon2; i++) printf("lon : %d %g\n", i+1, xvals2[i]);
      */
      gridDefXvals(gridID2, xvals2.data());
      gridDefYvals(gridID2, yvals2.data());
    }

  if (gridHasBounds(gridID1))
    {
      Varray<double> xbounds1, ybounds1, xbounds2, ybounds2;
      if (gridtype == GRID_CURVILINEAR)
        {
          xbounds1.resize(4 * nlon1 * nlat1);
          ybounds1.resize(4 * nlon1 * nlat1);
          xbounds2.resize(4 * nlon2 * nlat2);
          ybounds2.resize(4 * nlon2 * nlat2);
        }
      else
        {
          xbounds1.resize(2 * nlon1);
          ybounds1.resize(2 * nlat1);
          xbounds2.resize(2 * nlon2);
          ybounds2.resize(2 * nlat2);
        }

      double *pxbounds2 = xbounds2.data();
      double *pybounds2 = ybounds2.data();

      gridInqXbounds(gridID1, xbounds1.data());
      gridInqYbounds(gridID1, ybounds1.data());

      if (gridtype == GRID_CURVILINEAR)
        {
          gridDefNvertex(gridID2, 4);
          for (ilat = 0; ilat < nlat1; ilat++)
            {
              for (ilon = 4 * (nlon1 - lhalo); ilon < 4 * nlon1; ilon++)
                {
                  *pxbounds2++ = xbounds1[4 * ilat * nlon1 + ilon];
                  *pybounds2++ = ybounds1[4 * ilat * nlon1 + ilon];
                }

              for (ilon = 4 * nmin; ilon < 4 * nmax; ilon++)
                {
                  *pxbounds2++ = xbounds1[4 * ilat * nlon1 + ilon];
                  *pybounds2++ = ybounds1[4 * ilat * nlon1 + ilon];
                }

              for (ilon = 0; ilon < 4 * rhalo; ilon++)
                {
                  *pxbounds2++ = xbounds1[4 * ilat * nlon1 + ilon];
                  *pybounds2++ = ybounds1[4 * ilat * nlon1 + ilon];
                }
            }
        }
      else
        {
          gridDefNvertex(gridID2, 2);
          for (i = 2 * (nlon1 - lhalo); i < 2 * nlon1; i++) *pxbounds2++ = xbounds1[i] - cpi2;
          for (i = 2 * nmin; i < 2 * nmax; i++) *pxbounds2++ = xbounds1[i];
          for (i = 0; i < 2 * rhalo; i++) *pxbounds2++ = xbounds1[i] + cpi2;

          for (i = 0; i < 2 * nlat2; i++) ybounds2[i] = ybounds1[i];
        }

      gridDefXbounds(gridID2, xbounds2.data());
      gridDefYbounds(gridID2, ybounds2.data());
    }

  return gridID2;
}

static int
genindexgrid(int gridID1, long &lhalo, long &rhalo)
{
  operator_check_argc(2);

  lhalo = parameter_to_long(cdo_operator_argv(0));
  rhalo = parameter_to_long(cdo_operator_argv(1));

  const long nlon1 = gridInqXsize(gridID1);

  if (lhalo > nlon1)
    {
      lhalo = nlon1;
      cdo_warning("left halo out of range. Set to %d.", lhalo);
    }

  if (lhalo < 0 && -(lhalo) > nlon1 / 2)
    {
      lhalo = -nlon1 / 2;
      cdo_warning("left halo out of range. Set to %d.", lhalo);
    }

  if (rhalo > nlon1)
    {
      rhalo = nlon1;
      cdo_warning("right halo out of range. Set to %d.", rhalo);
    }

  if (rhalo < 0 && -(rhalo) > nlon1 / 2)
    {
      rhalo = -nlon1 / 2;
      cdo_warning("right halo out of range. Set to %d.", rhalo);
    }

  const int gridID2 = gengrid(gridID1, lhalo, rhalo);

  return gridID2;
}

static void
halo(double *array1, int gridID1, double *array2, long lhalo, long rhalo)
{
  long nlon1 = gridInqXsize(gridID1);
  long nlat = gridInqYsize(gridID1);

  long nmin = 0;
  long nmax = nlon1;
  if (lhalo < 0) nmin = -lhalo;
  if (rhalo < 0) nmax += rhalo;

  for (long ilat = 0; ilat < nlat; ilat++)
    {
      for (long ilon = nlon1 - lhalo; ilon < nlon1; ilon++) *array2++ = array1[ilat * nlon1 + ilon];

      for (long ilon = nmin; ilon < nmax; ilon++) *array2++ = array1[ilat * nlon1 + ilon];

      for (long ilon = 0; ilon < rhalo; ilon++) *array2++ = array1[ilat * nlon1 + ilon];
    }
}

static void
tpnhalo(double *array1, int gridID1, double *array2)
{
  const auto nlon = gridInqXsize(gridID1);
  const auto nlat = gridInqYsize(gridID1);

  for (size_t ilat = 0; ilat < nlat; ilat++)
    for (size_t ilon = 0; ilon < nlon; ilon++) array2[(ilat + 2) * nlon + ilon] = array1[ilat * nlon + ilon];

  for (size_t ilon = 0; ilon < nlon; ilon++)
    {
      const size_t ilonr = nlon - ilon - 1;
      array2[1 * nlon + ilon] = array2[2 * nlon + ilonr];  // syncronise line 2 with line 3
      array2[0 * nlon + ilon] = array2[3 * nlon + ilonr];  // syncronise line 1 with line 4
    }
}

void *
Sethalo(void *process)
{
  int gridID1 = -1, gridID2;
  int index;
  long lhalo = 0, rhalo = 0;

  cdo_initialize(process);

  // clang-format off
  const auto SETHALO = cdo_operator_add("sethalo", 0, 0, nullptr);
                       cdo_operator_add("tpnhalo", 0, 0, nullptr);
  // clang-format on

  const auto operatorID = cdo_operator_id();

  const auto streamID1 = cdo_open_read(0);
  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);

  auto ngrids = vlistNgrids(vlistID1);
  int ndiffgrids = 0;
  for (index = 1; index < ngrids; index++)
    if (vlistGrid(vlistID1, 0) != vlistGrid(vlistID1, index)) ndiffgrids++;

  for (index = 0; index < ngrids; index++)
    {
      gridID1 = vlistGrid(vlistID1, index);
      const auto gridtype = gridInqType(gridID1);
      if (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN) break;
      if (gridtype == GRID_CURVILINEAR) break;
      if (gridtype == GRID_GENERIC && gridInqXsize(gridID1) > 0 && gridInqYsize(gridID1) > 0) break;
    }

  if (gridInqType(gridID1) == GRID_GAUSSIAN_REDUCED)
    cdo_abort("Gaussian reduced grid found. Use option -R to convert it to a regular grid!");

  if (index == ngrids) cdo_abort("No regular grid found!");
  if (ndiffgrids > 0) cdo_abort("Too many different grids!");

  if (operatorID == SETHALO)
    {
      operator_input_arg("left and right halo");
      gridID2 = genindexgrid(gridID1, lhalo, rhalo);
    }
  else
    {
      gridID2 = gentpngrid(gridID1);
    }

  const auto vlistID2 = vlistDuplicate(vlistID1);

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

  ngrids = vlistNgrids(vlistID1);
  for (index = 0; index < ngrids; index++)
    {
      if (gridID1 == vlistGrid(vlistID1, index))
        {
          vlistChangeGridIndex(vlistID2, index, gridID2);
          break;
        }
    }

  const auto nvars = vlistNvars(vlistID1);
  std::vector<bool> vars(nvars);
  for (int varID = 0; varID < nvars; varID++) vars[varID] = (gridID1 == vlistInqVarGrid(vlistID1, varID));

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

  const auto gridsize = gridInqSize(gridID1);
  Varray<double> array1(gridsize);

  const auto gridsize2 = gridInqSize(gridID2);
  Varray<double> array2(gridsize2);

  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 recID = 0; recID < nrecs; recID++)
        {
          int varID, levelID;
          cdo_inq_record(streamID1, &varID, &levelID);

          if (vars[varID])
            {
              size_t nmiss;
              cdo_read_record(streamID1, array1.data(), &nmiss);

              if (operatorID == SETHALO)
                halo(array1.data(), gridID1, array2.data(), lhalo, rhalo);
              else
                tpnhalo(array1.data(), gridID1, array2.data());

              if (nmiss)
                {
                  const auto missval = vlistInqVarMissval(vlistID1, varID);
                  nmiss = varray_num_mv(gridsize2, array2, missval);
                }

              cdo_def_record(streamID2, varID, levelID);
              cdo_write_record(streamID2, array2.data(), nmiss);
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
