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

      Spectrum    spectrum         Spectrum
*/

#include <cdi.h>

#include "process_int.h"
#include "cdo_vlist.h"
#include "param_conversion.h"
#include "cdo_fft.h"
#include "field_functions.h"

static void
spectrum(int nrec, double *data, double *spectrum, double *real, double *imag, const double *window, double wssum, int detrend,
         int seg_n, int seg_l)
{
  int k;
  double sumx, sumkx;
  double a, b;
  int seg_i, offset;
  int bit;

  for (bit = seg_l; !(bit & 1); bit >>= 1)
    ;

  if (detrend == 1)
    {
      sumx = 0;
      for (k = 0; k < nrec; ++k) sumx += data[k];
      sumx /= nrec;
      for (k = 0; k < nrec; ++k) data[k] -= sumx;
    }
  else if (detrend == 2)
    {
      sumx = sumkx = 0;
      for (k = 0; k < nrec; ++k)
        {
          sumx += data[k];
          sumkx += k * data[k];
        }
      b = (sumkx - sumx * (nrec - 1) / 2.) / ((nrec + 1) * nrec * (nrec - 1) / 12.);
      a = sumx / nrec - b * (nrec - 1) / 2.;
      for (k = 0; k < nrec; ++k) data[k] -= a + b * k;
    }

  Varray<double> work_r, work_i;
  if (bit != 1)
    {
      work_r.resize(seg_l);
      work_i.resize(seg_l);
    }

  for (seg_i = 0; seg_i < seg_n; seg_i += 2)
    {
      offset = (seg_n == 1) ? 0 : (int) ((double) (nrec - seg_l) / (seg_n - 1) * seg_i);

      for (k = 0; k < seg_l; ++k) real[k] = data[offset + k];

      if (detrend == 3)
        {
          sumx = sumkx = 0;
          for (k = 0; k < seg_l; ++k)
            {
              sumx += real[k];
              sumkx += k * real[k];
            }

          b = (sumkx - sumx * (seg_l - 1) / 2.) / ((seg_l + 1) * seg_l * (seg_l - 1) / 12.);
          a = sumx / seg_l - b * (seg_l - 1) / 2.;

          for (k = 0; k < seg_l; ++k) real[k] -= a + b * k;
        }

      for (k = 0; k < seg_l; ++k) real[k] *= window[k];

      if (seg_i + 1 < seg_n)
        {
          offset = (seg_n == 1) ? 0 : (int) ((double) (nrec - seg_l) / (seg_n - 1) * (seg_i + 1));
          for (k = 0; k < seg_l; ++k) imag[k] = data[offset + k];
          if (detrend == 3)
            {
              sumx = sumkx = 0;
              for (k = 0; k < seg_l; ++k)
                {
                  sumx += imag[k];
                  sumkx += k * imag[k];
                }

              b = (sumkx - sumx * (seg_l - 1) / 2.) / ((seg_l + 1) * seg_l * (seg_l - 1) / 12.);
              a = sumx / seg_l - b * (seg_l - 1) / 2.;

              for (k = 0; k < seg_l; ++k) imag[k] -= a + b * k;
            }

          for (k = 0; k < seg_l; ++k) imag[k] *= window[k];
        }
      else
        for (k = 0; k < seg_l; ++k) imag[k] = 0;

      if (bit == 1)  // seg_l is a power of 2
        cdo::fft(real, imag, seg_l, 1);
      else
        cdo::ft_r(real, imag, seg_l, 1, work_r.data(), work_i.data());

      spectrum[0] += real[0] * real[0] + imag[0] * imag[0];

      for (k = 1; k < (seg_l + 1) / 2; ++k)
        spectrum[k]
            += real[k] * real[k] + imag[k] * imag[k] + real[seg_l - k] * real[seg_l - k] + imag[seg_l - k] * imag[seg_l - k];

      if (!(seg_l & 1)) spectrum[seg_l / 2] += real[seg_l / 2] * real[seg_l / 2] + imag[seg_l / 2] * imag[seg_l / 2];
    }

  for (k = 0; k <= seg_l / 2; ++k) spectrum[k] *= seg_l / (seg_n * wssum);

  spectrum[0] *= 2;

  if (!(seg_l & 1)) spectrum[seg_l / 2] *= 2;
}

class ModuleSpectrum
{
  int varID, levelID;
  int k;
  int nalloc = 0;
  size_t nmiss;
  int freq;
  int nvars;

  CdoStreamID streamID1;
  CdoStreamID streamID2;

  int taxisID1;
  int taxisID2;
  int vlistID1;

  VarList varList;
  FieldVector3D vars;
  std::vector<CdiDateTime> vDateTimes;

public:
  void
  init(void *process)
  {
    cdo_initialize(process);

    streamID1 = cdo_open_read(0);

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

    taxisID1 = vlistInqTaxis(vlistID1);
    taxisID2 = cdo_taxis_create(TAXIS_ABSOLUTE);
    vlistDefTaxis(vlistID2, taxisID2);

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

    varList_init(varList, vlistID1);

    nvars = vlistNvars(vlistID1);
  }

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

        if (tsID >= nalloc)
          {
            constexpr int NALLOC_INC = 1024;
            nalloc += NALLOC_INC;
            vDateTimes.resize(nalloc);
            vars.resize(nalloc);
          }

        vDateTimes[tsID] = taxisInqVdatetime(taxisID1);

        fields_from_vlist(vlistID1, vars[tsID]);

        for (int recID = 0; recID < nrecs; ++recID)
          {
            cdo_inq_record(streamID1, &varID, &levelID);
            const auto &var = varList[varID];
            vars[tsID][varID][levelID].resize(var.gridsize);
            cdo_read_record(streamID1, vars[tsID][varID][levelID].vec_d.data(), &nmiss);
            vars[tsID][varID][levelID].nmiss = nmiss;

            if (nmiss) cdo_abort("Missing values are not allowed!");
          }

        tsID++;
      }

    const int nts = tsID;

    operator_input_arg("detrend type, length of segments, number of segments, window type\n\n"
                       "       detrend type: 0 - data should be used unchanged\n"
                       "                     1 - the mean of the whole time series should be subtracted\n"
                       "                     2 - the whole time series should be detrended\n"
                       "                     3 - every segment should be detrended\n\n"
                       "        window type: 0 - no data windowing\n"
                       "                     1 - Hann window\n"
                       "                     2 - Bartlett window\n"
                       "                     3 - Welch window\n");

    operator_check_argc(4);

    const auto detrend = parameter_to_int(cdo_operator_argv(0));
    const auto seg_l = parameter_to_int(cdo_operator_argv(1));
    const auto seg_n = parameter_to_int(cdo_operator_argv(2));
    const auto which_window = parameter_to_int(cdo_operator_argv(3));

    if (detrend < 0 || detrend > 3) cdo_abort("Illegal value for detrend (=%d)!", detrend);

    if (seg_l <= 2 || seg_l > nts) cdo_abort("Length must be at least 3 and at most the number of timesteps (=%d)", nts);

    if (seg_n <= 0 || seg_n > nts - seg_l + 1)
      cdo_abort("Number of segments must be positiv and not greater than %d!", nts - seg_l + 1);

    const int nfreq = seg_l / 2 + 1;

    FieldVector3D vars2(nfreq);
    for (freq = 0; freq < nfreq; freq++) fields_from_vlist(vlistID1, vars2[freq], FIELD_VEC);

    Varray<double> array1(nts);
    Varray<double> array2(nfreq);
    Varray<double> real(seg_l);
    Varray<double> imag(seg_l);
    Varray<double> window(seg_l);

    switch (which_window)
      {
      case 0:
        for (k = 0; k < seg_l; ++k) window[k] = 1;
        break;
      case 1:
        for (k = 0; k < seg_l; ++k) window[k] = 1 - std::cos(2 * M_PI * (k + 1) / (seg_l + 1));
        break;
      case 2:
        for (k = 0; k < seg_l / 2; ++k) window[k] = window[seg_l - 1 - k] = k;
        break;
      case 3:
        for (k = 0; k < seg_l; ++k)
          {
            double temp = ((k + 1.) - 0.5 * (seg_l + 1.)) / (0.5 * (seg_l + 1.));
            window[k] = 1 - temp * temp;
          }
        break;
      default: cdo_abort("Invalid window type %d!", which_window); break;
      }

    double wssum = 0;
    for (k = 0; k < seg_l; ++k) wssum += window[k] * window[k];

    for (varID = 0; varID < nvars; ++varID)
      {
        const auto &var = varList[varID];
        for (levelID = 0; levelID < var.nlevels; ++levelID)
          {
            for (size_t i = 0; i < var.gridsize; ++i)
              {
                for (tsID = 0; tsID < nts; ++tsID) array1[tsID] = vars[tsID][varID][levelID].vec_d[i];

                for (freq = 0; freq < nfreq; freq++) array2[freq] = 0;

                spectrum(nts, array1.data(), array2.data(), real.data(), imag.data(), window.data(), wssum, detrend, seg_n, seg_l);

                for (freq = 0; freq < nfreq; freq++) vars2[freq][varID][levelID].vec_d[i] = array2[freq];
              }
          }
      }

    for (tsID = 0; tsID < nfreq; ++tsID)
      {
        taxisDefVdatetime(taxisID2, vDateTimes[0]);
        cdo_def_timestep(streamID2, tsID);

        for (varID = 0; varID < nvars; ++varID)
          {
            const auto &var = varList[varID];
            for (levelID = 0; levelID < var.nlevels; ++levelID)
              {
                if (!vars2[tsID][varID][levelID].empty())
                  {
                    nmiss = vars2[tsID][varID][levelID].nmiss;
                    cdo_def_record(streamID2, varID, levelID);
                    cdo_write_record(streamID2, vars2[tsID][varID][levelID].vec_d.data(), 0);
                  }
              }
          }
      }
  }
  void
  close()
  {
    cdo_stream_close(streamID2);
    cdo_stream_close(streamID1);

    cdo_finish();
  }
};
void *
Spectrum(void *process)
{
  ModuleSpectrum spectrum;
  spectrum.init(process);
  spectrum.run();
  spectrum.close();

  return nullptr;
}
