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

      Gridboxstat    gridboxrange        Gridbox range
      Gridboxstat    gridboxmin          Gridbox minimum
      Gridboxstat    gridboxmax          Gridbox maximum
      Gridboxstat    gridboxsum          Gridbox sum
      Gridboxstat    gridboxmean         Gridbox mean
      Gridboxstat    gridboxavg          Gridbox average
      Gridboxstat    gridboxstd          Gridbox standard deviation
      Gridboxstat    gridboxstd1         Gridbox standard deviation [Normalize by (n-1)]
      Gridboxstat    gridboxvar          Gridbox variance
      Gridboxstat    gridboxvar1         Gridbox variance [Normalize by (n-1)]
*/

#include <cdi.h>

#include "functs.h"
#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "cdo_options.h"
#include "progress.h"
#include "cimdOmp.h"

static void
genBoxGridReg2D(int gridID1, size_t xinc, size_t yinc, int gridID2)
{
  const auto nlon1 = gridInqXsize(gridID1);
  const auto nlat1 = gridInqYsize(gridID1);
  const auto nlon2 = gridInqXsize(gridID2);
  const auto nlat2 = gridInqYsize(gridID2);

  {
    Varray<double> xvals1(nlon1), yvals1(nlat1);
    Varray<double> xvals2(nlon2), yvals2(nlat2);
    gridInqXvals(gridID1, xvals1.data());
    gridInqYvals(gridID1, yvals1.data());

    size_t j = 0;
    for (size_t i = 0; i < nlon1; i += xinc)
      {
        auto i1 = i + (xinc - 1);
        if (i1 >= nlon1 - 1) i1 = nlon1 - 1;
        xvals2[j] = xvals1[i] + (xvals1[i1] - xvals1[i]) / 2.;
        j++;
      }

    j = 0;
    for (size_t i = 0; i < nlat1; i += yinc)
      {
        auto i1 = i + (yinc - 1);
        if (i1 >= nlat1 - 1) i1 = nlat1 - 1;
        yvals2[j] = yvals1[i] + (yvals1[i1] - yvals1[i]) / 2;
        j++;
      }

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

  if (gridHasBounds(gridID1))
    {
      Varray<double> grid1_corner_lon(2 * nlon1), grid1_corner_lat(2 * nlat1);
      Varray<double> grid2_corner_lon(2 * nlon2), grid2_corner_lat(2 * nlat2);
      gridInqXbounds(gridID1, grid1_corner_lon.data());
      gridInqYbounds(gridID1, grid1_corner_lat.data());

      size_t j = 0;
      for (size_t i = 0; i < nlon1; i += xinc)
        {
          auto i1 = i + (xinc - 1);
          if (i1 >= nlon1 - 1) i1 = nlon1 - 1;
          grid2_corner_lon[2 * j] = grid1_corner_lon[2 * i];
          grid2_corner_lon[2 * j + 1] = grid1_corner_lon[2 * i1 + 1];
          j++;
        }

      j = 0;
      for (size_t i = 0; i < nlat1; i += yinc)
        {
          auto i1 = i + (yinc - 1);
          if (i1 >= nlat1 - 1) i1 = nlat1 - 1;
          grid2_corner_lat[2 * j] = grid1_corner_lat[2 * i];
          grid2_corner_lat[2 * j + 1] = grid1_corner_lat[2 * i1 + 1];
          j++;
        }

      gridDefNvertex(gridID2, 2);
      gridDefXbounds(gridID2, grid2_corner_lon.data());
      gridDefYbounds(gridID2, grid2_corner_lat.data());
    }
}

static
int find_corner(size_t g1_add, size_t off1, size_t off2, const Varray<double> &grid1_corner_lon, const Varray<double> &grid1_corner_lat)
{
  int c_flag[4] = { 0 };
  for (int corner = 0; corner < 4; corner++)
    {
      const auto lon = grid1_corner_lon[4 * g1_add + corner];
      const auto lat = grid1_corner_lat[4 * g1_add + corner];
      const auto g1_add2 = g1_add + off1;
      const auto g1_add3 = g1_add + off2;
      for (int corner2 = 0; corner2 < 4; corner2++)
        {
          const auto lon2 = grid1_corner_lon[4 * g1_add2 + corner2];
          const auto lat2 = grid1_corner_lat[4 * g1_add2 + corner2];
          const auto lon3 = grid1_corner_lon[4 * g1_add3 + corner2];
          const auto lat3 = grid1_corner_lat[4 * g1_add3 + corner2];
          if ((IS_EQUAL(lon2, lon) && IS_EQUAL(lat2, lat)) || (IS_EQUAL(lon3, lon) && IS_EQUAL(lat3, lat)))
            c_flag[corner] = 1;
        }
    }

  if (c_flag[0] + c_flag[1] + c_flag[2] + c_flag[3] < 3) cdoWarning("Found two matching corners!");

  int corner = 0;
  for (corner = 0; corner < 4; corner++)
    if (!c_flag[corner]) break;

  return corner;
}

static void
genBoxGridCurv2D(int gridID1, size_t xinc, size_t yinc, int gridID2)
{
  const auto nlon1 = gridInqXsize(gridID1);
  const auto nlat1 = gridInqYsize(gridID1);
  const auto nlon2 = gridInqXsize(gridID2);
  const auto nlat2 = gridInqYsize(gridID2);
  const auto gridsize1 = gridInqSize(gridID1);
  const auto gridsize2 = gridInqSize(gridID2);

  const auto circular = gridIsCircular(gridID1);
  double xvals2_0 = 0.0;

  bool lGridHasBounds = gridHasBounds(gridID1);

  Varray<double> xvals1(gridsize1), yvals1(gridsize1);
  Varray<double> xvals2(gridsize2), yvals2(gridsize2);
  gridInqXvals(gridID1, xvals1.data());
  gridInqYvals(gridID1, yvals1.data());

  // Convert lat/lon units if required
  cdo_grid_to_degree(gridID1, CDI_XAXIS, gridsize1, xvals1.data(), "grid center lon");
  cdo_grid_to_degree(gridID1, CDI_YAXIS, gridsize1, yvals1.data(), "grid center lat");

  Varray<double> grid1_corner_lon, grid1_corner_lat;
  Varray<double> grid2_corner_lon, grid2_corner_lat;
  if (lGridHasBounds)
    {
      grid1_corner_lon.resize(4 * gridsize1);
      grid1_corner_lat.resize(4 * gridsize1);
      grid2_corner_lon.resize(4 * gridsize2);
      grid2_corner_lat.resize(4 * gridsize2);
      gridInqXbounds(gridID1, grid1_corner_lon.data());
      gridInqYbounds(gridID1, grid1_corner_lat.data());

      // Convert lat/lon units if required
      cdo_grid_to_degree(gridID1, CDI_XAXIS, 4 * gridsize1, grid1_corner_lon.data(), "grid corner lon");
      cdo_grid_to_degree(gridID1, CDI_YAXIS, 4 * gridsize1, grid1_corner_lat.data(), "grid corner lat");
   }

  // Process grid2 bounds
  const double area_norm = xinc * yinc;
  for (size_t y2 = 0; y2 < nlat2; y2++)
    {
      for (size_t x2 = 0; x2 < nlon2; x2++)
        {
          const auto g2_add = (y2 * nlon2 + x2);
          double on_up, on_lo, ox_up, ox_lo, an_le, an_ri, ax_le, ax_ri;
          on_up = on_lo = 360.0;
          ox_up = ox_lo = -360.0;
          an_ri = an_le = 90.0;
          ax_ri = ax_le = -90.0;

          for (size_t y1 = y2 * yinc; y1 < yinc * (y2 + 1); y1++)
            {
              auto use_y1 = (y1 >= nlat1) ? nlat1 - 1 : y1;
              for (size_t x1 = x2 * xinc; x1 < xinc * (x2 + 1); x1++)
                {
                  auto use_x1 = x1;
                  if (x1 >= nlon1)
                    {
                      if (circular && use_y1 == y1)
                        use_y1 -= 1;
                      else
                        use_x1 = nlon1 - 1;
                    }

                  const auto g1_add = (use_y1 * nlon1) + use_x1;
                  const auto xval1 = xvals1[g1_add];
                  const auto yval1 = yvals1[g1_add];

                  if (y1 == y2 * yinc && x1 == x2 * xinc)
                    {
                      xvals2_0 = xval1;
                      xvals2[g2_add] = xval1 / area_norm;
                      yvals2[g2_add] = yval1 / area_norm;
                    }
                  else if (std::fabs(xval1 - xvals2_0) > 270.)
                    {
                      if ((xval1 - xvals2_0) > 270.)
                        xvals2[g2_add] += (xval1 - 360.) / area_norm;
                      else if ((xval1 - xvals2_0) < -270.)
                        xvals2[g2_add] += (xval1 + 360.) / area_norm;
                      yvals2[g2_add] += yval1 / area_norm;
                    }
                  else
                    {
                      xvals2[g2_add] += xval1 / area_norm;
                      yvals2[g2_add] += yval1 / area_norm;
                    }

                  if (lGridHasBounds)
                    {
                      // upper left cell
                      if (y1 == y2 * yinc && x1 == x2 * xinc)
                        {
                          int corner = 0;
                          if (g1_add + nlon1 > gridsize1) cdoWarning("Can't find cell below upper left!");
                          else corner = find_corner(g1_add, 1, nlon1, grid1_corner_lon, grid1_corner_lat);
                          on_up = grid1_corner_lon[4 * g1_add + corner];
                          ax_le = grid1_corner_lat[4 * g1_add + corner];
                        }

                      // upper right cell
                      if ((y1 == y2 * yinc) && (x1 == (x2 + 1) * xinc - 1))
                        {
                          int corner = 0;
                          if (g1_add + nlon1 > gridsize1) cdoWarning("Can't find cell below upper right!");
                          else corner = find_corner(g1_add, -1, nlon1, grid1_corner_lon, grid1_corner_lat);
                          ox_up = grid1_corner_lon[4 * g1_add + corner];
                          ax_ri = grid1_corner_lat[4 * g1_add + corner];
                        }

                      // lower right cell
                      if ((y1 == (y2 + 1) * yinc - 1) && (x1 == (x2 + 1) * xinc - 1))
                        {
                          int corner = 0;
                          if (g1_add < nlon1) cdoWarning("Can't find cell above lower right!");
                          else corner = find_corner(g1_add, -1, -nlon1, grid1_corner_lon, grid1_corner_lat);
                          ox_lo = grid1_corner_lon[4 * g1_add + corner];
                          an_ri = grid1_corner_lat[4 * g1_add + corner];
                        }

                      // lower left cell
                      if ((y1 == (y2 + 1) * yinc - 1) && (x1 == x2 * xinc))
                        {
                          int corner = 0;
                          if (g1_add < nlon1) cdoWarning("Can't find cell above lower left!");
                          else corner = find_corner(g1_add, 1, -nlon1, grid1_corner_lon, grid1_corner_lat);
                          on_lo = grid1_corner_lon[4 * g1_add + corner];
                          an_le = grid1_corner_lat[4 * g1_add + corner];
                        }
                    }
                }     // for ( x1 = x2*xinc; x1 < xinc*(x2+1); x1++ )
            }         // for ( y1 = y2*yinc; y1 < yinc*(y2+1); y1++ )

          if (lGridHasBounds)
            {
              // upper left corner
              grid2_corner_lon[4 * g2_add + 3] = on_up;
              grid2_corner_lat[4 * g2_add + 3] = ax_le;
              // upper right corner
              grid2_corner_lon[4 * g2_add + 2] = ox_up;
              grid2_corner_lat[4 * g2_add + 2] = ax_ri;
              // lower right corner
              grid2_corner_lon[4 * g2_add + 1] = ox_lo;
              grid2_corner_lat[4 * g2_add + 1] = an_ri;
              // lower left corner
              grid2_corner_lon[4 * g2_add + 0] = on_lo;
              grid2_corner_lat[4 * g2_add + 0] = an_le;
            }

          //  while ( xvals2[g2_add] >  180. ) xvals2[g2_add] -= 360.;
          //  while ( xvals2[g2_add] < -180. ) xvals2[g2_add] += 360.;
        } // for ( x2 = 0; x2 < nlon2; x2++ )
    }     // for ( y2 = 0; y2 < nlat2; y2++ )

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

  if (lGridHasBounds)
    {
      gridDefNvertex(gridID2, 4);
      gridDefXbounds(gridID2, grid2_corner_lon.data());
      gridDefYbounds(gridID2, grid2_corner_lat.data());
    }
}

static int
genBoxGrid(int gridID1, size_t xinc, size_t yinc)
{
  if (xinc < 1 || yinc < 1) cdoAbort("xinc and yinc must not be smaller than 1!");

  int gridID2 = -1;
  const auto gridtype = gridInqType(gridID1);
  if (gridtype == GRID_GAUSSIAN || gridtype == GRID_LONLAT || gridtype == GRID_CURVILINEAR || gridtype == GRID_GENERIC)
    {
      const auto nlon1 = gridInqXsize(gridID1);
      const auto nlat1 = gridInqYsize(gridID1);
      if (xinc > nlon1 || yinc > nlat1) cdoAbort("xinc and/or yinc exceeds gridsize!");

      auto nlon2 = nlon1 / xinc;
      auto nlat2 = nlat1 / yinc;
      if (nlon1 % xinc) nlon2++;
      if (nlat1 % yinc) nlat2++;
      const auto gridsize2 = nlon2 * nlat2;

      gridID2 = gridCreate(gridtype, gridsize2);
      gridDefXsize(gridID2, nlon2);
      gridDefYsize(gridID2, nlat2);
    }
  else
    {
      cdoAbort("Unsupported grid: %s", gridNamePtr(gridtype));
    }

  if (gridtype == GRID_GAUSSIAN || gridtype == GRID_LONLAT)
    {
      genBoxGridReg2D(gridID1, xinc, yinc, gridID2);
    }
  else if (gridtype == GRID_GENERIC)
    {
    }
  else if (gridtype == GRID_CURVILINEAR)
    {
      genBoxGridCurv2D(gridID1, xinc, yinc, gridID2);
    }
  else
    {
      cdoAbort("Unsupported grid: %s", gridNamePtr(gridtype));
    }

  return gridID2;
}

static void
gridboxstat(const Field &field1, Field &field2, const size_t xinc, const size_t yinc, const int statfunc)
{
  const auto useWeight = !field1.weightv.empty();

  const auto gridsize = xinc * yinc;
  FieldVector field(Threading::ompNumThreads);
  for (int i = 0; i < Threading::ompNumThreads; i++)
    {
      field[i].resize(gridsize);
      if (useWeight) field[i].weightv.resize(gridsize);
      field[i].missval = field1.missval;
    }

  const auto gridID1 = field1.grid;
  const auto gridID2 = field2.grid;

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

  const auto nlon2 = gridInqXsize(gridID2);
  const auto nlat2 = gridInqYsize(gridID2);

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
  for (size_t ig = 0; ig < nlat2 * nlon2; ++ig)
    {
      const auto ompthID = cdo_omp_get_thread_num();

      const auto ilat = ig / nlon2;
      const auto ilon = ig - ilat * nlon2;

      size_t isize = 0;
      field[ompthID].nmiss = 0;
      for (size_t j = 0; j < yinc; ++j)
        {
          const auto jj = ilat * yinc + j;
          if (jj >= nlat1) break;
          for (size_t i = 0; i < xinc; ++i)
            {
              const auto ii = ilon * xinc + i;
              const auto index = jj * nlon1 + ii;
              if (ii >= nlon1) break;
              field[ompthID].vec_d[isize] = field1.vec_d[index];
              if (useWeight) field[ompthID].weightv[isize] = field1.weightv[index];
              if (DBL_IS_EQUAL(field[ompthID].vec_d[isize], field[ompthID].missval)) field[ompthID].nmiss++;
              isize++;
            }
        }

      field[ompthID].size = isize;
      field2.vec_d[ig] = vfldfun(field[ompthID], statfunc);
    }

  field2.nmiss = fieldNumMiss(field2);
}

static void
addOperators(void)
{
  // clang-format off
  cdoOperatorAdd("gridboxrange", func_range, 0, nullptr);
  cdoOperatorAdd("gridboxmin",   func_min,   0, nullptr);
  cdoOperatorAdd("gridboxmax",   func_max,   0, nullptr);
  cdoOperatorAdd("gridboxsum",   func_sum,   0, nullptr);
  cdoOperatorAdd("gridboxmean",  func_meanw, 1, nullptr);
  cdoOperatorAdd("gridboxavg",   func_avgw,  1, nullptr);
  cdoOperatorAdd("gridboxvar",   func_varw,  1, nullptr);
  cdoOperatorAdd("gridboxvar1",  func_var1w, 1, nullptr);
  cdoOperatorAdd("gridboxstd",   func_stdw,  1, nullptr);
  cdoOperatorAdd("gridboxstd1",  func_std1w, 1, nullptr);
  cdoOperatorAdd("gridboxskew",  func_skew,  0, nullptr);
  cdoOperatorAdd("gridboxkurt",  func_kurt,  0, nullptr);
  // clang-format on
}

void *
Gridboxstat(void *process)
{
  int lastgrid = -1;
  int nrecs;
  int varID, levelID;
  bool wstatus = false;
  char varname[CDI_MAX_NAME];

  cdoInitialize(process);

  operatorInputArg("xinc, yinc");
  operatorCheckArgc(2);
  const auto xinc = parameter2int(cdoOperatorArgv(0));
  const auto yinc = parameter2int(cdoOperatorArgv(1));

  addOperators();

  const auto operatorID = cdoOperatorID();
  const auto operfunc = cdoOperatorF1(operatorID);
  const auto needWeights = cdoOperatorF2(operatorID) != 0;

  const auto streamID1 = cdoOpenRead(0);

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

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

  const auto ngrids = vlistNgrids(vlistID1);
  if (ngrids > 1) cdoAbort("Too many different grids!");

  const auto gridID1 = vlistGrid(vlistID1, 0);
  if (gridInqType(gridID1) == GRID_GAUSSIAN_REDUCED)
    cdoAbort("Gaussian reduced grid found. Use option -R to convert it to a regular grid!");

  const auto gridID2 = genBoxGrid(gridID1, xinc, yinc);
  for (int index = 0; index < ngrids; index++) vlistChangeGridIndex(vlistID2, index, gridID2);

  const auto streamID2 = cdoOpenWrite(1);
  cdoDefVlist(streamID2, vlistID2);

  const auto gridsize1 = gridInqSize(gridID1);

  Field field1, field2;
  field1.resize(gridsize1);
  if (needWeights) field1.weightv.resize(gridsize1);

  const auto gridsize2 = gridInqSize(gridID2);
  field2.resize(gridsize2);

  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      taxisCopyTimestep(taxisID2, taxisID1);
      cdoDefTimestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; recID++)
        {
          cdoInqRecord(streamID1, &varID, &levelID);
          cdoReadRecord(streamID1, field1.vec_d.data(), &field1.nmiss);

          field1.grid = vlistInqVarGrid(vlistID1, varID);
          field1.size = gridInqSize(field1.grid);
          field1.missval = vlistInqVarMissval(vlistID1, varID);

          field2.grid = gridID2;
          field2.size = gridsize2;
          field2.missval = field1.missval;

          if (needWeights && field1.grid != lastgrid)
            {
              lastgrid = field1.grid;
              wstatus = gridWeights(field1.grid, field1.weightv.data());
            }
          if (wstatus != 0 && tsID == 0 && levelID == 0)
            {
              vlistInqVarName(vlistID1, varID, varname);
              cdoWarning("Grid cell bounds not available, using constant grid cell area weights for variable %s!", varname);
            }

          gridboxstat(field1, field2, xinc, yinc, operfunc);

          cdoDefRecord(streamID2, varID, levelID);
          cdoWriteRecord(streamID2, field2.vec_d.data(), field2.nmiss);
        }
      tsID++;
    }

  cdoStreamClose(streamID2);
  cdoStreamClose(streamID1);

  cdoFinish();

  return nullptr;
}
