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

#include "dmemory.h"
#include "functs.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "util_files.h"
#include "cdo_options.h"
#include "cdi_lockedIO.h"

static int globalGridType = CDI_UNDEFID;

struct EnsfileType
{
  CdoStreamID streamID;
  int vlistID;
  size_t nmiss;
  size_t gridsize;
  std::vector<long> gridindex;
  Varray<double> array;
};

struct xyinfoType
{
  double x, y;
  int id;
};

static int
cmpx(const void *a, const void *b)
{
  const auto x = ((const xyinfoType *) a)->x;
  const auto y = ((const xyinfoType *) b)->x;
  return ((x > y) - (x < y)) * 2 + (x > y) - (x < y);
}

static int
cmpxy_lt(const void *a, const void *b)
{
  const auto x1 = ((const xyinfoType *) a)->x;
  const auto x2 = ((const xyinfoType *) b)->x;
  const auto y1 = ((const xyinfoType *) a)->y;
  const auto y2 = ((const xyinfoType *) b)->y;

  int cmp = 0;
  if (y1 < y2 || (std::fabs(y1 - y2) <= 0 && x1 < x2))
    cmp = -1;
  else if (y1 > y2 || (std::fabs(y1 - y2) <= 0 && x1 > x2))
    cmp = 1;

  return cmp;
}

static int
cmpxy_gt(const void *a, const void *b)
{
  const auto x1 = ((const xyinfoType *) a)->x;
  const auto x2 = ((const xyinfoType *) b)->x;
  const auto y1 = ((const xyinfoType *) a)->y;
  const auto y2 = ((const xyinfoType *) b)->y;

  int cmp = 0;
  if (y1 > y2 || (std::fabs(y1 - y2) <= 0 && x1 < x2))
    cmp = -1;
  else if (y1 < y2 || (std::fabs(y1 - y2) <= 0 && x1 > x2))
    cmp = 1;

  return cmp;
}

static int
genGrid(int ngrids, int nfiles, std::vector<EnsfileType> &ef, bool ginit, int igrid, long nxblocks)
{
  bool lsouthnorth = true;
  bool lregular = false;
  bool lcurvilinear = false;
  bool lunstructured = false;
  int gridID2 = -1;

  long nx = -1;
  if (nxblocks != -1) nx = nxblocks;

  auto gridID = vlistGrid(ef[0].vlistID, igrid);
  auto gridtype0 = (globalGridType != CDI_UNDEFID) ? globalGridType : gridInqType(gridID);
  if (ngrids > 1 && gridtype0 == GRID_GENERIC && gridInqXsize(gridID) == 0 && gridInqYsize(gridID) == 0) return gridID2;
  if (gridtype0 == GRID_UNSTRUCTURED) lunstructured = true;
  const int nv = lunstructured ? gridInqNvertex(gridID) : 0;
  const bool lcenter = globalGridType == CDI_UNDEFID && gridInqXvals(gridID, nullptr) > 0 && gridInqYvals(gridID, nullptr) > 0;
  const bool lbounds = lunstructured && globalGridType == CDI_UNDEFID && gridInqXbounds(gridID, nullptr) > 0
                       && gridInqYbounds(gridID, nullptr) > 0;

  std::vector<long> xsize(nfiles), ysize(nfiles);
  xyinfoType *xyinfo = (xyinfoType *) Malloc(nfiles * sizeof(xyinfoType));
  Varray2D<double> xvals(nfiles), yvals(nfiles);
  Varray2D<double> xbounds(nfiles), ybounds(nfiles);

  for (int fileID = 0; fileID < nfiles; fileID++)
    {
      gridID = vlistGrid(ef[fileID].vlistID, igrid);
      const auto gridtype = (globalGridType != CDI_UNDEFID) ? globalGridType : gridInqType(gridID);
      if (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN)
        lregular = true;
      else if (gridtype == GRID_CURVILINEAR)
        lcurvilinear = true;
      else if (gridtype == GRID_UNSTRUCTURED)
        lunstructured = true;
      else if (gridtype == GRID_GENERIC /*&& gridInqXsize(gridID) > 0 && gridInqYsize(gridID) > 0*/)
        ;
      else
        cdoAbort("Unsupported grid type: %s!", gridNamePtr(gridtype));

      xsize[fileID] = lunstructured ? gridInqSize(gridID) : gridInqXsize(gridID);
      ysize[fileID] = lunstructured ? 1 : gridInqYsize(gridID);
      if (xsize[fileID] == 0) xsize[fileID] = 1;
      if (ysize[fileID] == 0) ysize[fileID] = 1;

      if (lregular)
        {
          xvals[fileID].resize(xsize[fileID]);
          yvals[fileID].resize(ysize[fileID]);
        }
      else if (lcurvilinear || lunstructured)
        {
          if (lcenter)
            {
              xvals[fileID].resize(xsize[fileID] * ysize[fileID]);
              yvals[fileID].resize(xsize[fileID] * ysize[fileID]);
            }
          if (lbounds)
            {
              xbounds[fileID].resize(nv * xsize[fileID] * ysize[fileID]);
              ybounds[fileID].resize(nv * xsize[fileID] * ysize[fileID]);
            }
        }

      if (lregular || lcurvilinear || lunstructured)
        {
          if (lcenter)
            {
              gridInqXvals(gridID, xvals[fileID].data());
              gridInqYvals(gridID, yvals[fileID].data());
            }
          if (lbounds)
            {
              gridInqXbounds(gridID, xbounds[fileID].data());
              gridInqYbounds(gridID, ybounds[fileID].data());
            }
        }
      // printf("fileID %d, gridID %d\n", fileID, gridID);

      if (lregular)
        {
          xyinfo[fileID].x = xvals[fileID][0];
          xyinfo[fileID].y = yvals[fileID][0];
          xyinfo[fileID].id = fileID;

          if (ysize[fileID] > 1)
            {
              if (yvals[fileID][0] > yvals[fileID][ysize[fileID] - 1]) lsouthnorth = false;
            }
        }
      else
        {
          xyinfo[fileID].x = 0;
          xyinfo[fileID].y = 0;
          xyinfo[fileID].id = fileID;
        }
    }

  if (Options::cdoVerbose && lregular)
    for (int fileID = 0; fileID < nfiles; fileID++) printf("1 %d %g %g \n", xyinfo[fileID].id, xyinfo[fileID].x, xyinfo[fileID].y);

  if (lregular)
    {
      std::qsort(xyinfo, nfiles, sizeof(xyinfoType), cmpx);

      if (Options::cdoVerbose)
        for (int fileID = 0; fileID < nfiles; fileID++)
          printf("2 %d %g %g \n", xyinfo[fileID].id, xyinfo[fileID].x, xyinfo[fileID].y);

      if (lsouthnorth)
        std::qsort(xyinfo, nfiles, sizeof(xyinfoType), cmpxy_lt);
      else
        std::qsort(xyinfo, nfiles, sizeof(xyinfoType), cmpxy_gt);

      if (Options::cdoVerbose)
        for (int fileID = 0; fileID < nfiles; fileID++)
          printf("3 %d %g %g \n", xyinfo[fileID].id, xyinfo[fileID].x, xyinfo[fileID].y);

      if (nx <= 0)
        {
          nx = 1;
          for (int fileID = 1; fileID < nfiles; fileID++)
            {
              if (DBL_IS_EQUAL(xyinfo[0].y, xyinfo[fileID].y))
                nx++;
              else
                break;
            }
        }
    }
  else
    {
      if (nx <= 0) nx = nfiles;
    }

  const long ny = nfiles / nx;
  if (nx * ny != nfiles) cdoAbort("Number of input files (%ld) and number of blocks (%ldx%ld) differ!", nfiles, nx, ny);

  long xsize2 = 0;
  for (long i = 0; i < nx; ++i) xsize2 += xsize[xyinfo[i].id];
  long ysize2 = 0;
  for (long j = 0; j < ny; ++j) ysize2 += ysize[xyinfo[j * nx].id];
  if (Options::cdoVerbose) cdoPrint("xsize2 %ld  ysize2 %ld", xsize2, ysize2);

  {  // verify size of data
    const long xs = xsize[xyinfo[0].id];
    for (long j = 1; j < ny; ++j)
      if (xsize[xyinfo[j * nx].id] != xs) cdoAbort("xsize=%ld differ from first file (xsize=%ld)!", xsize[xyinfo[j * nx].id], xs);
    const long ys = ysize[xyinfo[0].id];
    for (long i = 1; i < nx; ++i)
      if (ysize[xyinfo[i].id] != ys) cdoAbort("ysize=%ld differ from first file (ysize=%ld)!", ysize[xyinfo[i].id], ys);
  }

  Varray<double> xvals2, yvals2;
  Varray<double> xbounds2, ybounds2;
  if (lregular)
    {
      xvals2.resize(xsize2);
      yvals2.resize(ysize2);
    }
  else if (lcurvilinear || lunstructured)
    {
      if (lcenter)
        {
          xvals2.resize(xsize2 * ysize2);
          yvals2.resize(xsize2 * ysize2);
        }
      if (lbounds)
        {
          xbounds2.resize(nv * xsize2 * ysize2);
          ybounds2.resize(nv * xsize2 * ysize2);
        }
    }

  std::vector<long> xoff(nx + 1), yoff(ny + 1);

  xoff[0] = 0;
  for (long i = 0; i < nx; ++i)
    {
      const long idx = xyinfo[i].id;
      if (lregular) arrayCopy(xsize[idx], xvals[idx].data(), &xvals2[xoff[i]]);
      xoff[i + 1] = xoff[i] + xsize[idx];
    }

  yoff[0] = 0;
  for (long j = 0; j < ny; ++j)
    {
      const long idx = xyinfo[j * nx].id;
      if (lregular) arrayCopy(ysize[idx], yvals[idx].data(), &yvals2[yoff[j]]);
      yoff[j + 1] = yoff[j] + ysize[idx];
    }

  if (!ginit)
    {
      for (int fileID = 0; fileID < nfiles; fileID++)
        {
          const long idx = xyinfo[fileID].id;
          const long iy = fileID / nx;
          const long ix = fileID - iy * nx;

          const long offset = yoff[iy] * xsize2 + xoff[ix];
          // printf("fileID %d %d, iy %d, ix %d, offset %d\n", fileID, xyinfo[fileID].id, iy, ix, offset);

          long ij = 0;
          for (long j = 0; j < ysize[idx]; ++j)
            for (long i = 0; i < xsize[idx]; ++i)
              {
                if (lcurvilinear || lunstructured)
                  {
                    if (lcenter)
                      {
                        xvals2[offset + j * xsize2 + i] = xvals[idx][ij];
                        yvals2[offset + j * xsize2 + i] = yvals[idx][ij];
                      }
                    if (lbounds)
                      {
                        for (long k = 0; k < nv; ++k)
                          {
                            xbounds2[(offset + j * xsize2 + i) * nv + k] = xbounds[idx][ij * nv + k];
                            ybounds2[(offset + j * xsize2 + i) * nv + k] = ybounds[idx][ij * nv + k];
                          }
                      }
                  }
                ef[idx].gridindex[ij++] = offset + j * xsize2 + i;
              }
        }
    }

  gridID2 = gridCreate(gridtype0, xsize2 * ysize2);
  if (!lunstructured)
    {
      gridDefXsize(gridID2, xsize2);
      gridDefYsize(gridID2, ysize2);
    }
  else if (nv > 0)
    {
      gridDefNvertex(gridID2, nv);
    }

  if (lregular || lcurvilinear || lunstructured)
    {
      if (lcenter)
        {
          gridDefXvals(gridID2, xvals2.data());
          gridDefYvals(gridID2, yvals2.data());
        }
      if (lbounds)
        {
          gridDefXbounds(gridID2, xbounds2.data());
          gridDefYbounds(gridID2, ybounds2.data());
        }
    }

  Free(xyinfo);

  gridID = vlistGrid(ef[0].vlistID, igrid);

  grid_copy_attributes(gridID, gridID2);

  if (gridtype0 == GRID_PROJECTION) grid_copy_mapping(gridID, gridID2);

  return gridID2;
}

void *
Collgrid(void *process)
{
  int nxblocks = -1;
  int varID, levelID;
  int nrecs0;

  cdoInitialize(process);

  const auto nfiles = cdoStreamCnt() - 1;
  const auto ofilename = cdoGetStreamName(nfiles);

  if (!Options::cdoOverwriteMode && fileExists(ofilename) && !userFileOverwrite(ofilename))
    cdoAbort("Outputfile %s already exists!", ofilename);

  std::vector<EnsfileType> ef(nfiles);

  for (int fileID = 0; fileID < nfiles; fileID++)
    {
      ef[fileID].streamID = cdoOpenRead(fileID);
      ef[fileID].vlistID = cdoStreamInqVlist(ef[fileID].streamID);
    }

  const auto vlistID1 = ef[0].vlistID;
  vlistClearFlag(vlistID1);

  /* check that the contents is always the same */
  for (int fileID = 1; fileID < nfiles; fileID++) vlistCompare(vlistID1, ef[fileID].vlistID, CMP_NAME | CMP_NLEVEL);

  const auto nvars = vlistNvars(vlistID1);
  std::vector<bool> vars(nvars);
  for (varID = 0; varID < nvars; varID++) vars[varID] = false;
  std::vector<bool> vars1(nvars);
  for (varID = 0; varID < nvars; varID++) vars1[varID] = false;

  auto nsel = operatorArgc();
  int noff = 0;

  if (nsel > 0)
    {
      const char *argument = cdoOperatorArgv(0).c_str();
      if (strcmp(argument, "gridtype=unstructured") == 0)
        {
          nsel--;
          noff++;
          globalGridType = GRID_UNSTRUCTURED;
        }
      else
        {
          int len = (int) strlen(argument);
          while (--len >= 0 && isdigit(argument[len]))
            ;

          if (len == -1)
            {
              nsel--;
              noff++;
              nxblocks = parameter2int(argument);
            }
        }
    }

  if (nsel == 0)
    {
      for (varID = 0; varID < nvars; varID++) vars1[varID] = true;
    }
  else
    {
      if (Options::cdoVerbose)
        for (int i = 0; i < nsel; i++) cdoPrint("name %d = %s\n", i + 1, cdoOperatorArgv(noff + i).c_str());

      std::vector<bool> selfound(nsel);
      for (int i = 0; i < nsel; i++) selfound[i] = false;

      char varname[CDI_MAX_NAME];
      for (varID = 0; varID < nvars; varID++)
        {
          vlistInqVarName(vlistID1, varID, varname);

          for (int isel = 0; isel < nsel; isel++)
            {
              if (cdoOperatorArgv(noff + isel).compare(varname) == 0)
                {
                  selfound[isel] = true;
                  vars1[varID] = true;
                }
            }
        }

      int err = 0;
      for (int isel = 0; isel < nsel; isel++)
        {
          if (selfound[isel] == false)
            {
              err++;
              cdoWarning("Variable name %s not found!", cdoOperatorArgv(noff + isel).c_str());
            }
        }
      if (err) cdoAbort("Could not find all requested variables: (%d/%d)", nsel - err, nsel);
    }

  for (varID = 0; varID < nvars; varID++)
    {
      if (vars1[varID])
        {
          const auto zaxisID = vlistInqVarZaxis(vlistID1, varID);
          const auto nlevs = zaxisInqSize(zaxisID);
          for (int levID = 0; levID < nlevs; levID++) vlistDefFlag(vlistID1, varID, levID, true);
        }
    }

  for (int fileID = 0; fileID < nfiles; fileID++)
    {
      const auto gridsize = vlistGridsizeMax(ef[fileID].vlistID);
      ef[fileID].gridsize = gridsize;
      ef[fileID].gridindex.resize(gridsize);
      ef[fileID].array.resize(gridsize);
    }

  const auto vlistID2 = vlistCreate();
  cdoVlistCopyFlag(vlistID2, vlistID1);
  /*
  if ( Options::cdoVerbose )
    {
      vlistPrint(vlistID1);
      vlistPrint(vlistID2);
    }
  */
  // auto vlistID2 = vlistDuplicate(vlistID1);
  const auto nvars2 = vlistNvars(vlistID2);

  const auto ngrids1 = vlistNgrids(vlistID1);
  const auto ngrids2 = vlistNgrids(vlistID2);

  std::vector<int> gridIDs(ngrids2);

  bool ginit = false;
  for (int i2 = 0; i2 < ngrids2; ++i2)
    {
      int i1;
      for (i1 = 0; i1 < ngrids1; ++i1)
        if (vlistGrid(vlistID1, i1) == vlistGrid(vlistID2, i2)) break;

      //   printf("i1 %d i2 %d\n", i1, i2);

      if (!ginit)
        {
          gridIDs[i2] = genGrid(ngrids2, nfiles, ef, ginit, i1, nxblocks);
          if (gridIDs[i2] != -1) ginit = true;
        }
      else
        gridIDs[i2] = genGrid(ngrids2, nfiles, ef, ginit, i1, nxblocks);
    }

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

  size_t gridsize2 = 0;
  for (int i = 0; i < ngrids2; ++i)
    {
      if (gridIDs[i] != -1)
        {
          if (gridsize2 == 0) gridsize2 = gridInqSize(gridIDs[i]);
          if (gridsize2 != gridInqSize(gridIDs[i]))
            cdoAbort("Target gridsize differ (gridID=%d  gridsize=%zu)!", i + 1, gridInqSize(gridIDs[i]));
          vlistChangeGridIndex(vlistID2, i, gridIDs[i]);
        }
    }

  for (varID = 0; varID < nvars2; varID++)
    {
      const auto gridID = vlistInqVarGrid(vlistID2, varID);

      for (int i = 0; i < ngrids2; ++i)
        {
          if (gridIDs[i] != -1)
            {
              if (gridID == vlistGrid(vlistID2, i) || gridsize2 == gridInqSize(gridID)) vars[varID] = true;
              break;
            }
        }
    }

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

  Varray<double> array2;
  if (gridsize2 > 0) array2.resize(gridsize2);

  int tsID = 0;
  do
    {
      nrecs0 = cdoStreamInqTimestep(ef[0].streamID, tsID);
      for (int fileID = 1; fileID < nfiles; fileID++)
        {
          const auto nrecs = cdoStreamInqTimestep(ef[fileID].streamID, tsID);
          if (nrecs != nrecs0)
            cdoAbort("Number of records at time step %d of %s and %s differ!", tsID + 1, cdoGetStreamName(0),
                     cdoGetStreamName(fileID));
        }

      taxisCopyTimestep(taxisID2, taxisID1);

      if (nrecs0 > 0) cdoDefTimestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs0; recID++)
        {
          cdoInqRecord(ef[0].streamID, &varID, &levelID);
          if (Options::cdoVerbose && tsID == 0) printf(" tsID, recID, varID, levelID %d %d %d %d\n", tsID, recID, varID, levelID);

          for (int fileID = 1; fileID < nfiles; fileID++)
            {
              int varIDx, levelIDx;
              cdoInqRecord(ef[fileID].streamID, &varIDx, &levelIDx);
            }

          if (vlistInqFlag(vlistID1, varID, levelID) == true)
            {
              const auto varID2 = vlistFindVar(vlistID2, varID);
              const auto levelID2 = vlistFindLevel(vlistID2, varID, levelID);
              if (Options::cdoVerbose && tsID == 0) printf("varID %d %d levelID %d %d\n", varID, varID2, levelID, levelID2);

              const auto missval = vlistInqVarMissval(vlistID2, varID2);
              for (size_t i = 0; i < gridsize2; i++) array2[i] = missval;

#ifdef _OPENMP
#pragma omp parallel for default(shared)
#endif
              for (int fileID = 0; fileID < nfiles; fileID++)
                {
                  cdoReadRecord(ef[fileID].streamID, &ef[fileID].array[0], &ef[fileID].nmiss);

                  if (vars[varID2])
                    {
                      const auto gridsize = ef[fileID].gridsize;
                      for (size_t i = 0; i < gridsize; ++i) array2[ef[fileID].gridindex[i]] = ef[fileID].array[i];
                    }
                }

              cdoDefRecord(streamID2, varID2, levelID2);

              if (vars[varID2])
                {
                  const auto nmiss = varrayNumMV(gridsize2, array2, missval);
                  cdoWriteRecord(streamID2, array2.data(), nmiss);
                }
              else
                cdoWriteRecord(streamID2, &ef[0].array[0], ef[0].nmiss);
            }
        }

      tsID++;
    }
  while (nrecs0 > 0);

  for (int fileID = 0; fileID < nfiles; fileID++) cdoStreamClose(ef[fileID].streamID);

  cdoStreamClose(streamID2);

  cdoFinish();

  return nullptr;
}