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

  Author: Uwe Schulzweida

*/

#include <algorithm>

#include <cdi.h>

#include "cdo_options.h"
#include "cdo_rlimit.h"
#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "util_files.h"

struct GridInfo1
{
  int gridID;
  int gridtype;
  size_t gridsize;
  size_t nx, ny;
  bool isReg2d;
};

struct DistgridInfo
{
  int gridID;
  size_t gridsize;
  size_t nx, ny;
  size_t offset;
  std::vector<size_t> cellindex;
};

static void
gen_dist_grids(const GridInfo1 &gridInfo1, std::vector<DistgridInfo> &distgridInfo, size_t nxvals, size_t nyvals, size_t nxblocks,
               size_t nyblocks, size_t nsplit)
{
  const auto gridID1 = gridInfo1.gridID;
  const auto gridtype = gridInfo1.gridtype;
  const auto isReg2d = gridInfo1.isReg2d;
  const auto lunstructured = (gridtype == GRID_UNSTRUCTURED);
  const auto lcurvilinear = (gridtype == GRID_CURVILINEAR);
  const auto lregular
      = (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN || gridtype == GRID_GENERIC || gridtype == GRID_PROJECTION);
  if (!lregular && !lcurvilinear && !lunstructured) cdo_abort("Unsupported grid type: %s!", gridNamePtr(gridtype));

  const auto nx = gridInqXsize(gridID1);
  const size_t ny = lunstructured ? 1 : gridInqYsize(gridID1);

  const auto lxcoord = (gridInqXvals(gridID1, nullptr) > 0);
  const auto lycoord = (gridInqYvals(gridID1, nullptr) > 0);
  const auto withBounds = (!lregular && gridHasBounds(gridID1));

  std::vector<size_t> xlsize(nxblocks), ylsize(nyblocks);

  for (size_t ix = 0; ix < nxblocks; ++ix) xlsize[ix] = nxvals;
  if (nx % nxblocks != 0) xlsize[nxblocks - 1] = nx - (nxblocks - 1) * nxvals;
  if (Options::cdoVerbose)
    for (size_t ix = 0; ix < nxblocks; ++ix) cdo_print("xblock %zu: %zu", ix, xlsize[ix]);

  for (size_t iy = 0; iy < nyblocks; ++iy) ylsize[iy] = nyvals;
  if (ny % nyblocks != 0) ylsize[nyblocks - 1] = ny - (nyblocks - 1) * nyvals;
  if (Options::cdoVerbose)
    for (size_t iy = 0; iy < nyblocks; ++iy) cdo_print("yblock %zu: %zu", iy, ylsize[iy]);

  const auto nxvmax = std::max(nxvals, xlsize[nxblocks - 1]);
  const auto nyvmax = std::max(nyvals, ylsize[nyblocks - 1]);

  Varray<double> xpvals, ypvals;
  Varray<double> xvals, yvals, xbounds, ybounds;
  Varray<double> xvals2, yvals2, xbounds2, ybounds2;

  if (lxcoord)
    {
      xvals.resize(lregular ? nx : nx * ny);
      gridInqXvals(gridID1, xvals.data());

      if (!lregular) xvals2.resize(nxvmax * nyvmax);
    }

  if (lycoord)
    {
      yvals.resize(lregular ? ny : nx * ny);
      gridInqYvals(gridID1, yvals.data());

      if (!lregular) yvals2.resize(nxvmax * nyvmax);
    }

  size_t nv = 0;
  if (withBounds)
    {
      if (!lregular)
        {
          nv = gridInqNvertex(gridID1);
          xbounds.resize(nx * ny * nv);
          ybounds.resize(nx * ny * nv);
          xbounds2.resize(nxvmax * nyvmax * nv);
          ybounds2.resize(nxvmax * nyvmax * nv);
        }

      gridInqXbounds(gridID1, xbounds.data());
      gridInqYbounds(gridID1, ybounds.data());
    }

  size_t index = 0;
  for (size_t iy = 0; iy < nyblocks; ++iy)
    for (size_t ix = 0; ix < nxblocks; ++ix)
      {
        const auto offset = iy * nyvals * nx + ix * nxvals;
        auto gridsize2 = xlsize[ix] * ylsize[iy];
        if (!isReg2d) distgridInfo[index].cellindex.resize(gridsize2);
        auto &cellindex = distgridInfo[index].cellindex;

        distgridInfo[index].nx = xlsize[ix];
        distgridInfo[index].ny = ylsize[iy];
        distgridInfo[index].offset = offset;

        gridsize2 = 0;
        // printf("iy %d, ix %d offset %d\n", iy, ix,  offset);
        for (size_t j = 0; j < ylsize[iy]; ++j)
          for (size_t i = 0; i < xlsize[ix]; ++i)
            {
              // printf(">> %d %d %d\n", j, i, offset + j*nx + i);
              if (!lregular)
                {
                  if (lxcoord) xvals2[gridsize2] = xvals[offset + j * nx + i];
                  if (lycoord) yvals2[gridsize2] = yvals[offset + j * nx + i];
                  if (withBounds)
                    {
                      for (size_t k = 0; k < nv; ++k)
                        {
                          xbounds2[gridsize2 * nv + k] = xbounds[(offset + j * nx + i) * nv + k];
                          ybounds2[gridsize2 * nv + k] = ybounds[(offset + j * nx + i) * nv + k];
                        }
                    }
                }
              if (!isReg2d) cellindex[gridsize2] = offset + j * nx + i;
              gridsize2++;
            }
        // printf("gridsize2 %d\n", gridsize2);

        const auto gridID2 = gridCreate(gridtype, gridsize2);
        if (gridtype != GRID_UNSTRUCTURED)
          {
            gridDefXsize(gridID2, xlsize[ix]);
            gridDefYsize(gridID2, ylsize[iy]);

            gridDefNP(gridID2, gridInqNP(gridID1));
          }

        if (withBounds) gridDefNvertex(gridID2, nv);

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

        grid_copy_keys(gridID1, gridID2);

        if (gridtype == GRID_PROJECTION) grid_copy_mapping(gridID1, gridID2);

        if (lregular)
          {
            if (lxcoord) gridDefXvals(gridID2, &xvals[ix * nxvals]);
            if (lycoord) gridDefYvals(gridID2, &yvals[iy * nyvals]);
          }
        else
          {
            if (lxcoord) gridDefXvals(gridID2, xvals2.data());
            if (lycoord) gridDefYvals(gridID2, yvals2.data());
            if (withBounds)
              {
                gridDefXbounds(gridID2, xbounds2.data());
                gridDefYbounds(gridID2, ybounds2.data());
              }
          }

        const auto projID1 = gridInqProj(gridID1);
        if (projID1 != CDI_UNDEFID && gridInqType(projID1) == GRID_PROJECTION)
          {
            const auto projID2 = gridCreate(GRID_PROJECTION, gridsize2);
            gridDefXsize(projID2, xlsize[ix]);
            gridDefYsize(projID2, ylsize[iy]);

            grid_copy_keys(projID1, projID2);
            grid_copy_mapping(projID1, projID2);

            const auto lxpcoord = (gridInqXvals(projID1, nullptr) > 0);
            if (lxpcoord)
              {
                if (!xpvals.size())
                  {
                    xpvals.resize(nx);
                    gridInqXvals(projID1, xpvals.data());
                  }
                gridDefXvals(projID2, &xpvals[ix * nxvals]);
              }
            const auto lypcoord = (gridInqYvals(projID1, nullptr) > 0);
            if (lypcoord)
              {
                if (!ypvals.size())
                  {
                    ypvals.resize(ny);
                    gridInqYvals(projID1, ypvals.data());
                  }
                gridDefYvals(projID2, &ypvals[iy * nyvals]);
              }

            gridDefProj(gridID2, projID2);
          }

        distgridInfo[index].gridID = gridID2;
        distgridInfo[index].gridsize = gridsize2;

        index++;
        if (index > nsplit) cdo_abort("Internal problem, index exceeded bounds!");
      }
}

static void
dist_cells_reg2d(const Field &field1, Field &field2, const DistgridInfo &distgridInfo, size_t nlon)
{
  const auto nx = distgridInfo.nx;
  const auto ny = distgridInfo.ny;

  for (size_t j = 0; j < ny; ++j)
    {
      const auto offset1 = distgridInfo.offset + j * nlon;
      const auto offset2 = j * nx;

      if (field1.memType == MemType::Float)
        for (size_t i = 0; i < nx; ++i) field2.vec_f[offset2 + i] = field1.vec_f[offset1 + i];
      else
        for (size_t i = 0; i < nx; ++i) field2.vec_d[offset2 + i] = field1.vec_d[offset1 + i];
    }
}

static void
dist_cells(const Field &field1, Field &field2, const DistgridInfo &distgridInfo)
{
  const auto gridsize = distgridInfo.gridsize;
  const auto &cellidx = distgridInfo.cellindex;

  if (field1.memType == MemType::Float)
    for (size_t i = 0; i < gridsize; ++i) field2.vec_f[i] = field1.vec_f[cellidx[i]];
  else
    for (size_t i = 0; i < gridsize; ++i) field2.vec_d[i] = field1.vec_d[cellidx[i]];
}

void *
Distgrid(void *process)
{
  constexpr size_t MaxBlocks = 99999;
  int gridID1;
  int gridtype = -1;

  cdo_initialize(process);

  if (process_self().m_ID != 0) cdo_abort("This operator can't be combined with other operators!");

  operator_input_arg("nxblocks, [nyblocks]");
  if (cdo_operator_argc() < 1) cdo_abort("Too few arguments!");
  if (cdo_operator_argc() > 2) cdo_abort("Too many arguments!");
  size_t nxblocks = parameter_to_int(cdo_operator_argv(0));
  size_t nyblocks = 1;
  if (cdo_operator_argc() == 2) nyblocks = parameter_to_int(cdo_operator_argv(1));

  if (nxblocks == 0) cdo_abort("nxblocks has to be greater than 0!");
  if (nyblocks == 0) cdo_abort("nyblocks has to be greater than 0!");

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

  const auto ngrids = vlistNgrids(vlistID1);

  std::vector<GridInfo1> gridInfo1(ngrids);
  for (int index = 0; index < ngrids; ++index)
    {
      const auto gridID = vlistGrid(vlistID1, index);
      gridInfo1[index].gridID = gridID;
      gridInfo1[index].gridtype = gridInqType(gridID);
      gridInfo1[index].gridsize = gridInqSize(gridID);
      gridInfo1[index].nx = gridInqXsize(gridID);
      gridInfo1[index].ny = gridInqYsize(gridID);
      gridInfo1[index].isReg2d = (gridInfo1[index].gridtype != GRID_UNSTRUCTURED);
    }

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

    if (index == ngrids)
      cdo_abort("No Lon/Lat, Gaussian, curvilinear or generic grid found (%s data unsupported)!", gridNamePtr(gridtype));
  }

  const auto lunstruct = (gridtype == GRID_UNSTRUCTURED);

  gridID1 = vlistGrid(vlistID1, 0);
  const auto gridsize = gridInqSize(gridID1);
  const auto nx = gridInqXsize(gridID1);
  const size_t ny = lunstruct ? 1 : gridInqYsize(gridID1);
  for (int i = 1; i < ngrids; i++)
    {
      gridID1 = vlistGrid(vlistID1, i);
      if (gridsize != gridInqSize(gridID1)) cdo_abort("Gridsize must not change!");
    }

  if (nxblocks > nx)
    {
      cdo_print("nxblocks (%zu) greater than nx (%zu), set to %zu!", nxblocks, nx, nx);
      nxblocks = nx;
    }
  if (nyblocks > ny)
    {
      cdo_print("nyblocks (%zu) greater than ny (%zu), set to %zu!", nyblocks, ny, ny);
      nyblocks = ny;
    }

  auto xinc = nx / nxblocks;
  auto yinc = ny / nyblocks;
  if (nx % xinc && nx % (xinc + 1) && nxblocks * (xinc + 1) <= nx) xinc++;
  if (ny % yinc && ny % (yinc + 1) && nyblocks * (yinc + 1) <= ny) yinc++;

  const auto nsplit = nxblocks * nyblocks;
  if (nsplit > MaxBlocks) cdo_abort("Too many blocks (max = %d)!", MaxBlocks);

  cdo::set_numfiles(nsplit + 8);

  Field field1, field2;

  VarList varList1;
  varListInit(varList1, vlistID1);

  std::vector<int> vlistIDs(nsplit);
  std::vector<CdoStreamID> streamIDs(nsplit);

  std::vector<std::vector<DistgridInfo>> distgridInfo(ngrids);
  for (int i = 0; i < ngrids; i++) distgridInfo[i].resize(nsplit);

  for (size_t index = 0; index < nsplit; index++) vlistIDs[index] = vlistDuplicate(vlistID1);

  if (Options::cdoVerbose) cdo_print("ngrids=%d  nsplit=%zu", ngrids, nsplit);

  for (int i = 0; i < ngrids; i++)
    {
      gen_dist_grids(gridInfo1[i], distgridInfo[i], xinc, yinc, nxblocks, nyblocks, nsplit);
      /*
      if ( Options::cdoVerbose )
        for ( size_t index = 0; index < nsplit; index++ )
          cdo_print("Block %d,  gridID %d,  gridsize %zu", index+1,
      distgridInfo[i].gridID[index], gridInqSize(grids[i].gridID[index]));
      */
      for (size_t index = 0; index < nsplit; index++) vlistChangeGridIndex(vlistIDs[index], i, distgridInfo[i][index].gridID);
    }

  char filename[8192];
  strcpy(filename, cdo_get_obase().c_str());
  const int nchars = strlen(filename);

  auto refname = cdo_get_stream_name(0);
  char filesuffix[32] = { 0 };
  FileUtils::gen_suffix(filesuffix, sizeof(filesuffix), cdo_inq_filetype(streamID1), vlistID1, refname);

  for (size_t index = 0; index < nsplit; index++)
    {
      sprintf(filename + nchars, "%05ld", (long) index);
      if (filesuffix[0]) sprintf(filename + nchars + 5, "%s", filesuffix);

      streamIDs[index] = cdo_open_write(filename);
      cdo_def_vlist(streamIDs[index], vlistIDs[index]);
    }

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

      for (auto &streamID : streamIDs) cdo_def_timestep(streamID, tsID);

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

          for (size_t index = 0; index < nsplit; index++)
            {
              auto var = varList1[varID];
              var.gridID = distgridInfo[0][index].gridID;
              var.gridsize = distgridInfo[0][index].gridsize;
              field2.init(var);

              if (gridInfo1[0].isReg2d)
                dist_cells_reg2d(field1, field2, distgridInfo[0][index], gridInfo1[0].nx);
              else
                dist_cells(field1, field2, distgridInfo[0][index]);

              if (field1.nmiss) field_num_mv(field2);

              cdo_def_record(streamIDs[index], varID, levelID);
              cdo_write_record(streamIDs[index], field2);
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID1);

  for (auto &streamID : streamIDs) cdo_stream_close(streamID);
  for (auto &vlistID : vlistIDs) vlistDestroy(vlistID);

  for (int i = 0; i < ngrids; i++)
    for (size_t index = 0; index < nsplit; index++) gridDestroy(distgridInfo[i][index].gridID);

  cdo_finish();

  return nullptr;
}
