#include <optional>
#define DEBUGTRACE_ENABLED
#include "debugtrace.hpp"
#include "lasp_avpowerspectra.h"
#include <cmath>

using std::runtime_error;

PowerSpectra::PowerSpectra(const us nfft, const Window::WindowType w)
    : PowerSpectra(Window::create(w, nfft)) {}

PowerSpectra::PowerSpectra(const vd &window)
    : nfft(window.size()), _fft(nfft), _window(window) {

  d win_pow = arma::sum(window % window);

  /* Scale fft such that power is easily computed */
  _scale_fac = sqrt(2 / win_pow) / nfft;
}

arma::Cube<c> PowerSpectra::compute(const dmat &input) {

  dmat input_tmp = input;

  // Multiply each column of the inputs element-wise with the window.
  input_tmp.each_col() %= _window;

  cmat rfft = _fft.fft(input_tmp) * _scale_fac;

  arma::cx_cube output(rfft.n_rows, input.n_cols, input.n_cols);
  for (us i = 0; i < input.n_cols; i++) {
    for (us j = 0; j < input.n_cols; j++) {
      output.slice(j).col(i) = rfft.col(i) % arma::conj(rfft.col(j));
    }
  }
  return output;
}

AvPowerSpectra::AvPowerSpectra(const us nfft, const Window::WindowType w,
                               const d overlap_percentage,
                               const int time_constant)
    : _ps(nfft, w) {

  DEBUGTRACE_ENTER;
  if (overlap_percentage >= 100 || overlap_percentage < 0) {
    throw runtime_error("Overlap percentage should be >= 0 and < 100");
  }

  _overlap_keep = (nfft * overlap_percentage) / 100;
  if (_overlap_keep >= nfft) {
    throw runtime_error("Overlap is too high. Results in no jump. Please "
                        "choose a smaller overlap percentage or a higher nfft");
  }
  if (time_constant < 0) {
    _mode = Mode::Averaging;
  } else if (time_constant == 0) {
    _mode = Mode::Spectrogram;
  } else {
    _mode = Mode::Leaking;
    _alpha = exp(-static_cast<d>(nfft - _overlap_keep) / time_constant);
  }
}

std::optional<arma::cx_cube> AvPowerSpectra::compute(const dmat &timedata) {

  _timeBuf.push(timedata);
  std::optional<arma::cx_cube> res;

  while (auto samples = _timeBuf.pop(_ps.nfft, _overlap_keep)) {
    switch (_mode) {
    case (Mode::Spectrogram): {
      res.emplace(_ps.compute(samples.value()));
    } break;
    case (Mode::Averaging): {
      n_averages++;
      if (n_averages == 1) {
        _est = _ps.compute(samples.value());
      } else {
        _est = _est * (n_averages - 1) / n_averages +
               _ps.compute(samples.value()) / n_averages;
      }
      res = _est;
    } break;
    case (Mode::Leaking): {
      if (arma::size(_est) == arma::size(0, 0, 0)) {
        _est = _ps.compute(samples.value());
      } else {
        _est = _alpha * _est + (1 - _alpha) * _ps.compute(samples.value());
      }
      res = _est;
    } break;
    } // end switch mode
  }
  return res;
}
std::optional<arma::cx_cube> AvPowerSpectra::get_est() {
  if (_est.n_cols > 0)
    return _est;
  return std::nullopt;
}