#include "meter.h"
#include "simd.h"
#include "util.h"
#include <catch2/catch.hpp>
#include <emscripten/atomic.h>
#include <emscripten/bind.h>
#include <emscripten/threading.h>
#include <wasm_simd128.h>

// accessed only from main runtime thread
static std::map<int,
                std::unique_ptr<PFFFT_Setup, decltype(&pffft_destroy_setup)>>
    pffft_setup_cache;

FFTObserver::FFTObserver(int channelCount, int binCount, int stride,
                         float decayAlpha)
    : stride(stride), enabled(0), inputQueue(channelCount, 2 * binCount),
      windowFunction(1, 2 * binCount), inputQueueWriteIndex(0),
      decayAlpha(decayAlpha), buffer(std::make_shared<BufferF32>(1, binCount)) {
  assert(decayAlpha >= 0.f && decayAlpha <= 1.f);
  assert(binCount % 8 == 0);
  assert(stride <= binCount);
  assert(emscripten_is_main_runtime_thread());

  auto it = pffft_setup_cache.find(2 * binCount);
  if (it == pffft_setup_cache.end()) {
    it = pffft_setup_cache
             .emplace(
                 2 * binCount,
                 std::unique_ptr<PFFFT_Setup, decltype(&pffft_destroy_setup)>(
                     pffft_new_setup(2 * binCount, PFFFT_REAL),
                     pffft_destroy_setup))
             .first;
  }
  setup = it->second.get();

  const auto N = (int)windowFunction.getFrameCount() - 1;
  auto windowData = windowFunction.getChannelData(0);

  // Hann window
  // weird indexing to be consistent with wikipedia
  for (int i = 0; i <= N; i++) {
    windowData[i] = 0.5 * (1.0 - std::cos((2.0 * M_PI * i) / N));
  }

  buffer->fill(0);
}

FFTObserver::~FFTObserver() {}

void FFTObserver::update(const BufferF32 &buf) {
  assert(buf.getChannelCount() == inputQueue.getChannelCount());

  bool enabled = emscripten_atomic_load_u32(&this->enabled) != 0;
  if (!enabled) {
    return;
  }

  int bufReadPosition = 0;

  while (true) {
    const auto bufFramesRemaining = (int)buf.getFrameCount() - bufReadPosition;
    if (bufFramesRemaining == 0) {
      break;
    }
    if (inputQueueWriteIndex == inputQueue.getFrameCount()) {
      doFrame();
    }
    int framesToCopy =
        std::min(bufFramesRemaining,
                 (int)inputQueue.getFrameCount() - inputQueueWriteIndex);

    auto sourceSlice =
        buf.slice(bufReadPosition, bufReadPosition + framesToCopy);
    auto destSlice = inputQueue.slice(inputQueueWriteIndex,
                                      inputQueueWriteIndex + framesToCopy);
    destSlice.set(0, sourceSlice);

    bufReadPosition += framesToCopy;
    inputQueueWriteIndex += framesToCopy;
  }
}

float FFTObserver::getDecayAlpha() const {
  return emscripten_atomic_load_f32(&this->decayAlpha);
}

void FFTObserver::setDecayAlpha(float decayAlpha) {
  assert(decayAlpha >= 0.f && decayAlpha <= 1.f);
  emscripten_atomic_store_f32(&this->decayAlpha, decayAlpha);
}

bool FFTObserver::getEnabled() const {
  return emscripten_atomic_load_u32(&this->enabled) != 0;
}

void FFTObserver::setEnabled(bool enabled) {
  emscripten_atomic_store_u32(&this->enabled, enabled ? 1 : 0);
}

void FFTObserver::doFrame() {
  // mix down to mono and apply window function in VLA
  float inputData[inputQueue.getFrameCount()];
  memset(inputData, 0, inputQueue.getFrameCount() * sizeof(float));

  auto windowData = windowFunction.getChannelData(0);
  float mixdownScale = 1.0f / inputQueue.getChannelCount();

  for (int channel = 0; channel < inputQueue.getChannelCount(); channel++) {
    auto queueData = inputQueue.getChannelData(channel);
#pragma clang loop vectorize(enable)
    for (int i = 0; i < inputQueue.getFrameCount(); i++) {
      inputData[i] += queueData[i] * windowData[i] * mixdownScale;
    }
  }

  // fft
  pffft_transform_ordered(setup, inputData, inputData, nullptr, PFFFT_FORWARD);

  // convert to magnitude, decay output buffer, and apply max(decayed output,
  // this output)
  auto outputData = buffer->getChannelData(0);

  const v128_t decayAlphaV =
      wasm_f32x4_splat(emscripten_atomic_load_f32(&decayAlpha));
  const v128_t scalingV = wasm_f32x4_splat(1.0f / inputQueue.getFrameCount());
  const int outputFrameCount = buffer->getFrameCount();
  for (int i = 0; i < outputFrameCount; i += 4) {
    v128_t outputMagnitude = wasm_v128_load(outputData + i);
    v128_t inter0 = wasm_v128_load(inputData + i * 2);
    v128_t inter1 = wasm_v128_load(inputData + i * 2 + 4);
    v128_t re = wasm_i32x4_shuffle(inter0, inter1, 0, 2, 4, 6);
    v128_t im = wasm_i32x4_shuffle(inter0, inter1, 1, 3, 5, 7);
    v128_t newMagnitude = wasm_f32x4_sqrt(
        wasm_f32x4_mul(scalingV, wasm_f32x4_add(wasm_f32x4_mul(re, re),
                                                wasm_f32x4_mul(im, im))));

    v128_t decayedOutputMagnitude =
        wasm_f32x4_mul(outputMagnitude, decayAlphaV);
    v128_t maxMagnitude = wasm_f32x4_max(decayedOutputMagnitude, newMagnitude);
    wasm_v128_store(outputData + i, maxMagnitude);
  }

// move inputQueue content left by stride
#pragma clang loop vectorize(enable)
  for (int channel = 0; channel < inputQueue.getChannelCount(); channel++) {
    auto queueData = inputQueue.getChannelData(channel);
    for (int i = 0; i < inputQueue.getFrameCount() - stride; i++) {
      queueData[i] = queueData[i + stride];
    }
  }

  inputQueueWriteIndex -= stride;
}

Meter::Meter(int channelCount, int sampleRate, float ppmHoldUpdateHz,
             int ppmHoldKeepPeriods, float vuAlpha)
    : channelCount(channelCount), sampleRate(sampleRate),
      ppmHoldUpdateHz(ppmHoldUpdateHz), ppmHoldKeepPeriods(ppmHoldKeepPeriods),
      vuAlpha(vuAlpha), ppmHoldValues(channelCount, ppmHoldKeepPeriods),
      ppmHoldValuesWriteIndex(0),
      ppmHoldBuffer(channelCount, sampleRate / ppmHoldUpdateHz),
      ppmHoldBufferRemaining(ppmHoldBuffer), ppmHoldIndicated(-200.f),
      ppmIndicated(-200.f), vuIndicated(-200.f),
      fftObserver(channelCount, 1024, 256, 0.95f) {
  assert(ppmHoldKeepPeriods > 0);
  assert(ppmHoldUpdateHz > 0);
  assert(sampleRate > 0);
  assert(vuAlpha >= 0.f && vuAlpha <= 1.f);
}

void Meter::update(const BufferF32 &buf) {
  assert(buf.getChannelCount() == channelCount);

  fftObserver.update(buf);

  float bufPeaks[channelCount];
  for (int channel = 0; channel < channelCount; channel++) {
    bufPeaks[channel] = buf.sliceChannel(channel).peak();
    ppmIndicated[channel] = signalValueToDbFS(bufPeaks[channel]);
  }

  int copyToPpmOffset = 0;
  while (copyToPpmOffset < buf.getFrameCount()) {
    const auto copySize = std::min(buf.getFrameCount() - copyToPpmOffset,
                                   ppmHoldBufferRemaining.getFrameCount());
    ppmHoldBufferRemaining.set(
        0, buf.slice(copyToPpmOffset, copyToPpmOffset + copySize));
    copyToPpmOffset += copySize;

    if (copySize == ppmHoldBufferRemaining.getFrameCount()) {
      // Buffer is full, copy to ppmHoldValues
      for (int channel = 0; channel < channelCount; channel++) {
        ppmHoldValues[channel][ppmHoldValuesWriteIndex] =
            ppmHoldBuffer.sliceChannel(channel).peak();
      }
      ppmHoldValuesWriteIndex =
          (ppmHoldValuesWriteIndex + 1) % ppmHoldValues.getFrameCount();
      ppmHoldBufferRemaining = ppmHoldBuffer;
    } else {
      ppmHoldBufferRemaining = ppmHoldBufferRemaining.slice(copySize);
      for (int channel = 0; channel < channelCount; channel++) {
        auto &ppmHoldValue = ppmHoldValues[channel][ppmHoldValuesWriteIndex];
        ppmHoldValue = std::max(ppmHoldValue, bufPeaks[channel]);
      }
    }
  }

  for (int channel = 0; channel < channelCount; channel++) {
    ppmHoldIndicated[channel] =
        signalValueToDbFS(ppmHoldValues.sliceChannel(channel).peak());

    const auto samples = buf.getChannelData(channel);
    const auto newVuValue = vfast_ema_squared_array(
        samples, vuAlpha, lastVuValue[channel], buf.getFrameCount());
    lastVuValue[channel] = newVuValue;
    vuIndicated[channel] = signalSquareValueToDbFS(newVuValue);
  }
}

float Meter::getFFTObserverDecayAlpha() const {
  return fftObserver.getDecayAlpha();
}

void Meter::setFFTObserverDecayAlpha(float decayAlpha) {
  fftObserver.setDecayAlpha(decayAlpha);
}

bool Meter::getFFTObserverEnabled() const { return fftObserver.getEnabled(); }

void Meter::setFFTObserverEnabled(bool enabled) {
  fftObserver.setEnabled(enabled);
}

std::shared_ptr<BufferF32> Meter::getFFTObserverBuffer() const {
  return fftObserver.buffer;
}

EMSCRIPTEN_BINDINGS(meter) {
  using namespace emscripten;

  emscripten::class_<Meter>("Meter")
      .smart_ptr<std::shared_ptr<Meter>>("Meter")
      .function("getPpmHoldIndicated(channel)", &Meter::getPpmHoldIndicated)
      .function("getPpmIndicated(channel)", &Meter::getPpmIndicated)
      .function("getVuIndicated(channel)", &Meter::getVuIndicated)
      .property("fftObserverEnabled", &Meter::getFFTObserverEnabled,
                &Meter::setFFTObserverEnabled)
      .property("fftObserverDecayAlpha", &Meter::getFFTObserverDecayAlpha,
                &Meter::setFFTObserverDecayAlpha)
      .function("getFFTObserverBuffer", &Meter::getFFTObserverBuffer,
                nonnull<ret_val>());
}

TEST_CASE("Meter memory tests", "[meter]") {
  Meter meter(1, 44100, 10, 10, 0.99f);
  BufferF32 buf(1, 128);
  for (int i = 0; i < 500; i++) {
    buf.noise();
    meter.update(buf);
  }

  FFTObserver fftObserver(1, 1024, 256, 0.95f);
  for (int i = 0; i < 100; i++) {
    buf.noise();
    fftObserver.update(buf);
  }
}