/*
  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 "cdo_cdi_wrapper.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;
}

void *
Spectrum(void *process)
{
  int varID, levelID;
  int k;
  int nalloc = 0;
  size_t nmiss;
  int freq;

  cdo_initialize(process);

  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 = cdo_taxis_create(TAXIS_ABSOLUTE);
  vlistDefTaxis(vlistID2, taxisID2);

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

  VarList varList;
  varListInit(varList, vlistID1);

  const auto nvars = vlistNvars(vlistID1);
  FieldVector3D vars;
  std::vector<CdiDateTime> vDateTimes;

  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);
          vars[tsID][varID][levelID].resize(varList[varID].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 size_t gridsize = varList[varID].gridsize;
      const int nlevels = varList[varID].nlevels;
      for (levelID = 0; levelID < nlevels; ++levelID)
        {
          for (size_t i = 0; i < 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 nlevels = varList[varID].nlevels;
          for (levelID = 0; levelID < 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);
                }
            }
        }
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
