/* #define DEBUGTRACE_ENABLED */
#include <memory>
#include <cassert>
#include "debugtrace.hpp"
#include "lasp_config.h"
#include "lasp_fft.h"
using rte = std::runtime_error;

#if LASP_FFT_BACKEND == Armadillo
class Fft_impl {
public:
  us nfft;
  Fft_impl(const us nfft) : nfft(nfft) { DEBUGTRACE_ENTER; }

  vc fft(const vd &time) {
    DEBUGTRACE_ENTER;
    vc res = arma::fft(time);
    return res.rows(0, nfft / 2);
  }
  vd ifft(const vc &f) {
    DEBUGTRACE_ENTER;
    vc f_full(nfft);
    vc f_above_nyq = arma::conj(arma::reverse(f));
    f_full.rows(0, nfft / 2) = f;
    f_full.rows(nfft / 2, nfft - 1) = f_above_nyq.rows(0, nfft/2-1);
    return arma::real(arma::ifft(f_full));
  }
};
#elif LASP_FFT_BACKEND == FFTW
#include <fftw3.h>

/**
 * @brief FFTW implementation
 */
class Fft_impl {
public:
  us nfft;
  fftw_plan forward_plan = nullptr;
  fftw_plan reverse_plan = nullptr;
  fftw_complex *frequencyDomain = nullptr;
  d *timeDomain = nullptr;

  Fft_impl(const us nfft) : nfft(nfft) {

    timeDomain = (d *)fftw_malloc(sizeof(d) * nfft);

    frequencyDomain =
        (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * (nfft / 2 + 1));

    forward_plan =
        fftw_plan_dft_r2c_1d(nfft, timeDomain, frequencyDomain, FFTW_MEASURE);

    reverse_plan =
        fftw_plan_dft_c2r_1d(nfft, frequencyDomain, timeDomain, FFTW_MEASURE);
    if (!forward_plan || !reverse_plan || !timeDomain || !frequencyDomain) {
      throw rte("Error allocating FFT");
    }
  };

  ~Fft_impl() {
    fftw_destroy_plan(forward_plan);
    fftw_destroy_plan(reverse_plan);
    fftw_free(frequencyDomain);
    fftw_free(timeDomain);
  }

  vc fft(const vd &time) {
    if (time.n_rows != nfft) {
      throw rte("Invalid size of input vector, should be equal to nfft");
    }

    vc res(nfft / 2 + 1);
    memcpy(timeDomain, time.memptr(), sizeof(d) * nfft);

    fftw_execute(forward_plan);

    memcpy(reinterpret_cast<void *>(res.memptr()), frequencyDomain,
           sizeof(c) * (nfft / 2 + 1));

    return res;
  }
  vd ifft(const vc &f) {
    DEBUGTRACE_ENTER;
    if (f.n_rows != nfft / 2 + 1) {
      throw rte("Invalid size of input vector, should be equal to nfft/2+1");
    }
    memcpy(frequencyDomain,
           reinterpret_cast<void *>(const_cast<c *>(f.memptr())),
           (nfft / 2 + 1) * sizeof(c));

    fftw_execute(reverse_plan);

    vd res(nfft);
    memcpy(res.memptr(), timeDomain, nfft * sizeof(d));

    /* Scale by dividing by nfft. Checked with numpy implementation
     * that this indeed needs to be done for FFTW. */
    res *= 1.0 / nfft;

    return res;
  }
};
#else
#error                                                                         \
    "Cannot compile lasp_ffc.c, no FFT backend specified. Should either be FFTPack, or FFTW"
#endif

Fft::Fft(const us nfft) {
  DEBUGTRACE_ENTER;
  if (nfft == 0) {
    throw rte("Invalid nfft: 0");
  }
  if (nfft >= LASP_MAX_NFFT) {
    throw rte("Invalid nfft, should be smaller than: " +
              std::to_string(LASP_MAX_NFFT));
  }
  _impl = std::make_unique<Fft_impl>(nfft);
}
Fft::~Fft() {}

us Fft::nfft() const {
  assert(_impl);
  return _impl->nfft;
}

vc Fft::fft(const vd &timedata) {
  assert(_impl);
  return _impl->fft(timedata);
}

cmat Fft::fft(const dmat &freqdata) {
  DEBUGTRACE_ENTER;
  assert(_impl);
  cmat res(_impl->nfft / 2 + 1, freqdata.n_cols);
  /// * WARNING *. This was source of a serious bug. It is not possible to run
  /// FFT's and IFFT's on the same _impl, as it overwrites the same memory.
  /// Uncommenting the line below results in faulty results.
  // #pragma omp parallel for
  for (us colno = 0; colno < freqdata.n_cols; colno++) {
    res.col(colno) = _impl->fft(freqdata.col(colno));
  }
  return res;
}

vd Fft::ifft(const vc &freqdata) {
  DEBUGTRACE_ENTER;
  assert(_impl);
  return _impl->ifft(freqdata);
}
dmat Fft::ifft(const cmat &freqdata) {
  dmat res(_impl->nfft, freqdata.n_cols);
  /// * WARNING *. This was source of a serious bug. It is not possible to run
  /// FFT's and IFFT's on the same _impl, as it overwrites the same memory.
  /// Uncommenting the line below results in faulty results.
  // #pragma omp parallel for

  for (us colno = 0; colno < freqdata.n_cols; colno++) {
    res.col(colno) = _impl->ifft(freqdata.col(colno));
  }
  return res;
}

void Fft::load_fft_wisdom(const std::string &wisdom) {
#if LASP_FFT_BACKEND == Armadillo
#elif LASP_FFT_BACKEND == FFTW
  if (wisdom.length() > 0) {
    int rv = fftw_import_wisdom_from_string(wisdom.c_str());
    if (rv != 1) {
      throw rte("Error loading FFTW wisdom");
    }
  }
#endif
}

std::string Fft::store_fft_wisdom() {
#if LASP_FFT_BACKEND == Armadillo
  return "";
#elif LASP_FFT_BACKEND == FFTW
  // It is not possible to let FFTW directly return a C++ string. We have to
  // put it in this container by copying in. Not a good solution if this has to
  // happen often.
  // Fortunately, this function is only called at the end of the program.
  char *wis = fftw_export_wisdom_to_string();
  std::string res{wis};
  free(wis);
  return res;
#endif
}