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

      Intlevel   intlevel3d      Linear level interpolation on a 3d vertical coordinates variable
      Intlevel   intlevelx3d     Linear level interpolation on a 3d vertical coordinates variable with extrapolation
*/

#include <cdi.h>

#include "cdo_options.h"
#include "process_int.h"
#include "cdo_vlist.h"
#include "cdi_lockedIO.h"

void vert_interp_lev3d(size_t gridsize, double missval, const Varray<double> &vardata1, Varray<double> &vardata2, int nlev2,
                       const Varray<int> &lev_idx1, const Varray<int> &lev_idx2, const Varray<double> &lev_wgt1, const Varray<double> &lev_wgt2);
void vert_gen_weights(int expol, int nlev1, const Varray<double> &lev1, int nlev2, const Varray<double> &lev2,
                      Varray<int> &lev_idx1, Varray<int> &lev_idx2, Varray<double> &lev_wgt1, Varray<double> &lev_wgt2);
void vert_init_level_0_and_N(int nlev, size_t gridsize, Varray<double> &zlevels);

/*
 * Create weights for the 3d vertical coordinate
 *
 * The resulting index sets lev_idx1 and lev_idx2 contain absolute numbers,i.e.
 * wrt. the given gridsize. They can directly be used to read values from 3d data fields.
 *
 * 3d version of vert_gen_weights() (src/Intlevel.cc)
 */
static void
vert_gen_weights3d(bool expol, int nlev1, size_t gridsize, Varray<double> &xlev1, int nlev2, Varray<double> &xlev2,
                   Varray<int> &xlev_idx1, Varray<int> &xlev_idx2, Varray<double> &xlev_wgt1, Varray<double> &xlev_wgt2)
{
  Varray<double> lev1(nlev1), lev2(nlev2);
  Varray<double> lev_wgt1(nlev2), lev_wgt2(nlev2);
  Varray<int> lev_idx1(nlev2), lev_idx2(nlev2);

  for (size_t i = 0; i < gridsize; i++)
    {
      for (int k = 0; k < nlev1; ++k) lev1[k] = xlev1[k * gridsize + i];
      for (int k = 0; k < nlev2; ++k) lev2[k] = xlev2[k * gridsize + i];

      vert_gen_weights(expol, nlev1, lev1, nlev2, lev2, lev_idx1, lev_idx2, lev_wgt1, lev_wgt2);

      for (int k = 0; k < nlev2; ++k) xlev_idx1[k * gridsize + i] = lev_idx1[k] * gridsize + i;
      for (int k = 0; k < nlev2; ++k) xlev_idx2[k * gridsize + i] = lev_idx2[k] * gridsize + i;
      for (int k = 0; k < nlev2; ++k) xlev_wgt1[k * gridsize + i] = lev_wgt1[k];
      for (int k = 0; k < nlev2; ++k) xlev_wgt2[k * gridsize + i] = lev_wgt2[k];
    }
}

void *
Intlevel3d(void *process)
{
  size_t gridsizeo;
  int nrecs;
  int i;
  int varID = -1, levelID = -1;
  int nvct;
  int zaxisID1 = -1, zaxisID3;
  int gridID3 = -1;
  int nlevo, maxlev;
  Varray<double> zlevels_out;
  size_t zlevels_out_miss = 0;
  char varname[CDI_MAX_NAME];

  cdoInitialize(process);

  // clang-format off
  const auto INTLEVEL3D  = cdoOperatorAdd("intlevel3d",  0, 0, nullptr);
  const auto INTLEVELX3D = cdoOperatorAdd("intlevelx3d", 0, 0, nullptr);
  // clang-format on

  (void) (INTLEVEL3D);  // CDO_UNUSED

  const auto operatorID = cdoOperatorID();

  const bool expol = (operatorID == INTLEVELX3D);

  operatorInputArg("icoordinate");

  const auto streamID1 = cdoOpenRead(0);  // input data
  const auto streamID2 = cdoOpenRead(1);  // 3d target vertical coordinate
  const auto streamID3 = cdoOpenWrite(2); // output stream

  /* Read filename from Parameter */
  operatorInputArg("filename for vertical source coordinates variable");
  operatorCheckArgc(1);

  const auto vlistID1 = cdoStreamInqVlist(streamID1);
  const auto taxisID1 = vlistInqTaxis(vlistID1);

  const auto vlistID2 = cdoStreamInqVlist(streamID2);
  const auto vlistID3 = vlistDuplicate(vlistID1);
  const auto taxisID3 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID3, taxisID1);

  VarList varList1;
  varListInit(varList1, vlistID1);

  /*
   * Read 3d input coordinate (streamID0)
   * * two additional levels are added (top + bottom) for later extrapolation checking
   */
  size_t gridsizei;
  int nlevi;
  Varray<double> zlevels_in;
  size_t zlevels_in_miss = 0;

  {
    const auto streamID0 = streamOpenReadLocked(cdoOperatorArgv(0).c_str()); /*  3d vertical input coordinate */
    const auto vlistID0 = streamInqVlist(streamID0);

    if (vlistNvars(vlistID0) != 1) cdoAbort("Only one single variable is allowed!");

    const auto gridID = vlistInqVarGrid(vlistID0, 0);
    const auto zaxisID = vlistInqVarZaxis(vlistID0, 0);
    gridsizei = gridInqSize(gridID);  // horizontal gridsize of input z coordinate
    nlevi = zaxisInqSize(zaxisID);    // number of input levels for later use

    zlevels_in.resize(gridsizei * (nlevi + 2));
    nrecs = streamInqTimestep(streamID0, 0);
    if (Options::cdoVerbose) cdoPrint("%d records input 3d vertical height", nrecs);

    for (int recID = 0; recID < nrecs; recID++)
      {
        streamInqRecord(streamID0, &varID, &levelID);
        const auto offset = gridsizei + gridsizei * levelID;
        streamReadRecord(streamID0, &zlevels_in[offset], &zlevels_in_miss);
      }

    streamClose(streamID0);
  }

  /*
   * Read 3d output coordinate (streamID2)
   */
  {
    if (vlistNvars(vlistID2) != 1) cdoAbort("Only one single variable is allowed!");

    const auto gridID = vlistInqVarGrid(vlistID2, varID);
    gridID3 = gridID;
    const auto zaxisID = vlistInqVarZaxis(vlistID2, varID);
    const auto gridsize = gridInqSize(gridID);
    const auto nlevel = zaxisInqSize(zaxisID);

    zlevels_out.resize(gridsize * nlevel);
    nlevo = nlevel;       // number of output levels for later use
    gridsizeo = gridsize; // horizontal gridsize of output z coordinate
    nrecs = cdoStreamInqTimestep(streamID2, 0);
    if (Options::cdoVerbose) cdoPrint("%d records target 3d vertical height and gridsize %zu", nrecs, gridsize);

    for (int recID = 0; recID < nrecs; recID++)
      {
        cdoInqRecord(streamID2, &varID, &levelID);
        const auto offset = gridsize * levelID;
        cdoReadRecord(streamID2, &zlevels_out[offset], &zlevels_out_miss);
      }
  }

  /* Missing values are not allowed for coordinate variables */
  if (0 != zlevels_in_miss)
    cdoAbort("Input vertical coordinate variables are not allowed to contain missing values.");
  else if (Options::cdoVerbose)
    cdoPrint("Input vertical coordinate has no missing values.");

  if (0 != zlevels_out_miss)
    cdoAbort("Output vertical coordinate variables are not allowd to contain missing values.");
  else if (Options::cdoVerbose)
    cdoPrint("Output vertical coordinate has no missing values.");

  /*
   * gridsize of input and output vertical coordinate must be equal (later use of gridsizeo ONLY)
   */
  if (gridsizei != gridsizeo) cdoAbort("Input and output vertical coordinate must have the same gridsize!");

  auto gridSize = gridsizeo;

  /* input and output vertical coordinates must have exactly the same horizontal grid */
  if (gridsizei != gridsizeo)
    {
      /* i =0; printf ( "lonIn:%g latIn:%g lonOut:%g
       * latOut:%g\n",lonIn[i],latIn[i],lonOut[i],latOut[i] ); */
      cdoAbort("Input and output vertical coordinates do NOT exactly have the same horizontal grid.");
    }

  /*
   * Check for the correct vertical axis in the input: Variables with the same
   * number of levels as the input vertical levels from operators parameter
   * (streamID0). Variables with a different z-axis should be copied into output.
   */
  const auto nzaxis = vlistNzaxis(vlistID1);
  for (i = 0; i < nzaxis; ++i)
    {
      const auto zaxisID = vlistZaxis(vlistID1, i);
      const auto nlevel = zaxisInqSize(zaxisID);
      if (nlevel == nlevi)
        {
          zaxisID1 = zaxisID;
          break;
        }
    }
  if (i == nzaxis) cdoAbort("No processable variable found (vertical coordinate differ)!");

  const auto ngrids = vlistNgrids(vlistID1);
  for (i = 0; i < ngrids; ++i)
    {
      const auto gridsize = gridInqSize(vlistGrid(vlistID1, i));
      if (gridsize == gridSize) break;
    }
  if (i == nzaxis) cdoAbort("No processable variable found (grid coordinate differ)!");

  vert_init_level_0_and_N(nlevi, gridSize, zlevels_in);

  /*
   * Create weights for later interpolation - assumption: input vertical correct is constant in time
   */
  Varray<int> lev_idx1(nlevo * gridSize), lev_idx2(nlevo * gridSize);
  Varray<double> lev_wgt1(nlevo * gridSize), lev_wgt2(nlevo * gridSize);

  vert_gen_weights3d(expol, nlevi + 2, gridSize, zlevels_in, nlevo, zlevels_out, lev_idx1, lev_idx2, lev_wgt1, lev_wgt2);

  /*
   * Copy z-axis information to output z-axis
   */
  zaxisID3 = zaxisCreate(zaxisInqType(zaxisID1), nlevo);
  Varray<double> lev2(nlevo);
  /* fill values with its indices */
  for (i = 0; i < nlevo; i++) lev2[i] = (double) i + 1;
  zaxisDefLevels(zaxisID3, lev2.data());
  cdiDefKeyString(zaxisID3, CDI_GLOBAL, CDI_KEY_NAME, "lev");
  /* copy VCT from input vlistID1 to output vlistID3 if there is one */
  nvct = zaxisInqVctSize(zaxisID1);
  if (nvct > 0) zaxisDefVct(zaxisID3, zaxisInqVctSize(zaxisID1), zaxisInqVctPtr(zaxisID1));

  for (i = 0; i < nzaxis; i++)
    if (zaxisID1 == vlistZaxis(vlistID1, i)) vlistChangeZaxisIndex(vlistID3, i, zaxisID3);
  /* add the vertical output field to the output stream */
  auto oz3dvarID = vlistDefVar(vlistID3, gridID3, zaxisID3, TIME_VARYING);
  {
    char str[CDI_MAX_NAME];
    vlistInqVarName(vlistID2, 0, str);
    cdiDefKeyString(vlistID3, oz3dvarID, CDI_KEY_NAME, str);
    str[0] = 0;
    vlistInqVarLongname(vlistID2, 0, str);
    if (str[0]) cdiDefKeyString(vlistID3, oz3dvarID, CDI_KEY_LONGNAME, str);
    str[0] = 0;
    vlistInqVarUnits(vlistID2, 0, str);
    if (str[0]) cdiDefKeyString(vlistID3, oz3dvarID, CDI_KEY_UNITS, str);
  }

  cdoDefVlist(streamID3, vlistID3);

  VarList varList3;
  varListInit(varList3, vlistID3);

  maxlev = nlevi > nlevo ? nlevi : nlevo;
  const auto nvars = vlistNvars(vlistID1);

  std::vector<bool> vars(nvars);
  std::vector<bool> varinterp(nvars); // marker for variables to be interpolated
  std::vector<std::vector<size_t>> varnmiss(nvars); // can for missing values of arbitrary variables
  Varray2D<double> vardata1(nvars);   // input
  Varray2D<double> vardata2(nvars);   // output

  // by default no variable should be interpolated
  for (varID = 0; varID < nvars; varID++) varinterp[varID] = false;

  for (varID = 0; varID < nvars; varID++)
    {
      const auto zaxisID = varList1[varID].zaxisID;
      const auto gridsize = varList1[varID].gridsize;
      const auto nlevel = varList1[varID].nlevels;

      vlistInqVarName(vlistID1, varID, varname);

      vardata1[varID].resize(gridsize * nlevel);
      /*  variabls for interpolation:
       *  * have the required vertical axis, i.e. the correct number of levels (nlevi)
       *  * have the same number of horizontal grid points (i.e. same gridSize) like the two vertical coordinates
       *  * are NOT the output vertical coordinates itself
       */
      if (zaxisID == zaxisID1 && varID != oz3dvarID && gridsize == gridSize && gridsizeo == gridsize)
        {
          varinterp[varID] = true;
          vardata2[varID].resize(gridsize * nlevo);
          varnmiss[varID].resize(maxlev, 0);
        }
      else
        {
          varinterp[varID] = false;
          varnmiss[varID].resize(nlevel);
          if (Options::cdoVerbose) cdoPrint("Ignore variable %s (levels=%d gridsize=%zu)!", varname, nlevel, gridsize);
        }
    }

  for (varID = 0; varID < nvars; varID++)
    {
      if (varinterp[varID]) break;
    }
  if (varID == nvars) cdoAbort("No processable variable found!");

  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      for (varID = 0; varID < nvars; ++varID) vars[varID] = false;

      taxisCopyTimestep(taxisID3, taxisID1);
      cdoDefTimestep(streamID3, tsID);

      /*
       * Read the whole 3d data field
       */
      for (int recID = 0; recID < nrecs; recID++)
        {
          cdoInqRecord(streamID1, &varID, &levelID);
          const auto gridsize = varList1[varID].gridsize;
          const auto offset = gridsize * levelID;
          cdoReadRecord(streamID1, &vardata1[varID][offset], &varnmiss[varID][levelID]);
          vars[varID] = true;
        }

      /* Perform the interpolation on all valid data variables */
      for (varID = 0; varID < nvars; varID++)
        {
          if (vars[varID] && varinterp[varID])
            {
              const auto gridsize = varList1[varID].gridsize;
              const auto missval = varList1[varID].missval;

              vert_interp_lev3d(gridsize, missval, vardata1[varID], vardata2[varID], nlevo, lev_idx1, lev_idx2, lev_wgt1, lev_wgt2);

              for (levelID = 0; levelID < nlevo; levelID++)
                {
                  const auto offset = gridsize * levelID;
                  varnmiss[varID][levelID] = arrayNumMV(gridsize, &vardata2[varID][offset], missval);
                }
            }
          else
            {
              vlistInqVarName(vlistID1, varID, varname);
              if (Options::cdoVerbose && tsID <= 1) cdoPrint("Perform no interpolation on variable %s", varname);
            }
        }

      /* write the output */
      for (varID = 0; varID < nvars; varID++)
        {
          if (vars[varID])
            {
              for (levelID = 0; levelID < varList3[varID].nlevels; levelID++)
                {
                  const auto offset = varList3[varID].gridsize * levelID;
                  auto single = varinterp[varID] ? &vardata2[varID][offset] : &vardata1[varID][offset];
                  cdoDefRecord(streamID3, varID, levelID);
                  cdoWriteRecord(streamID3, single, varnmiss[varID][levelID]);
                }
            }
        }

      /* copy output z coordinate to output stream */
      for (levelID = 0; levelID < varList3[oz3dvarID].nlevels; levelID++)
        {
          const auto offset = varList3[oz3dvarID].gridsize * levelID;
          cdoDefRecord(streamID3, oz3dvarID, levelID);
          cdoWriteRecord(streamID3, &zlevels_out[offset], 0);
        }

      tsID++;
    }

  cdoStreamClose(streamID1);
  cdoStreamClose(streamID2);
  cdoStreamClose(streamID3);

  vlistDestroy(vlistID3);

  cdoFinish();

  return nullptr;
}
