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

      Setgrid    setgrid         Set grid
      Setgrid    setgridtype     Set grid type
      Setgrid    setgridarea     Set grid area
      Setgrid    setgridmask     Set grid mask
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "cdi_lockedIO.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "griddes.h"
#include "gridreference.h"
#include "util_files.h"

static void
changeGrid(const int vlistID1, const int vlistID2, const int gridID2)
{
  int found = 0;
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      if (gridInqSize(gridID1) == gridInqSize(gridID2))
        {
          vlistChangeGridIndex(vlistID2, index, gridID2);
          found++;
        }
    }
  if (!found) cdoWarning("No horizontal grid with %zu cells found!", gridInqSize(gridID2));
}

static void
setGridtype(const int vlistID1, const int vlistID2, const int gridtype, const std::string &gridname, const bool lregular,
            const bool lregularnn, const bool ldereference, const bool lbounds, std::vector<int> &grid2_vgpm)
{
  int gridID2;
  bool lrgrid = false;
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      const auto gridtype1 = gridInqType(gridID1);
      gridID2 = -1;

      if (gridtype1 == GRID_GENERIC && gridInqSize(gridID1) == 1) continue;

      if (lregular || lregularnn)
        {
          if (gridtype1 == GRID_GAUSSIAN_REDUCED) gridID2 = gridToRegular(gridID1);
        }
      else if (ldereference)
        {
          const auto reference = dereferenceGrid(gridID1);
          if (reference.isValid) gridID2 = reference.gridID;
          if (reference.notFound) cdoAbort("Reference to horizontal grid not found!");
        }
      else
        {
          if (gridtype == GRID_CURVILINEAR)
            {
              gridID2 = (gridtype1 == GRID_CURVILINEAR) ? gridID1 : gridToCurvilinear(gridID1, lbounds);
            }
          else if (gridtype == GRID_UNSTRUCTURED)
            {
              gridID2 = gridToUnstructured(gridID1, 1);
              if (gridtype1 == GRID_GME)
                {
                  const auto grid2_nvgp = gridInqSize(gridID2);
                  grid2_vgpm.resize(grid2_nvgp);
                  gridInqMaskGME(gridID2, grid2_vgpm.data());
                  gridCompress(gridID2);
                }
            }
          else if (gridtype == GRID_LONLAT && gridtype1 == GRID_CURVILINEAR)
            {
              gridID2 = gridCurvilinearToRegular(gridID1);
              if (gridID2 == -1) cdoWarning("Conversion of curvilinear grid to regular grid failed!");
            }
          else if (gridtype == GRID_LONLAT && gridtype1 == GRID_UNSTRUCTURED)
            {
              gridID2 = -1;
              if (gridID2 == -1) cdoWarning("Conversion of unstructured grid to regular grid failed!");
            }
          else if (gridtype == GRID_LONLAT && gridtype1 == GRID_GENERIC)
            {
              gridID2 = -1;
              if (gridID2 == -1) cdoWarning("Conversion of generic grid to regular grid failed!");
            }
          else if (gridtype == GRID_LONLAT && gridtype1 == GRID_LONLAT)
            {
              gridID2 = gridID1;
            }
          else if (gridtype == GRID_PROJECTION)
            {
              if (gridInqType(gridID1) == GRID_PROJECTION)
                gridID2 = gridID1;
              else
                {
                  const int projID = gridInqProj(gridID1);
                  if (projID != CDI_UNDEFID && gridInqType(projID) == GRID_PROJECTION) gridID2 = projID;
                }
              if (gridID2 == -1) cdoWarning("Projection not found!");
            }
          else
            cdoAbort("Unsupported grid name: %s", gridname.c_str());
        }

      if (gridID2 == -1)
        {
          if (!(lregular || lregularnn)) cdoAbort("Unsupported grid type!");
        }

      if (gridID2 != -1)
        {
          if (lregular || lregularnn) lrgrid = true;
          vlistChangeGridIndex(vlistID2, index, gridID2);
        }
    }

  if ((lregular || lregularnn) && !lrgrid) cdoWarning("No reduced Gaussian grid found!");
}

static void
setGridcellArea(const int vlistID1, const int vlistID2, Varray<double> &gridcellArea)
{
  const auto areasize = gridcellArea.size();
  int numGridsChanges = 0;
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      const auto gridsize = gridInqSize(gridID1);
      if (gridsize == areasize)
        {
          const auto gridID2 = gridDuplicate(gridID1);
          gridDefArea(gridID2, gridcellArea.data());
          vlistChangeGridIndex(vlistID2, index, gridID2);
          numGridsChanges++;
        }
    }
  if (!numGridsChanges) cdoWarning("No grid with %zu cells found!", areasize);
}

static void
setGridMask(const int vlistID1, const int vlistID2, Varray<double> &gridmask)
{
  const auto masksize = gridmask.size();
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      const auto gridsize = gridInqSize(gridID1);
      if (gridsize == masksize)
        {
          std::vector<int> mask(masksize);
          for (size_t i = 0; i < masksize; i++)
            {
              mask[i] = (gridmask[i] < 0 || gridmask[i] > 255) ? 0 : (int) std::lround(gridmask[i]);
            }
          const auto gridID2 = gridDuplicate(gridID1);
          gridDefMask(gridID2, mask.data());
          vlistChangeGridIndex(vlistID2, index, gridID2);
        }
    }
}

static void
unsetGridMask(const int vlistID1, const int vlistID2)
{
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      const auto gridID2 = gridDuplicate(gridID1);
      gridDefMask(gridID2, nullptr);
      vlistChangeGridIndex(vlistID2, index, gridID2);
    }
}

static void
setProjParams(const int vlistID1, const int vlistID2, const char *projparams)
{
  const auto ngrids = vlistNgrids(vlistID1);
  for (int index = 0; index < ngrids; index++)
    {
      const auto gridID1 = vlistGrid(vlistID1, index);
      if (gridInqType(gridID1) == GRID_PROJECTION)
        {
          const auto gridID2 = gridDuplicate(gridID1);
          cdiDefAttTxt(gridID2, CDI_GLOBAL, "proj4_params", (int) (strlen(projparams)), projparams);
          vlistChangeGridIndex(vlistID2, index, gridID2);
        }
    }
}

static void
readAreaFromFile(const std::string &areafile, Varray<double> &gridcellArea)
{
  bool searchName = false;
  std::string filename = areafile;
  std::string varname;

  if (!fileExists(areafile.c_str()))
    {
      auto const pos = filename.find_last_of(':');
      if (pos > 1 && pos < (filename.size() - 1))
        {
          varname = filename.substr(pos + 1);
          filename = filename.substr(0, pos);
          searchName = true;
        }
    }

  const auto streamID = streamOpenReadLocked(filename.c_str());
  const auto vlistID = streamInqVlist(streamID);

  int svarID = 0;
  if (searchName)
    {
      char name[CDI_MAX_NAME];
      const auto nvars = vlistNvars(vlistID);
      int varID;
      for (varID = 0; varID < nvars; ++varID)
        {
          vlistInqVarName(vlistID, varID, name);
          if (cstrIsEqual(name, varname.c_str()))
            {
              svarID = varID;
              break;
            }
        }
      if (varID == nvars) cdoAbort("Variable %s not found in %s\n", varname.c_str(), filename.c_str());
    }

  const auto nrecs = streamInqTimestep(streamID, 0);
  int varID, levelID;
  for (int recID = 0; recID < nrecs; ++recID)
    {
      streamInqRecord(streamID, &varID, &levelID);
      if (varID == svarID)
        {
          const auto gridID = vlistInqVarGrid(vlistID, varID);
          const auto areasize = gridInqSize(gridID);
          gridcellArea.resize(areasize);

          size_t nmiss;
          streamReadRecord(streamID, gridcellArea.data(), &nmiss);
          break;
        }
    }

  streamClose(streamID);
}

void *
Setgrid(void *process)
{
  int nrecs;
  int varID, levelID;
  int gridID2 = -1;
  int gridtype = -1;
  size_t nmiss;
  bool lregular = false;
  bool lregularnn = false;
  bool ldereference = false;
  int number = 0, position = 0;
  bool lbounds = true;
  std::string gridname;
  const char *griduri = nullptr;
  const char *projparams = nullptr;
  std::vector<int> grid2_vgpm;
  Varray<double> gridmask, gridcellArea;

  cdoInitialize(process);

  // clang-format off
  const auto SETGRID       = cdoOperatorAdd("setgrid",       0, 0, "grid description file or name");
  const auto SETGRIDTYPE   = cdoOperatorAdd("setgridtype",   0, 0, "grid type");
  const auto SETGRIDAREA   = cdoOperatorAdd("setgridarea",   0, 0, "filename with area weights");
  const auto SETGRIDMASK   = cdoOperatorAdd("setgridmask",   0, 0, "filename with grid mask");
  const auto UNSETGRIDMASK = cdoOperatorAdd("unsetgridmask", 0, 0, nullptr);
  const auto SETGRIDNUMBER = cdoOperatorAdd("setgridnumber", 0, 0, "grid number and optionally grid position");
  const auto SETGRIDURI    = cdoOperatorAdd("setgriduri",    0, 0, "reference URI of the horizontal grid");
  const auto USEGRIDNUMBER = cdoOperatorAdd("usegridnumber", 0, 0, "use existing grid identified by grid number");
  const auto SETPROJPARAMS = cdoOperatorAdd("setprojparams", 0, 0, "proj library parameter (e.g.: +init=EPSG:3413)");
  // clang-format on

  const auto operatorID = cdoOperatorID();

  if (operatorID != UNSETGRIDMASK) operatorInputArg(cdoOperatorEnter(operatorID));

  if (operatorID == SETGRID)
    {
      operatorCheckArgc(1);
      gridID2 = cdoDefineGrid(cdoOperatorArgv(0));
    }
  else if (operatorID == SETGRIDTYPE)
    {
      operatorCheckArgc(1);
      gridname = cdoOperatorArgv(0);

      // clang-format off
      if      (gridname == "curvilinear0")  {gridtype = GRID_CURVILINEAR; lbounds = false;}
      else if (gridname == "curvilinear")   {gridtype = GRID_CURVILINEAR; lbounds = true;}
      else if (gridname == "cell")           gridtype = GRID_UNSTRUCTURED;
      else if (gridname == "unstructured0") {gridtype = GRID_UNSTRUCTURED; lbounds = false;}
      else if (gridname == "unstructured")  {gridtype = GRID_UNSTRUCTURED; lbounds = true;}
      else if (gridname == "generic")        gridtype = GRID_GENERIC;
      else if (gridname == "dereference")    ldereference = true;
      else if (gridname == "lonlat")         gridtype = GRID_LONLAT;
      else if (gridname == "gaussian")       gridtype = GRID_GAUSSIAN;
      else if (gridname == "regularnn")     {gridtype = GRID_GAUSSIAN; lregularnn = true;}
      else if (gridname == "regular")       {gridtype = GRID_GAUSSIAN; lregular = true;}
      else if (gridname == "projection")    {gridtype = GRID_PROJECTION;}
      else cdoAbort("Unsupported grid name: %s", gridname.c_str());
      // clang-format on
    }
  else if (operatorID == SETGRIDAREA)
    {
      operatorCheckArgc(1);

      readAreaFromFile(cdoOperatorArgv(0), gridcellArea);

      if (Options::cdoVerbose)
        {
          const auto areasize = gridcellArea.size();
          auto mmm = varrayMinMaxMean(areasize, gridcellArea);
          cdoPrint("gridcellAreas: %zu %#12.5g%#12.5g%#12.5g", areasize, mmm.min, mmm.mean, mmm.max);
        }
    }
  else if (operatorID == SETGRIDMASK)
    {
      operatorCheckArgc(1);
      const auto maskfile = cdoOperatorArgv(0).c_str();
      const auto streamID = streamOpenReadLocked(maskfile);
      const auto vlistID = streamInqVlist(streamID);

      (void) streamInqTimestep(streamID, 0);
      streamInqRecord(streamID, &varID, &levelID);

      const auto missval = vlistInqVarMissval(vlistID, varID);
      const auto gridID = vlistInqVarGrid(vlistID, varID);
      const auto masksize = gridInqSize(gridID);
      gridmask.resize(masksize);

      streamReadRecord(streamID, gridmask.data(), &nmiss);
      streamClose(streamID);

      for (size_t i = 0; i < masksize; i++)
        if (DBL_IS_EQUAL(gridmask[i], missval)) gridmask[i] = 0;
    }
  else if (operatorID == USEGRIDNUMBER)
    {
      operatorCheckArgc(1);
      number = parameter2int(cdoOperatorArgv(0));
    }
  else if (operatorID == SETGRIDNUMBER)
    {
      if (operatorArgc() >= 1 && operatorArgc() <= 2)
        {
          number = parameter2int(cdoOperatorArgv(0));
          if (operatorArgc() == 2) position = parameter2int(cdoOperatorArgv(1));
        }
      else
        {
          operatorCheckArgc(1);
        }
    }
  else if (operatorID == SETGRIDURI)
    {
      operatorCheckArgc(1);
      griduri = cdoOperatorArgv(0).c_str();
    }
  else if (operatorID == SETPROJPARAMS)
    {
      operatorCheckArgc(1);
      projparams = cdoOperatorArgv(0).c_str();
    }

  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);

  if (operatorID == SETGRID)
    {
      changeGrid(vlistID1, vlistID2, gridID2);
    }
  else if (operatorID == SETGRIDNUMBER || operatorID == SETGRIDURI || operatorID == USEGRIDNUMBER)
    {
      int gridID2x = -1;
      if (operatorID == SETGRIDNUMBER)
        {
          const auto gridID1 = vlistGrid(vlistID1, 0);
          gridID2x = gridCreate(GRID_UNSTRUCTURED, gridInqSize(gridID1));
          cdiDefKeyInt(gridID2x, CDI_GLOBAL, CDI_KEY_NUMBEROFGRIDUSED, number);
          cdiDefKeyInt(gridID2x, CDI_GLOBAL, CDI_KEY_NUMBEROFGRIDINREFERENCE, position);
        }
      else if (operatorID == USEGRIDNUMBER)
        {
          if (number < 1 || number > vlistNgrids(vlistID1))
            cdoAbort("Invalid grid number: %d (max = %d)!", number, vlistNgrids(vlistID1));

          gridID2x = vlistGrid(vlistID1, number - 1);
        }
      else
        {
          const auto gridID1 = vlistGrid(vlistID1, 0);
          gridID2x = gridDuplicate(gridID1);
          cdiDefKeyString(gridID2x, CDI_GLOBAL, CDI_KEY_REFERENCEURI, griduri);
        }

      changeGrid(vlistID1, vlistID2, gridID2x);
    }
  else if (operatorID == SETGRIDTYPE)
    {
      setGridtype(vlistID1, vlistID2, gridtype, gridname, lregular, lregularnn, ldereference, lbounds, grid2_vgpm);
    }
  else if (operatorID == SETGRIDAREA)
    {
      setGridcellArea(vlistID1, vlistID2, gridcellArea);
    }
  else if (operatorID == SETGRIDMASK)
    {
      setGridMask(vlistID1, vlistID2, gridmask);
    }
  else if (operatorID == UNSETGRIDMASK)
    {
      unsetGridMask(vlistID1, vlistID2);
    }
  else if (operatorID == SETPROJPARAMS)
    {
      setProjParams(vlistID1, vlistID2, projparams);
    }

  const auto streamID2 = cdoOpenWrite(1);

  cdoDefVlist(streamID2, vlistID2);
  // vlistPrint(vlistID2);

  auto gridsizemax = (lregular || lregularnn) ? vlistGridsizeMax(vlistID2) : vlistGridsizeMax(vlistID1);

  if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;
  Varray<double> array(gridsizemax);

  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);
          cdoDefRecord(streamID2, varID, levelID);

          cdoReadRecord(streamID1, array.data(), &nmiss);

          const auto gridID1 = vlistInqVarGrid(vlistID1, varID);
          if (lregular || lregularnn)
            {
              if (gridInqType(gridID1) == GRID_GAUSSIAN_REDUCED)
                {
                  const auto missval = vlistInqVarMissval(vlistID1, varID);
                  const int lnearst = lregularnn ? 1 : 0;
                  field2regular(gridID1, vlistInqVarGrid(vlistID2, varID), missval, array.data(), nmiss, lnearst);
                }
            }
          else if (gridInqType(gridID1) == GRID_GME)
            {
              const auto gridsize = gridInqSize(gridID1);
              size_t j = 0;
              for (size_t i = 0; i < gridsize; i++)
                if (grid2_vgpm[i]) array[j++] = array[i];
            }

          cdoWriteRecord(streamID2, array.data(), nmiss);
        }

      tsID++;
    }

  cdoStreamClose(streamID2);
  cdoStreamClose(streamID1);

  cdoFinish();

  return nullptr;
}
