#include "complexflux.h"
#include "buffer.h"
#include "fft/pffft.h"
#include "remotebuffer.h"
#include <Eigen/Dense>
#include <catch2/catch.hpp>
#include <cmath>
#include <vector>

auto filterbankFrequencies(int bandsPerOctave, double fmin, double fmax,
                           double a) {
  const auto factor = std::pow(2.0, 1.0 / bandsPerOctave);
  auto freq = a;
  std::vector<double> res;
  res.reserve(bandsPerOctave * 10);
  while (freq <= fmax) {
    freq *= factor;
    res.push_back(freq);
  }
  freq = a;
  while (freq >= fmin) {
    freq /= factor;
    res.push_back(freq);
  }
  std::sort(res.begin(), res.end());
  return res;
}

static auto triangularFilter(int start, int mid, int stop, bool equal) {
  Eigen::VectorXf filter(stop - start);
  const auto height = equal ? (2.0 / (stop - start)) : 1.0;
  filter(Eigen::seq(0, mid - start - 1))
      .setEqualSpaced(0, height / (mid - start));
  filter(Eigen::seq(mid - start, stop - start - 1))
      .setEqualSpaced(height, -height / (stop - mid));
  return filter;
}

Eigen::MatrixXf filterbank(int numFFTBins, double sampleRate, int bands,
                           double fmin, double fmax, bool equal) {
  assert(numFFTBins > 0);
  assert(sampleRate > 0);

  fmax = std::min(sampleRate / 2.0, fmax);

  auto freqs = filterbankFrequencies(bands, fmin, fmax);
  const auto freqPerBin = (sampleRate / 2.0) / numFFTBins;
  std::vector<int> freqBins(freqs.size());
  std::transform(freqs.begin(), freqs.end(), freqBins.begin(),
                 [=](auto f) { return (int)std::round(f / freqPerBin); });
  freqBins.erase(std::unique(freqBins.begin(), freqBins.end()), freqBins.end());
  freqBins.erase(std::remove_if(freqBins.begin(), freqBins.end(),
                                [=](auto i) { return i >= numFFTBins; }),
                 freqBins.end());
  bands = freqBins.size() - 2;

  assert(bands >= 3);

  Eigen::MatrixXf res = Eigen::MatrixXf::Zero(numFFTBins, bands);
  for (int i = 0; i < bands; i++) {
    const auto start = freqBins[i];
    const auto mid = freqBins[i + 1];
    const auto stop = freqBins[i + 2];
    res(Eigen::seq(start, stop - 1), i) =
        triangularFilter(start, mid, stop, equal);
  }

  return res;
}

static Eigen::ArrayXf hanningWindow(int frameSize) {
  Eigen::ArrayXf res(frameSize);
  for (int i = 0; i < frameSize; i++) {
    res(i) =
        0.5 * (1.0 - std::cos((2.0 * M_PI / (double)frameSize) * (double)i));
  }
  return res;
}

static inline void unwrap_array(float *in, float *out, int len) {
  out[0] = in[0];
  for (int i = 1; i < len; i++) {
    float d = in[i] - in[i - 1];
    d = d > (float)M_PI ? d - 2.f * (float)M_PI
                        : (d < -(float)M_PI ? d + 2.f * (float)M_PI : d);
    out[i] = out[i - 1] + d;
  }
}

static inline float atan_scalar_approximation(float x) {
  float a1 = 0.99997726f;
  float a3 = -0.33262347f;
  float a5 = 0.19354346f;
  float a7 = -0.11643287f;
  float a9 = 0.05265332f;
  float a11 = -0.01172120f;

  float x_sq = x * x;
  return x *
         (a1 +
          x_sq * (a3 + x_sq * (a5 + x_sq * (a7 + x_sq * (a9 + x_sq * a11)))));
}

static void fastarg(const Eigen::ArrayXcf &in, Eigen::ArrayXf &out) {
  const float pi = M_PI;
  const float pi2 = M_PI / 2.0;
  auto inData = in.data();
  auto outData = out.data();
  int N = in.size();
#pragma clang loop vectorize(enable)
  for (int i = 0; i < N; i++) {
    const auto c = inData[i];
    float y = c.imag();
    float x = c.real();
    if (x == 0.f && y == 0.f) {
      // IEEE-754 compliance
      outData[i] = 0.f;
      continue;
    }
    bool swap = fabs(x) < fabs(y);
    float atan_input = (swap ? x : y) / (swap ? y : x);
    float res = atan_scalar_approximation(atan_input);

    res = swap ? (atan_input >= 0.0f ? pi2 : -pi2) - res : res;
    if (x < 0.0f) {
      res = (y >= 0.0f ? pi : -pi) + res;
    }

    outData[i] = res;
  }
}

SpectrogramResult spectrogram(std::shared_ptr<RandomAccessAudioReadable> audio,
                              int frameSize, double fps,
                              const Eigen::MatrixXf &filterbank, bool log,
                              float mul, float add, int blockSize,
                              bool useLgd) {
  assert(add > 0);
  assert(mul > 0);
  assert(audio->getChannelCount() > 0);

  SpectrogramResult res;

  res.hopSize = (double)audio->getSampleRate() / fps;
  int numFrames = std::ceil(audio->getFrameCount() / res.hopSize);
  int numFFTBins = int(frameSize / 2);
  int numBins = filterbank.cols();
  blockSize = std::min(blockSize, numFrames);

  res.spectrogram = Eigen::MatrixXf::Zero(numFrames, numBins);
  int block = 0;
  Eigen::MatrixXf specPrefilter = Eigen::MatrixXf::Zero(blockSize, numFFTBins);
  if (useLgd) {
    res.lgd = Eigen::MatrixXf::Zero(numFrames, numFFTBins);
  }

  res.window = hanningWindow(frameSize);

  BufferF32 audioFrame(audio->getChannelCount(), frameSize);
  audioFrame.fill(0.f);

  std::unique_ptr<PFFFT_Setup, void (*)(PFFFT_Setup *)> pffft_setup(
      pffft_new_setup(frameSize, PFFFT_REAL), pffft_destroy_setup);

  Eigen::ArrayXcf spectrum = Eigen::ArrayXcf::Zero(numFFTBins);
  Eigen::ArrayXf spectrumAngles = Eigen::ArrayXf::Zero(numFFTBins);
  Eigen::ArrayXf spectrumUnwrappedAngles = Eigen::ArrayXf::Zero(numFFTBins);
  Eigen::Map<Eigen::ArrayXf> audioFrameVector(audioFrame.getChannelData(0),
                                              audioFrame.getFrameCount());

  for (int frame = 0; frame < numFrames; frame++) {
    int seek = (double)frame * res.hopSize - (double)frameSize / 2.0;
    audio->readZeroPadded(seek, audioFrame);
    audioFrame.mixDownInPlace();

    audioFrameVector *= res.window;

    pffft_transform_ordered(pffft_setup.get(), audioFrameVector.data(),
                            reinterpret_cast<float *>(spectrum.data()), nullptr,
                            PFFFT_FORWARD);

    if (useLgd) {
      fastarg(spectrum, spectrumAngles);
      unwrap_array(spectrumAngles.data(), spectrumUnwrappedAngles.data(),
                   numFFTBins);
      res.lgd(frame, Eigen::seq(0, numFFTBins - 2)) =
          spectrumUnwrappedAngles(Eigen::seq(0, numFFTBins - 2)) -
          spectrumUnwrappedAngles(Eigen::seq(1, numFFTBins - 1));
    }

    specPrefilter.row(frame % blockSize) = spectrum.abs(); // 10%
    if ((frame + 1) / blockSize > block || (frame + 1) == numFrames) {
      auto start = block * blockSize;
      auto stop = std::min(start + blockSize, numFrames);
      res.spectrogram(Eigen::seq(start, stop - 1), Eigen::indexing::all) =
          specPrefilter(Eigen::seq(0, stop - start - 1), Eigen::indexing::all) *
          filterbank;
      block += 1;
    }
  }

  if (log) {
    res.spectrogram = (res.spectrogram.array() * mul + add).log().matrix();
  }

  return res;
}

static inline int reflect(int high, int val) {
  if (val < 0) {
    return -val;
  }
  if (val >= high) {
    return 2 * high - val - 1;
  }
  return val;
}

static Eigen::MatrixXf maximumFilter(const Eigen::MatrixXf &mat, int takeRows,
                                     int takeCols, int rowOrigin = 0,
                                     int columnOrigin = 0) {
  Eigen::MatrixXf res(mat.rows(), mat.cols());
  rowOrigin -= takeRows / 2;
  columnOrigin -= takeCols / 2;
  for (int mrow = 0; mrow < mat.rows(); mrow++) {
    for (int mcol = 0; mcol < mat.cols(); mcol++) {
      float maxVal = -INFINITY;
      int mkrowOrigin = mrow + rowOrigin;
      int mkcolOrigin = mcol + columnOrigin;
      for (int krow = 0; krow < takeRows; krow++) {
        for (int kcol = 0; kcol < takeCols; kcol++) {
          int mkrow = reflect(mat.rows(), krow + mkrowOrigin);
          int mkcol = reflect(mat.cols(), kcol + mkcolOrigin);
          float val = mat(mkrow, mkcol);
          if (val > maxVal) {
            maxVal = val;
          }
        }
      }
      res(mrow, mcol) = maxVal;
    }
  }
  return res;
}

static Eigen::MatrixXf superfluxDiffSpec(const SpectrogramResult &spectrogram,
                                         int diffFrames, int maxBins) {
  const auto &spec = spectrogram.spectrogram;
  Eigen::MatrixXf res = Eigen::MatrixXf::Zero(spec.rows(), spec.cols());
  auto maxSpec = maximumFilter(spec, 1, maxBins, 0, 0);
  res(Eigen::seq(diffFrames, spec.rows() - 1), Eigen::indexing::all) =
      spec(Eigen::seq(diffFrames, spec.rows() - 1), Eigen::indexing::all) -
      maxSpec(Eigen::seq(0, spec.rows() - diffFrames - 1),
              Eigen::indexing::all);
  res = res.cwiseMax(0.0);
  return res;
}

Eigen::ArrayXf superflux(const SpectrogramResult &spectrogram, float ratio) {
  int firstSampleLarger = 0;
  const auto &window = spectrogram.window;
  for (int i = 1; i < window.size(); i++) {
    if (window[i] > ratio) {
      firstSampleLarger = i;
      break;
    }
  }
  const auto diffSamples = (int)window.size() / 2 - firstSampleLarger;
  const int diffFrames = std::round(diffSamples / spectrogram.hopSize);

  auto diffSpec = superfluxDiffSpec(spectrogram, diffFrames, 3);

  return diffSpec.rowwise().sum();
}

static Eigen::MatrixXf lgdMask(const Eigen::MatrixXf &spec,
                               const Eigen::MatrixXcf &lgd,
                               const Eigen::MatrixXf &filterbank,
                               int temporalFilter, int temporalOrigin) {
  Eigen::MatrixXf lgdAbs = lgd.array().abs();
  if (temporalFilter > 0) {
    lgdAbs = maximumFilter(lgdAbs, temporalFilter, 1, temporalOrigin, 0);
  }

  Eigen::MatrixXf mask = Eigen::MatrixXf::Zero(spec.rows(), spec.cols());
  int numBins = lgd.cols();
  for (int b = 0; b < mask.cols(); b++) {
    int startBin = 0;
    int stopBin = filterbank.rows() - 1;
    while (filterbank(startBin, b) == 0.f && startBin < filterbank.rows()) {
      ++startBin;
    }
    while (filterbank(stopBin, b) == 0.f && stopBin >= 0) {
      --stopBin;
    }
    startBin = std::max(0, startBin - 1);
    stopBin = std::min(numBins, stopBin + 1);
    mask.col(b) =
        lgdAbs(Eigen::indexing::all, Eigen::seq(startBin, stopBin - 1))
            .rowwise()
            .minCoeff();
  }
  return mask / (float)M_PI;
}

Eigen::ArrayXf complexflux(const SpectrogramResult &spectrogram,
                           const Eigen::MatrixXf &filterbank, float ratio) {
  int firstSampleLarger = 0;
  const auto &window = spectrogram.window;
  for (int i = 1; i < window.size(); i++) {
    if (window[i] > ratio) {
      firstSampleLarger = i;
      break;
    }
  }
  const auto diffSamples = (int)window.size() / 2 - firstSampleLarger;
  const int diffFrames = std::round(diffSamples / spectrogram.hopSize);

  auto diffSpec = superfluxDiffSpec(spectrogram, diffFrames, 3);

  auto mask =
      lgdMask(spectrogram.spectrogram, spectrogram.lgd, filterbank, 3, 0);
  diffSpec.array() *= mask.array();

  return diffSpec.rowwise().sum();
}

TEST_CASE("complexflux doesn't crash", "[complexflux]") {
  auto testOnData = [](auto data) {
    auto fb = filterbank(1024, 44100);
    auto remotebuf = std::make_shared<BufferAsRemoteAudioBuffer>(44100, data);
    auto res =
        spectrogram(remotebuf, 2048, 200, fb, true, 1.f, 1.f, 2048, true);

    REQUIRE(res.spectrogram.rows() == 200);
    REQUIRE(!res.lgd.hasNaN());
    REQUIRE(!res.window.hasNaN());
    REQUIRE(!res.spectrogram.hasNaN());

    auto resc = complexflux(res, fb, 0.5);
    REQUIRE(!resc.hasNaN());
  };

  auto somedata = std::make_shared<BufferF32>(2, 44100);
  somedata->fill(0.f);
  testOnData(somedata);
  somedata->noise();
  testOnData(somedata);

  auto somemonodata = std::make_shared<BufferF32>(1, 44100);
  somemonodata->fill(0.f);
  testOnData(somemonodata);
}

TEST_CASE("maximumFilterInPlace", "[complexflux]") {
  Eigen::MatrixXf m(4, 4);
  m << 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16;

  {
    auto m1 = maximumFilter(m, 1, 1, 0, 0);
    REQUIRE(m1.isApprox(m));
  }

  {
    auto m2 = maximumFilter(m, 1, 3, 0, 0);
    Eigen::MatrixXf m2t(4, 4);
    m2t << 2, 3, 4, 4, 6, 7, 8, 8, 10, 11, 12, 12, 14, 15, 16, 16;
    REQUIRE(m2.isApprox(m2t));
  }

  {
    auto m3 = maximumFilter(m, 3, 1, 0, 0);
    Eigen::MatrixXf m3t(4, 4);
    m3t << 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16;
    REQUIRE(m3.isApprox(m3t));
  }

  {
    auto m4 = maximumFilter(m, 1, 1, 1, 2);
    Eigen::MatrixXf m4t(4, 4);
    m4t << 7, 8, 8, 7, 11, 12, 12, 11, 15, 16, 16, 15, 15, 16, 16, 15;
    REQUIRE(m4.isApprox(m4t));
  }
}