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

*/

#include <cdi.h>

#include "process_int.h"
#include "param_conversion.h"
#include "cdo_wtime.h"
#include <mpim_grid.h>
#include "gridreference.h"
#include "grid_point_search.h"
#include "cdo_options.h"
#include "progress.h"
#include "cimdOmp.h"

void
fillmiss(Field &field1, Field &field2, int nfill)
{
  int i, j;
  size_t nmiss2 = 0;
  int kr, ku, kl, ko;
  int ir, iu, il, io;
  int kh, kv, k1, k2, kk;
  double s1, s2;
  double xr, xu, xl, xo;

  const int gridID = field1.grid;
  const size_t nmiss1 = field1.nmiss;
  const double missval = field1.missval;
  double *array1 = field1.vec_d.data();
  double *array2 = field2.vec_d.data();

  const int nx = gridInqXsize(gridID);
  const int ny = gridInqYsize(gridID);
  const bool globgrid = (bool) gridIsCircular(gridID);

  const int gridtype = gridInqType(gridID);
  if (!(gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN)) cdoAbort("Unsupported grid type: %s!", gridNamePtr(gridtype));

  std::vector<double *> matrix1(ny), matrix2(ny);

  for (j = 0; j < ny; j++)
    {
      matrix1[j] = array1 + j * nx;
      matrix2[j] = array2 + j * nx;
    }

  for (j = 0; j < ny; j++)
    for (i = 0; i < nx; i++)
      {
        if (DBL_IS_EQUAL(matrix1[j][i], missval))
          {
            nmiss2++;

            kr = ku = kl = ko = 0;
            xr = xu = xl = xo = 0.;

            for (ir = i + 1; ir < nx; ir++)
              if (!DBL_IS_EQUAL(matrix1[j][ir], missval))
                {
                  kr = ir - i;
                  xr = matrix1[j][ir];
                  break;
                }

            if (globgrid && ir == nx)
              {
                for (ir = 0; ir < i; ir++)
                  if (!DBL_IS_EQUAL(matrix1[j][ir], missval))
                    {
                      kr = nx + ir - i;
                      xr = matrix1[j][ir];
                      break;
                    }
              }

            for (il = i - 1; il >= 0; il--)
              if (!DBL_IS_EQUAL(matrix1[j][il], missval))
                {
                  kl = i - il;
                  xl = matrix1[j][il];
                  break;
                }

            if (globgrid && il == -1)
              {
                for (il = nx - 1; il > i; il--)
                  if (!DBL_IS_EQUAL(matrix1[j][il], missval))
                    {
                      kl = nx + i - il;
                      xl = matrix1[j][il];
                      break;
                    }
              }

            for (iu = j + 1; iu < ny; iu++)
              if (!DBL_IS_EQUAL(matrix1[iu][i], missval))
                {
                  ku = iu - j;
                  xu = matrix1[iu][i];
                  break;
                }

            for (io = j - 1; io >= 0; io--)
              if (!DBL_IS_EQUAL(matrix1[io][i], missval))
                {
                  ko = j - io;
                  xo = matrix1[io][i];
                  break;
                }

            // printf("%d %d %d %d %d %d %g %g %g %g\n", j,i,kr,kl,ku,ko,xr,xl,xu,xo);

            kh = kl + kr;
            kv = ko + ku;
            // clang-format off
            if      (kh == 0) { k1 = 0; s1 = 0.; }
            else if (kl == 0) { k1 = 1; s1 = xr; }
            else if (kr == 0) { k1 = 1; s1 = xl; }
            else              { k1 = 2; s1 = xr * kl / kh + xl * kr / kh; }

            if      (kv == 0) { k2 = 0; s2 = 0.; }
            else if (ku == 0) { k2 = 1; s2 = xo; }
            else if (ko == 0) { k2 = 1; s2 = xu; }
            else              { k2 = 2; s2 = xu * ko / kv + xo * ku / kv; }
            // clang-format on

            kk = k1 + k2;
            if (kk >= nfill)
              {
                if (kk == 0)
                  cdoAbort("no point found!");
                else if (k1 == 0)
                  matrix2[j][i] = s2;
                else if (k2 == 0)
                  matrix2[j][i] = s1;
                else
                  matrix2[j][i] = s1 * k2 / kk + s2 * k1 / kk;
              }
            else
              matrix2[j][i] = matrix1[j][i];

            // matrix1[j][i] = matrix2[j][i];
          }
        else
          {
            matrix2[j][i] = matrix1[j][i];
          }
      }

  if (nmiss1 != nmiss2) cdoAbort("found only %zu of %zu missing values!", nmiss2, nmiss1);
}

void
fillmiss_one_step(Field &field1, Field &field2, int maxfill)
{
  int gridID, nx, ny, i, j;
  size_t nmiss2 = 0;
  int kr, ku, kl, ko;
  int ir, iu, il, io;
  int kh, kv, k1, k2, kk;
  double s1, s2;
  double xr, xu, xl, xo;
  double missval;
  double *array1, *array2;

  gridID = field1.grid;
  missval = field1.missval;
  array1 = field1.vec_d.data();
  array2 = field2.vec_d.data();

  nx = gridInqXsize(gridID);
  ny = gridInqYsize(gridID);

  std::vector<double *> matrix1(ny), matrix2(ny);

  for (j = 0; j < ny; j++)
    {
      matrix1[j] = array1 + j * nx;
      matrix2[j] = array2 + j * nx;
    }

  for (int fill_iterations = 0; fill_iterations < maxfill; fill_iterations++)
    {
      for (j = 0; j < ny; j++)
        for (i = 0; i < nx; i++)
          {
            if (DBL_IS_EQUAL(matrix1[j][i], missval))
              {
                nmiss2++;

                kr = ku = kl = ko = 0;
                xr = xu = xl = xo = 0.;

                for (ir = i + 1; ir < nx; ir++)
                  if (!DBL_IS_EQUAL(matrix1[j][ir], missval))
                    {
                      kr = ir - i;
                      xr = matrix1[j][ir];
                      break;
                    }

                for (il = i - 1; il >= 0; il--)
                  if (!DBL_IS_EQUAL(matrix1[j][il], missval))
                    {
                      kl = i - il;
                      xl = matrix1[j][il];
                      break;
                    }

                for (iu = j + 1; iu < ny; iu++)
                  if (!DBL_IS_EQUAL(matrix1[iu][i], missval))
                    {
                      ku = iu - j;
                      xu = matrix1[iu][i];
                      break;
                    }

                for (io = j - 1; io >= 0; io--)
                  if (!DBL_IS_EQUAL(matrix1[io][i], missval))
                    {
                      ko = j - io;
                      xo = matrix1[io][i];
                      break;
                    }

                kh = kl + kr;
                kv = ko + ku;
                if (kh == 0)
                  {
                    s1 = 0.;
                    k1 = 0;
                  }
                else if (kl == 0)
                  {
                    s1 = xr;
                    k1 = kr;
                  }
                else if (kr == 0)
                  {
                    s1 = xl;
                    k1 = kl;
                  }
                else
                  {
                    if (kl < kr)
                      {
                        s1 = xl;
                        k1 = kl;
                      }
                    else
                      {
                        s1 = xr;
                        k1 = kr;
                      }
                  }

                if (kv == 0)
                  {
                    s2 = 0.;
                    k2 = 0;
                  }
                else if (ku == 0)
                  {
                    s2 = xo;
                    k2 = ko;
                  }
                else if (ko == 0)
                  {
                    s2 = xu;
                    k2 = ku;
                  }
                else
                  {
                    if (ku < ko)
                      {
                        s2 = xu;
                        k2 = ku;
                      }
                    else
                      {
                        s2 = xo;
                        k2 = ko;
                      }
                  }

                kk = k1 + k2;
                if (kk == 0)
                  matrix2[j][i] = matrix1[j][i];
                else if (k1 == 0)
                  matrix2[j][i] = s2;
                else if (k2 == 0)
                  matrix2[j][i] = s1;
                else
                  {
                    if (k1 <= k2)
                      {
                        matrix2[j][i] = s1;
                      }
                    else
                      {
                        matrix2[j][i] = s2;
                      }
                  }

                // printf("%d %d %2d %2d %2d %2d %2g %2g %2g %2g %2g %2g %2g\n",
                // j,i,kr,kl,ku,ko,xr,xl,xu,xo,s1,s2,matrix2[j][i]);
                /* matrix1[j][i] = matrix2[j][i]; */
              }
            else
              {
                matrix2[j][i] = matrix1[j][i];
              }
          }
      for (j = 0; j < ny; j++)
        for (i = 0; i < nx; i++) matrix1[j][i] = matrix2[j][i];
    }
}

static void
setmisstodis(Field &field1, Field &field2, int numNeighbors)
{
  auto gridID = field1.grid;
  const auto gridID0 = gridID;
  const auto missval = field1.missval;
  auto array1 = field1.vec_d.data();
  auto array2 = field2.vec_d.data();

  const auto gridtype = gridInqType(gridID);
  const auto gridsize = gridInqSize(gridID);

  auto nmiss = field1.nmiss;
  const auto nvals = gridsize - nmiss;

  if (gridtype == GRID_GME) gridID = gridToUnstructured(gridID, 0);
  if (gridtype != GRID_UNSTRUCTURED && gridtype != GRID_CURVILINEAR) gridID = gridToCurvilinear(gridID, 0);

  if (gridtype == GRID_UNSTRUCTURED && !gridHasCoordinates(gridID))
    {
      const auto reference = dereferenceGrid(gridID);
      if (reference.isValid) gridID = reference.gridID;
      if (reference.notFound) cdoAbort("Reference to source grid not found!");
    }

  if (!gridHasCoordinates(gridID)) cdoAbort("Cell center coordinates missing!");

  Varray<double> xvals(gridsize), yvals(gridsize);
  gridInqXvals(gridID, xvals.data());
  gridInqYvals(gridID, yvals.data());

  // Convert lat/lon units if required
  cdo_grid_to_radian(gridID, CDI_XAXIS, gridsize, xvals.data(), "grid center lon");
  cdo_grid_to_radian(gridID, CDI_YAXIS, gridsize, yvals.data(), "grid center lat");

  std::vector<size_t> mindex(nmiss, 1), vindex(nvals, 1);
  Varray<double> lons(nvals), lats(nvals);

  size_t nv = 0, nm = 0;
  for (size_t i = 0; i < gridsize; ++i)
    {
      array2[i] = array1[i];
      if (DBL_IS_EQUAL(array1[i], missval))
        {
          mindex[nm] = i;
          nm++;
        }
      else
        {
          lons[nv] = xvals[i];
          lats[nv] = yvals[i];
          vindex[nv] = i;
          nv++;
        }
    }

  if (nv != nvals) cdoAbort("Internal problem, number of valid values differ!");

  std::vector<knnWeightsType> knnWeights;
  for (int i = 0; i < Threading::ompNumThreads; ++i) knnWeights.push_back(knnWeightsType(numNeighbors));

  double start = Options::cdoVerbose ? cdo_get_wtime() : 0;

  GridPointSearch gps;

  if (nmiss)
    {
      const auto xIsCyclic = false;
      size_t dims[2] = { nvals, 0 };
      gridPointSearchCreate(gps, xIsCyclic, dims, nvals, lons, lats);
      gridPointSearchExtrapolate(gps);
    }

  if (Options::cdoVerbose) cdoPrint("Point search created: %.2f seconds", cdo_get_wtime() - start);

  progress::init();

  start = Options::cdoVerbose ? cdo_get_wtime() : 0;

  double findex = 0;

#ifdef _OPENMP
#pragma omp parallel for default(none) \
    shared(nmiss, knnWeights, findex, mindex, vindex, array1, array2, xvals, yvals, gps, numNeighbors)
#endif
  for (size_t i = 0; i < nmiss; ++i)
    {
#ifdef _OPENMP
#pragma omp atomic
#endif
      findex++;
      if (cdo_omp_get_thread_num() == 0) progress::update(0, 1, findex / nmiss);

      const auto ompthID = cdo_omp_get_thread_num();

      gridSearchPoint(gps, xvals[mindex[i]], yvals[mindex[i]], knnWeights[ompthID]);

      // Compute weights based on inverse distance if mask is false, eliminate those points
      const auto nadds = knnWeights[ompthID].computeWeights();
      if (nadds)
        {
          double result = 0;
          for (size_t n = 0; n < nadds; ++n)
            result += array1[vindex[knnWeights[ompthID].m_addr[n]]] * knnWeights[ompthID].m_dist[n];
          array2[mindex[i]] = result;
        }
    }

  progress::update(0, 1, 1);

  if (Options::cdoVerbose) cdoPrint("Point search nearest: %.2f seconds", cdo_get_wtime() - start);

  gridPointSearchDelete(gps);

  if (gridID0 != gridID) gridDestroy(gridID);
}

void *
Fillmiss(void *process)
{
  int nrecs, varID, levelID;
  void (*fill_method)(Field & fin, Field & fout, int) = &setmisstodis;

  cdoInitialize(process);

  // clang-format off
  const auto FILLMISS        = cdoOperatorAdd("fillmiss"   ,   0, 0, "nfill");
  const auto FILLMISSONESTEP = cdoOperatorAdd("fillmiss2"  ,   0, 0, "nfill");
  const auto SETMISSTONN     = cdoOperatorAdd("setmisstonn" ,  0, 0, "");
  const auto SETMISSTODIS    = cdoOperatorAdd("setmisstodis" , 0, 0, "number of neighbors");
  // clang-format on

  const auto operatorID = cdoOperatorID();

  int nfill = 1;
  if (operatorID == FILLMISS)
    {
      fill_method = &fillmiss;
    }
  else if (operatorID == FILLMISSONESTEP)
    {
      fill_method = &fillmiss_one_step;
    }
  else if (operatorID == SETMISSTONN)
    {
      fill_method = &setmisstodis;
    }
  else if (operatorID == SETMISSTODIS)
    {
      nfill = 4;
      fill_method = &setmisstodis;
    }

  /* Argument handling */
  {
    const auto oargc = operatorArgc();
    if (oargc == 1)
      {
        nfill = parameter2int(cdoOperatorArgv(0));
        if (operatorID == FILLMISS)
          {
            if (nfill < 1 || nfill > 4) cdoAbort("nfill out of range!");
          }
      }
    else if (oargc > 1)
      cdoAbort("Too many arguments!");
  }

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

  const auto streamID2 = cdoOpenWrite(1);

  cdoDefVlist(streamID2, vlistID2);

  const auto gridsizemax = vlistGridsizeMax(vlistID1);

  Field field1, field2;
  field1.resize(gridsizemax);
  field2.resize(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);
          cdoReadRecord(streamID1, field1.vec_d.data(), &field1.nmiss);

          cdoDefRecord(streamID2, varID, levelID);

          if (field1.nmiss == 0)
            {
              cdoWriteRecord(streamID2, field1.vec_d.data(), 0);
            }
          else
            {
              const auto gridID = vlistInqVarGrid(vlistID1, varID);

              if ((operatorID == FILLMISS || operatorID == FILLMISSONESTEP)
                  && (gridInqType(gridID) == GRID_GME || gridInqType(gridID) == GRID_UNSTRUCTURED))
                cdoAbort("%s data unsupported!", gridNamePtr(gridInqType(gridID)));

              field1.grid = gridID;
              field1.size = gridInqSize(field1.grid);
              field1.missval = vlistInqVarMissval(vlistID1, varID);

              field2.grid = field1.grid;
              field2.size = field1.size;
              field2.missval = field1.missval;

              fill_method(field1, field2, nfill);

              cdoWriteRecord(streamID2, field2.vec_d.data(), fieldNumMiss(field2));
            }
        }

      tsID++;
    }

  cdoStreamClose(streamID2);
  cdoStreamClose(streamID1);

  cdoFinish();

  return nullptr;
}
