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

#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"


static void
genGrids(int gridID1, int *gridIDs, size_t nxvals, size_t nyvals, size_t nxblocks, size_t nyblocks,
         std::vector<std::vector<size_t>> &gridindex, size_t *ogridsize, size_t nsplit)
{
  Varray<double> xpvals, ypvals;
  auto gridtype = gridInqType(gridID1);
  bool lunstructured = gridtype == GRID_UNSTRUCTURED;
  bool lregular = true;
  if (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN || gridtype == GRID_GENERIC)
    lregular = true;
  else if (gridtype == GRID_CURVILINEAR)
    lregular = false;
  else if (lunstructured)
    lregular = false;
  else
    cdoAbort("Unsupported grid type: %s!", gridNamePtr(gridtype));

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

  const bool lxcoord = gridInqXvals(gridID1, nullptr) > 0;
  const bool lycoord = gridInqYvals(gridID1, nullptr) > 0;
  const bool lbounds = !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) cdoPrint("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) cdoPrint("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> 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 (lbounds)
    {
      if (lregular)
        {
        }
      else
        {
          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];
        gridindex[index].resize(gridsize2);

        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 (lbounds)
                    {
                      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];
                        }
                    }
                }
              gridindex[index][gridsize2++] = offset + j * nx + i;
            }
        // 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 (lbounds) gridDefNvertex(gridID2, nv);

        gridDefDatatype(gridID2, gridInqDatatype(gridID1));

        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 (lbounds)
              {
                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 bool lxpcoord = gridInqXvals(projID1, nullptr) > 0;
            if (lxpcoord)
              {
                if (!xpvals.size())
                  {
                    xpvals.resize(nx);
                    gridInqXvals(projID1, xpvals.data());
                  }
                gridDefXvals(projID2, &xpvals[ix * nxvals]);
              }
            const bool 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);
          }

        gridIDs[index] = gridID2;
        ogridsize[index] = gridsize2;

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

static void
window_cell(const Field &field1, Field &field2, size_t gridsize2, size_t *cellidx)
{
  if (field1.memType == MemType::Float)
    for (size_t i = 0; i < gridsize2; ++i) field2.vec_f[i] = field1.vec_f[cellidx[i]];
  else
    for (size_t i = 0; i < gridsize2; ++i) field2.vec_d[i] = field1.vec_d[cellidx[i]];
}

struct GridInfo
{
  std::vector<int> gridIDs;
  std::vector<size_t> gridsize;
  std::vector<std::vector<size_t>> gridindex;
};

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

  cdoInitialize(process);

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

  operatorInputArg("nxblocks, [nyblocks]");
  if (operatorArgc() < 1) cdoAbort("Too few arguments!");
  if (operatorArgc() > 2) cdoAbort("Too many arguments!");
  size_t nxblocks = parameter2int(cdoOperatorArgv(0));
  size_t nyblocks = 1;
  if (operatorArgc() == 2) nyblocks = parameter2int(cdoOperatorArgv(1));

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

  const auto streamID1 = cdoOpenRead(0);
  const auto vlistID1 = cdoStreamInqVlist(streamID1);

  const auto ngrids = vlistNgrids(vlistID1);

  {
    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_GENERIC && gridInqXsize(gridID1) > 0 && gridInqYsize(gridID1) > 0))
          break;
      }

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

  const bool 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)) cdoAbort("Gridsize must not change!");
    }

  if (nxblocks > nx)
    {
      cdoPrint("nxblocks (%zu) greater than nx (%zu), set to %zu!", nxblocks, nx, nx);
      nxblocks = nx;
    }
  if (nyblocks > ny)
    {
      cdoPrint("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) cdoAbort("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<GridInfo> gridinfo(ngrids);
  for (int i = 0; i < ngrids; i++)
    {
      gridinfo[i].gridIDs.resize(nsplit);
      gridinfo[i].gridsize.resize(nsplit);
      gridinfo[i].gridindex.resize(nsplit);
    }

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

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

  for (int i = 0; i < ngrids; i++)
    {
      gridID1 = vlistGrid(vlistID1, i);
      genGrids(gridID1, gridinfo[i].gridIDs.data(), xinc, yinc, nxblocks, nyblocks, gridinfo[i].gridindex,
               gridinfo[i].gridsize.data(), nsplit);
      /*
      if ( Options::cdoVerbose )
        for ( size_t index = 0; index < nsplit; index++ )
          cdoPrint("Block %d,  gridID %d,  gridsize %zu", index+1,
      gridinfo[i].gridIDs[index], gridInqSize(grids[i].gridIDs[index]));
      */
      for (size_t index = 0; index < nsplit; index++) vlistChangeGridIndex(vlistIDs[index], i, gridinfo[i].gridIDs[index]);
    }

  char filename[8192];
  strcpy(filename, cdoGetObase());
  const int nchars = strlen(filename);

  auto refname = cdoGetStreamName(0);
  char filesuffix[32] = { 0 };
  cdoGenFileSuffix(filesuffix, sizeof(filesuffix), cdoInqFiletype(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] = cdoOpenWrite(filename);
      cdoDefVlist(streamIDs[index], vlistIDs[index]);
    }

  if (ngrids > 1) cdoPrint("Baustelle: number of different grids > 1!");
  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      for (auto &streamID : streamIDs) cdoDefTimestep(streamID, tsID);

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

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

              window_cell(field1, field2, gridinfo[0].gridsize[index], gridinfo[0].gridindex[index].data());
              if (field1.nmiss) fieldNumMV(field2);

              cdoDefRecord(streamIDs[index], varID, levelID);
              cdoWriteRecord(streamIDs[index], field2);
            }
        }

      tsID++;
    }

  cdoStreamClose(streamID1);

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

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

  cdoFinish();

  return nullptr;
}
