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

  Copyright (C) 2003-2021 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:

      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 varID, levelID;
  int gridID1 = -1, gridID2 = -1;
  size_t nmiss;
  std::vector<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 bool lgp2sp = (operatorID == GP2SP || operatorID == GP2SPL);
  const bool lsp2gp = (operatorID == SP2GP || operatorID == SP2GPL);
  const bool linear = (operatorID == GP2SPL || operatorID == SP2GPL);

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

  if ((lgp2sp || lsp2gp) && cdo_operator_argc() == 1)
    {
      std::string type = parameter_to_word(cdo_operator_argv(0));
      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.c_str());
    }
  // 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 Gaussian grid found!");

      gridID1 = gridIDgp;

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

          long ntr = nlat2ntr(nlat);

          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 && gridInqType(vlistGrid(vlistID1, 0)) == GRID_GAUSSIAN_REDUCED)
            cdo_abort("Gaussian reduced grid found. Use option -R to convert it to a regular grid!");

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

          gridID2 = gridIDsp;

          ntr = gridInqTrunc(gridID2);
          spTrans.init(nlon, nlat, ntr, 0);
        }
    }
  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, 0);
        }
    }
  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> vars(nvars);
  for (varID = 0; varID < nvars; varID++) vars[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++)
        {
          cdo_inq_record(streamID1, &varID, &levelID);

          if (vars[varID])
            {
              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.data(), gridID2, array2.data());
              else if (lsp2gp)
                spec2grid(spTrans, gridID1, array1.data(), gridID2, array2.data());
              else if (operatorID == SP2SP)
                spec2spec(gridID1, array1.data(), gridID2, array2.data());
              else if (operatorID == SPCUT)
                speccut(gridID1, array1.data(), array2.data(), waves.data());

              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
                {
                  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;
}
