/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Uwe Schulzweida

*/

/*
   This module contains the following operators:

      Spectral   sp2gp           Spectral to gridpoint
      Spectral   sp2gpl          Spectral to gridpoint linear (sp2gp,linear)
      Spectral   gp2sp           Gridpoint to spectral
      Spectral   gp2spl          Gridpoint to spectral linear (gp2sp,linear)
      Spectral   sp2sp           Spectral to spectral
      Spectral   spcut           Cut spectral wave number
*/

#include <cdi.h>

#include "cdo_vlist.h"
#include "process_int.h"
#include "param_conversion.h"
#include <mpim_grid.h>
#include "griddes.h"
#include "specspace.h"

void *
Spectral(void *process)
{
  int gridID1 = -1, gridID2 = -1;
  int defaultTrunc = 0;
  Varray<int> waves;
  SP_Transformation spTrans;

  cdo_initialize(process);

  const auto dataIsUnchanged = data_is_unchanged();

  // clang-format off
  const auto GP2SP  = cdo_operator_add("gp2sp",  0, 0, nullptr);
  const auto GP2SPL = cdo_operator_add("gp2spl", 0, 0, nullptr);
  const auto SP2GP  = cdo_operator_add("sp2gp",  0, 0, nullptr);
  const auto SP2GPL = cdo_operator_add("sp2gpl", 0, 0, nullptr);
  const auto SP2SP  = cdo_operator_add("sp2sp",  0, 0, nullptr);
  const auto SPCUT  = cdo_operator_add("spcut",  0, 0, nullptr);

  const auto operatorID = cdo_operator_id();

  const auto lgp2sp = (operatorID == GP2SP || operatorID == GP2SPL);
  const auto lsp2gp = (operatorID == SP2GP || operatorID == SP2GPL);
  const auto linear = (operatorID == GP2SPL || operatorID == SP2GPL);

  int (*nlat2ntr)(int) = linear ? nlat_to_ntr_linear : nlat_to_ntr;
  const char *ctype = linear ? "l" : "";

  const auto paramArgc = cdo_operator_argc();
  if ((lgp2sp || lsp2gp) && paramArgc == 1)
    {
      const auto &parg = cdo_operator_argv(0);
      const auto pos = parg.find("=");
      if (lgp2sp && pos > 0 && parg.substr(0, pos) == "trunc")
        {
          defaultTrunc = parameter_to_int(parg.substr(pos + 1));
        }
      else
        {
          auto type = parameter_to_word((pos > 0 && parg.substr(0, pos) == "type") ? parg.substr(pos + 1) : parg);
          if      (type == "linear")    { nlat2ntr = nlat_to_ntr_linear; ctype = "l"; }
          else if (type == "cubic")     { nlat2ntr = nlat_to_ntr_cubic; ctype = "c"; }
          else if (type == "quadratic") { nlat2ntr = nlat_to_ntr; }
          else cdo_abort("Unsupported type: %s\n", type);
        }
    }
  else if (paramArgc > 0)
    {
      cdo_abort("Too many parameters");
    }
  // clang-format on

  const auto streamID1 = cdo_open_read(0);

  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);
  const auto vlistID2 = vlistDuplicate(vlistID1);

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

  auto gridIDsp = vlist_get_first_spectral_grid(vlistID1);
  auto gridIDgp = vlist_get_first_gaussian_grid(vlistID1);

  // define output grid
  if (lgp2sp)
    {
      if (gridIDgp == -1) cdo_warning("No data on regular Gaussian grid found!");

      gridID1 = gridIDgp;

      if (gridID1 != -1)
        {
          const long nlon = gridInqXsize(gridID1);
          const long nlat = gridInqYsize(gridID1);

          long ntr = (defaultTrunc > 0) ? defaultTrunc : nlat2ntr(nlat);
          if (Options::cdoVerbose) cdo_print("trunc=%ld\n", ntr);

          if (gridIDsp != -1)
            if (ntr != gridInqTrunc(gridIDsp)) gridIDsp = -1;

          if (gridIDsp == -1)
            {
              gridIDsp = gridCreate(GRID_SPECTRAL, (ntr + 1) * (ntr + 2));
              gridDefTrunc(gridIDsp, ntr);
              gridDefComplexPacking(gridIDsp, 1);
            }

          if (gridIDsp == -1) cdo_abort("Computation of spherical harmonics failed!");

          gridID2 = gridIDsp;

          ntr = gridInqTrunc(gridID2);
          spTrans.init(nlon, nlat, ntr, PolFlag::FC2SP);
        }
    }
  else if (lsp2gp)
    {
      if (gridIDsp == -1) cdo_warning("No spectral data found!");

      gridID1 = gridIDsp;

      if (gridID1 != -1)
        {
          if (gridIDgp != -1)
            {
              const long nlat = gridInqYsize(gridIDgp);
              const long ntr = nlat2ntr(nlat);
              if (gridInqTrunc(gridIDsp) != ntr) gridIDgp = -1;
            }

          if (gridIDgp == -1)
            {
              char gridname[20];
              snprintf(gridname, sizeof(gridname), "t%s%dgrid", ctype, gridInqTrunc(gridIDsp));
              gridIDgp = grid_from_name(gridname);
            }

          gridID2 = gridIDgp;

          const long ntr = gridInqTrunc(gridID1);
          const long nlon = gridInqXsize(gridID2);
          const long nlat = gridInqYsize(gridID2);
          spTrans.init(nlon, nlat, ntr, PolFlag::SP2FC);
        }
    }
  else if (operatorID == SP2SP)
    {
      gridID1 = gridIDsp;

      operator_input_arg("truncation");
      if (gridID1 != -1)
        {
          if (!isdigit(cdo_operator_argv(0)[0])) cdo_abort("parameter truncation must comprise only digits [0-9]!");
          const long ntr = parameter_to_int(cdo_operator_argv(0));
          const long nsp = (ntr + 1) * (ntr + 2);
          gridIDsp = gridCreate(GRID_SPECTRAL, nsp);
          gridDefTrunc(gridIDsp, ntr);
          gridDefComplexPacking(gridIDsp, 1);
        }
      else
        cdo_abort("No spectral data found!");

      gridID2 = gridIDsp;
    }
  else if (operatorID == SPCUT)
    {
      gridID1 = gridIDsp;

      operator_input_arg("wave numbers");
      if (gridID1 != -1)
        {
          const long maxntr = 1 + gridInqTrunc(gridID1);
          const auto wnums = cdo_argv_to_int(cdo_get_oper_argv());
          const long ncut = wnums.size();
          waves.resize(maxntr);
          for (long i = 0; i < maxntr; ++i) waves[i] = 1;
          for (long i = 0; i < ncut; ++i)
            {
              const long j = wnums[i] - 1;
              if (j < 0 || j >= maxntr) cdo_abort("wave number %ld out of range (min=1, max=%l qd)!", wnums[i], maxntr);
              waves[j] = 0;
            }
        }
      else
        cdo_abort("No spectral data found!");

      gridID2 = gridIDsp;
    }

  const auto nvars = vlistNvars(vlistID2);
  std::vector<bool> processVars(nvars);
  for (int varID = 0; varID < nvars; ++varID) processVars[varID] = (gridID1 == vlistInqVarGrid(vlistID1, varID));

  if (gridID1 != -1) vlistChangeGrid(vlistID2, gridID1, gridID2);

  const auto streamID2 = cdo_open_write(1);
  cdo_def_vlist(streamID2, vlistID2);

  const auto gridsizemax = vlistGridsizeMax(vlistID1);
  Varray<double> array1(gridsizemax), array2;
  if (gridID2 != -1) array2.resize(gridInqSize(gridID2));

  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      cdo_taxis_copy_timestep(taxisID2, taxisID1);
      cdo_def_timestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; ++recID)
        {
          int varID, levelID;
          cdo_inq_record(streamID1, &varID, &levelID);

          if (processVars[varID])
            {
              size_t nmiss;
              cdo_read_record(streamID1, array1.data(), &nmiss);
              if (nmiss) cdo_abort("Missing values unsupported for spectral data!");

              gridID1 = vlistInqVarGrid(vlistID1, varID);
              if (lgp2sp)
                grid2spec(spTrans, gridID1, array1, gridID2, array2);
              else if (lsp2gp)
                spec2grid(spTrans, gridID1, array1, gridID2, array2);
              else if (operatorID == SP2SP)
                spec2spec(gridID1, array1, gridID2, array2);
              else if (operatorID == SPCUT)
                speccut(gridID1, array1, array2, waves);

              cdo_def_record(streamID2, varID, levelID);
              cdo_write_record(streamID2, array2.data(), nmiss);
            }
          else
            {
              cdo_def_record(streamID2, varID, levelID);
              if (dataIsUnchanged)
                {
                  cdo_copy_record(streamID2, streamID1);
                }
              else
                {
                  size_t nmiss;
                  cdo_read_record(streamID1, array1.data(), &nmiss);
                  cdo_write_record(streamID2, array1.data(), nmiss);
                }
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
