// #define DEBUGTRACE_ENABLED
#include "lasp_freqsmooth.h"

#include <cassert>

#include "debugtrace.hpp"

using rte = std::runtime_error;

vd freqSmooth(const vd& freq, const vd& X, const unsigned w,
              bool power_correct) {
  DEBUGTRACE_ENTER;
  if (freq.size() < 2) {
    throw rte("Invalid frequency size. Should be > 2");
  }
  if (freq.size() != X.size()) {
    throw rte("Sizes of freq and X do not match");
  }
  if (freq.size() > std::numeric_limits<long>::max() / 2) {
    throw rte("Frequency size limit for smoothing is 2^30");
  }
  if (w == 0) {
    throw rte("Invalid number of octaves");
  }
  const us Nfreq = freq.size();

  // Smoothing width in unit of number of octaves
  const d Delta = 1 / d(w);

  // Minimum frequency and maximum frequency to smooth on (frequency range that
  // is interpolated to a log scale)
  d freq_min;
  const d freq_max = freq(Nfreq - 1);
  const bool firstFreqEqZero = (d_abs(freq(0)) < 1e-15);

  // AC-signal power
  d ac_pwr;
  if (firstFreqEqZero) {
    freq_min = freq(1);
    if (power_correct) {
      ac_pwr = arma::sum(X.subvec(1, Nfreq - 1));
    }
  } else {
    freq_min = freq(0);
    if (power_correct) {
      ac_pwr = arma::sum(X);
    }
  }
  DEBUGTRACE_PRINT(freq_min);
  DEBUGTRACE_PRINT(freq_max);
  const vd freq_log =
      arma::logspace(d_log10(freq_min), d_log10(freq_max), 10 * Nfreq);
  DEBUGTRACE_PRINT("freq_log = ");

  const long Nfreq_sm = freq_log.size();

  // Interpolate X to logscale
  vd X_log;
  DEBUGTRACE_PRINT("X_log = :");
  arma::interp1(freq, X, freq_log, X_log, "*linear");

  // First and last point are not interpolated well, could be minimally out of
  // the interpolation range, due to roundoff errors. Armadillo sets these
  // points to nan, so we have to manually "interpolate" them.
  X_log(Nfreq_sm - 1) = X(X.size() - 1);
  if (firstFreqEqZero) {
    X_log(0) = X(1);
  } else {
    X_log(0) = X(0);
  }

  // Allocate space for smoothed X on log scale
  vd Xsm_log(freq_log.size());
  const d beta = d_log10(Nfreq_sm) / d_log10(2) / (Nfreq_sm - 1);
  // int rounds down
  const long mu = int(Delta / d(2) / beta);
  DEBUGTRACE_PRINT(mu);

  // Long is at least 32 bits. So +/- 2M points length
  for (long k = 0; k < Nfreq_sm; k++) {
    // const d fcur = freq_log(k);
    long idx_start = std::max(k - mu, 0l);
    long idx_stop = std::min(k + mu, Nfreq_sm - 1);

    // Make window smaller at the sides (close to the end of the array)
    if (idx_start == 0 || idx_stop == Nfreq_sm - 1) {
      const long mu_edge = std::min(k - idx_start, idx_stop - k);
      idx_start = k - mu_edge;
      idx_stop = k + mu_edge;
    }
    assert(idx_stop < Nfreq_sm);
    assert(idx_start < Nfreq_sm);

    DEBUGTRACE_PRINT(idx_start)
    DEBUGTRACE_PRINT(idx_stop);

    Xsm_log(k) = arma::mean(X_log.subvec(idx_start, idx_stop));
  }
  DEBUGTRACE_PRINT("Xsm_log:");
  //   std::cerr << Xsm_log << std::endl;

  // Back-interpolate to a linear scale, and be wary of nans at the start end
  // and range. Also interpolates power
  vd Xsm(Nfreq);
  if (firstFreqEqZero) {
    vd Xsm_gt0;
    arma::interp1(freq_log, Xsm_log, freq.subvec(1, Nfreq - 1), Xsm_gt0,
                  "*linear");
    Xsm(0) = X(0);
    Xsm.subvec(1, Nfreq - 1) = Xsm_gt0;
    Xsm(1) = Xsm_log(1);
    Xsm(Nfreq - 1) = Xsm_log(Nfreq_sm - 1);

    // Final step: power-correct smoothed spectrum
    if (power_correct) {
      d new_acpwr = arma::sum(Xsm.subvec(1, Nfreq - 1));
      Xsm.subvec(1, Nfreq - 1) *= ac_pwr / new_acpwr;
    }

  } else {
    arma::interp1(freq_log, Xsm_log, freq, Xsm, "*linear");
    Xsm(0) = X(0);
    Xsm(Nfreq - 1) = Xsm_log(Nfreq_sm - 1);

    // Final step: power-correct smoothed spectrum
    if (power_correct) {
      d new_acpwr = arma::sum(Xsm);
      Xsm *= ac_pwr / new_acpwr;
    }
  }

  return Xsm;
}