#include "timestretch.h"
#include "audioclip.h"
#include "remotebuffer.h"
#include "rendercontext.h"
#include "ringbuffercache.h"
#include "timing/operators.h"
#include <algorithm>
#include <bungee/Bungee.h>
#include <catch2/catch.hpp>
#include <cmath>
#include <emscripten/atomic.h>
#include <emscripten/bind.h>

using namespace emscripten;
using namespace Bungee;

template <typename ImplType, int prerollT, int blockDelayT>
struct BungeeStretcherVirtualWrap : public BungeeStretcherBase {
  BungeeStretcherVirtualWrap(Bungee::SampleRates sampleRates, int channelCount,
                             int log2SynthesisHopOverride = 0)
      : BungeeStretcherBase(prerollT, blockDelayT),
        stretcher(sampleRates, channelCount, log2SynthesisHopOverride) {}

  Bungee::InputChunk specifyGrain(const Bungee::Request &request) {
    return stretcher.specifyGrain(request);
  }
  void analyseGrain(const float *data, intptr_t channelStride) {
    stretcher.analyseGrain(data, channelStride);
  }
  void synthesiseGrain(Bungee::OutputChunk &outputChunk) {
    stretcher.synthesiseGrain(outputChunk);
  }

private:
  Stretcher<ImplType> stretcher;
};

using BungeeStretcherBasic = BungeeStretcherVirtualWrap<Bungee::Basic, 3, 2>;
using BungeeStretcherPro = BungeeStretcherVirtualWrap<Bungee::Pro, 3, 3>;

TimestretchReader::TimestretchReader(
    std::shared_ptr<RandomAccessAudioReadable> underlyingBuffer,
    std::shared_ptr<time_transform::WarpMap<double>> warpMap,
    int playbackSampleRate, bool useProAlgorithm)
    : TimestretchReader(underlyingBuffer->getChannelCount(),
                        underlyingBuffer->getSampleRate(), playbackSampleRate,
                        useProAlgorithm) {
  assert(underlyingBuffer != nullptr);
  assert(warpMap != nullptr);

  this->underlyingBuffer->setUnderlyingReadable(std::move(underlyingBuffer));
  this->warpMap = std::move(warpMap);
}

TimestretchReader::TimestretchReader(int channelCount, int underlyingSampleRate,
                                     int playbackSampleRate,
                                     bool useProAlgorithm)
    : underlyingBuffer(
          std::make_shared<RingBufferCache>(channelCount, 2048, 20)),
      underlyingSampleRate(underlyingSampleRate),
      playbackSampleRate(playbackSampleRate),
      synthesisHopSize(std::pow(2,
                                std::floor(std::log2(underlyingSampleRate)) -
                                    6)), // see Bungee Timing.cpp
      grainSize(synthesisHopSize * 8), state(State::Reset),
      currentPosition(time_units::Beats<double>(0.0)),
      lastLocalPosition(
          time_units::Beats<double>(std::numeric_limits<double>::lowest())),
      // You may assume excessOutput is large enough to hold Bungee's output.
      excessOutput(
          channelCount,
          (1 << ((int)std::ceil(-AudioClip::TRANSPOSE_MINIMUM_SEMITONES / 12) +
                 1)) *
              synthesisHopSize),
      excessOutputOffset(0), excessOutputCount(0), transposition(0.0f),
      underlyingChunk(channelCount, grainSize),
      lastAnalysisPosition(time_units::Seconds<double>(0.0)) {
  assert(underlyingSampleRate > 0);
  assert(playbackSampleRate > 0);

  const auto sampleRates =
      SampleRates{underlyingSampleRate, playbackSampleRate};

  if (useProAlgorithm) {
    stretcher = std::make_unique<BungeeStretcherPro>(sampleRates, channelCount);
  } else {
    stretcher =
        std::make_unique<BungeeStretcherBasic>(sampleRates, channelCount);
  }

  calculatePitchCoefAndHopSize();
}

TimestretchReader::~TimestretchReader() {}

void TimestretchReader::setUnderlying(
    std::shared_ptr<RandomAccessAudioReadable> underlyingBuffer,
    std::shared_ptr<time_transform::WarpMap<double>> warpMap) {
  assert(underlyingBuffer != nullptr);
  assert(warpMap != nullptr);

  assert(underlyingBuffer->getSampleRate() == underlyingSampleRate);
  this->underlyingBuffer->setUnderlyingReadable(std::move(underlyingBuffer));
  this->warpMap = std::move(warpMap);
  state = State::Reset;
}

bool TimestretchReader::readSegment(const RenderContext *renderContext,
                                    time_units::Beats<double> localPosition,
                                    float incomingTransposition,
                                    BufferF32 output) {
  if (incomingTransposition != transposition) {
    transposition = incomingTransposition;
    calculatePitchCoefAndHopSize();
    state = State::Reset;
  }

  int framesToProduce = output.getFrameCount();
  int framesProduced = 0;
  const double dBeatPerFrame = renderContext->bps / (double)playbackSampleRate;

  const auto isSeek = state == State::Reset || !renderContext->isContinuous ||
                      (localPosition != lastLocalPosition);

  if (isSeek) {
    reset(renderContext, localPosition);
  }

  while (framesProduced < framesToProduce) {
    // Take from excessOutput if available
    if (excessOutputCount > 0) {
      auto framesToCopy =
          std::min(framesToProduce - framesProduced, excessOutputCount);
      output.set(framesProduced,
                 excessOutput.slice(excessOutputOffset,
                                    excessOutputOffset + framesToCopy));
      excessOutputOffset += framesToCopy;
      excessOutputCount -= framesToCopy;
      framesProduced += framesToCopy;
      continue;
    }

    // Synthesize a grain - this will fill excessOutput
    synthesize(renderContext, localPosition, false);
  }

  lastLocalPosition =
      localPosition +
      time_units::BeatsDelta(dBeatPerFrame * (double)framesProduced);

  return isSeek;
}

void TimestretchReader::reset(const RenderContext *renderContext,
                              time_units::Beats<double> localPosition) {
  state = State::ContinuousPlayback;

  const auto preroll = stretcher->preroll;
  const auto blockDelay = stretcher->blockDelay;

  const auto warmupSamples = preroll * resamplingSynthesisHopSize;
  const auto warmupSeconds = warmupSamples / (double)playbackSampleRate;

  // Find global beat position delaySeconds before render position
  currentPosition = findDomainPointForCodomainDelta(
      renderContext->globalBeatsToSeconds, renderContext->position,
      time_units::SecondsDelta(-warmupSeconds));

  // adjust run-in for variable rate resampling used by bungee
  int framesRendered = 0;

  const auto framesExpected =
      resamplingSynthesisHopSize * (preroll + blockDelay + 1);

  synthesize(renderContext, localPosition, true);
  framesRendered += excessOutputCount;
  while (framesExpected - framesRendered >= resamplingSynthesisHopSize) {
    synthesize(renderContext, localPosition, false);
    framesRendered += excessOutputCount;
  }

  if (framesExpected - framesRendered < excessOutputCount) {
    excessOutputOffset += framesExpected - framesRendered;
    excessOutputCount -= framesExpected - framesRendered;
  } else {
    synthesize(renderContext, localPosition, false);
    framesRendered += excessOutputCount;
  }
}

void TimestretchReader::synthesize(const RenderContext *renderContext,
                                   time_units::Beats<double> localPosition,
                                   bool reset) {
  const auto globalToLocal = localPosition - renderContext->position;
  const auto localSynthesisPosition = currentPosition + globalToLocal;
  const auto analysisGrainSegment =
      warpMap->getSegmentIteratorAt(localSynthesisPosition);
  assert(analysisGrainSegment != warpMap->end());
  const auto analysisGrainSeconds =
      analysisGrainSegment->map_point(localSynthesisPosition);
  assert(analysisGrainSeconds.has_value());
  const auto analysisGrainSamples =
      analysisGrainSeconds->raw() * (double)underlyingBuffer->getSampleRate();

  lastAnalysisPosition = *analysisGrainSeconds;

  const auto analysisGrain = Request{.position = analysisGrainSamples,
                                     .speed = 1.0, // this is ignored by bungee
                                     .pitch = pitchFreqCoef,
                                     .reset = reset};
  const auto inChunk = stretcher->specifyGrain(analysisGrain);

  const auto underlyingSlice =
      underlyingChunk.slice(0, inChunk.end - inChunk.begin);
  underlyingBuffer->readZeroPadded(inChunk.begin, underlyingSlice);

  if (underlyingSlice.getChannelCount() < 2) {
    stretcher->analyseGrain(underlyingSlice.getChannelData(0), 0);
  } else {
    stretcher->analyseGrain(underlyingSlice.getChannelData(0),
                            underlyingSlice.getChannelData(1) -
                                underlyingSlice.getChannelData(0));
  }

  OutputChunk output{};
  stretcher->synthesiseGrain(output);

  for (int channel = 0; channel < underlyingSlice.getChannelCount();
       channel++) {
    excessOutput.set(channel, 0, output.frameCount,
                     output.data + channel * output.channelStride);
  }
  excessOutputOffset = 0;
  excessOutputCount = output.frameCount;

  const auto synthesizedSeconds =
      (double)output.frameCount / (double)playbackSampleRate;
  currentPosition = findDomainPointForCodomainDelta(
      renderContext->globalBeatsToSeconds, currentPosition,
      time_units::SecondsDelta(synthesizedSeconds));
}

void TimestretchReader::calculatePitchCoefAndHopSize() {
  pitchFreqCoef = std::pow(2.0, transposition / 12.0);

  resamplingSynthesisHopSize =
      (double)synthesisHopSize *
      ((double)playbackSampleRate / (double)underlyingSampleRate) /
      pitchFreqCoef;
}

time_units::Seconds<double> TimestretchReader::getLastAnalysisPosition() const {
  return lastAnalysisPosition;
}

TEST_CASE("timestretch tests", "[timestretch]") {
  auto testWithParams = [=](bool pro, int channelCount, int inputSampleRate,
                            int outputSampleRate, float transposition,
                            double stretch) {
    using namespace time_units;
    using namespace time_transform;
    auto inputBuf = std::make_shared<BufferF32>(channelCount, 20000);
    inputBuf->noise();
    BufferF32 outBuf(channelCount, 1024);
    auto inputRd =
        std::make_shared<BufferAsRemoteAudioBuffer>(inputSampleRate, inputBuf);
    auto warpMap = std::make_shared<time_transform::WarpMap<double>>(
        std::vector<WarpMap<double>::WarpMarker>{
            {Beats(0.0), Seconds(0.0)}, {Beats(stretch), Seconds(1.0)}});
    auto reader = TimestretchReader(inputRd, warpMap, outputSampleRate, pro);
    TimelineTempoMap tm(2.0, {});
    for (Beats i(0.0); i < Beats(1.0); i += BeatsDelta(0.1)) {
      RenderContext rc{i, 2.0, true, tm, nullptr, BeatsDelta(0.1)};
      reader.readSegment(&rc, i, transposition, outBuf);
    }
  };

  auto alignmentTestWithParams = [=](bool pro, int channelCount,
                                     int inputSampleRate, int outputSampleRate,
                                     float transposition, double stretch) {
    using namespace time_units;
    using namespace time_transform;

    const int peakPosition = 23232;
    auto inputBuf =
        std::make_shared<BufferF32>(channelCount, 4 * inputSampleRate);
    auto concatOutBuf = BufferF32(channelCount, 4 * outputSampleRate);
    inputBuf->fill(0.0f);
    concatOutBuf.fill(0.0f);
    for (int channel = 0; channel < inputBuf->getChannelCount(); channel++) {
      inputBuf->getChannelData(channel)[peakPosition] = 0.8f;
    }
    BufferF32 outBuf(channelCount, 1024);
    auto inputRd =
        std::make_shared<BufferAsRemoteAudioBuffer>(inputSampleRate, inputBuf);
    auto warpMap = std::make_shared<time_transform::WarpMap<double>>(
        std::vector<WarpMap<double>::WarpMarker>{
            {Beats(0.0), Seconds(0.0)}, {Beats(stretch), Seconds(1.0)}});
    auto reader = TimestretchReader(inputRd, warpMap, outputSampleRate, pro);
    TimelineTempoMap tm(1.0, {});

    const auto renderDelta =
        BeatsDelta((double)outBuf.getFrameCount() / (double)outputSampleRate);
    auto renderPositionSamples = 0;
    for (Beats i(0.0); i < Beats(3.0); i += renderDelta) {
      RenderContext rc{i, 1.0, i > Beats(0.0), tm, nullptr, renderDelta};
      reader.readSegment(&rc, i, transposition, outBuf);
      concatOutBuf.set(renderPositionSamples, outBuf);
      renderPositionSamples += outBuf.getFrameCount();
    }

    concatOutBuf.mixDownInPlace();
    auto data = concatOutBuf.getChannelData(0);
    auto maxIndex = 0;
    auto maxValue = data[0];
    for (int i = 1; i < concatOutBuf.getFrameCount(); i++) {
      if (data[i] > maxValue) {
        maxValue = data[i];
        maxIndex = i;
      }
    }

    const auto sampleRateRatio =
        (double)inputSampleRate / (double)outputSampleRate;

    REQUIRE(std::abs((double)maxIndex * sampleRateRatio -
                     (double)peakPosition * stretch) < 20);

    // seek back
    concatOutBuf.fill(0.0f);
    renderPositionSamples = 0;
    for (Beats i(0.1); i < Beats(3.0); i += renderDelta) {
      RenderContext rc{i, 1.0, i > Beats(0.1), tm, nullptr, renderDelta};
      reader.readSegment(&rc, i, transposition, outBuf);
      concatOutBuf.set(renderPositionSamples, outBuf);
      renderPositionSamples += outBuf.getFrameCount();
    }

    concatOutBuf.mixDownInPlace();
    maxIndex = 0;
    maxValue = data[0];
    for (int i = 1; i < concatOutBuf.getFrameCount(); i++) {
      if (data[i] > maxValue) {
        maxValue = data[i];
        maxIndex = i;
      }
    }

    REQUIRE(std::abs((double)maxIndex * sampleRateRatio -
                     (double)peakPosition * stretch +
                     0.1 * (double)inputSampleRate) < 20);
  };

  for (bool pro : {false, true}) {
    for (int channelCount : {1, 2}) {
      for (float transposition : {
               -5.f,
               -3.f,
               -1.f,
               0.f,
               1.f,
               3.f,
               5.f,
           }) {
        for (double stretch : {0.5, 0.83, 1.0, 1.25, 2.0}) {
          alignmentTestWithParams(pro, channelCount, 44100, 44100,
                                  transposition, stretch);
          /*
          TODO transposition plus sample rate conversion doesn't align properly

          alignmentTestWithParams(pro, channelCount, 44100, 48000,
                                  transposition, stretch);
          alignmentTestWithParams(pro, channelCount, 48000, 44100,
                                  transposition, stretch);
          alignmentTestWithParams(pro, channelCount, 48000, 16000,
                                  transposition, stretch);
          alignmentTestWithParams(pro, channelCount, 16000, 16000,
                                  transposition, stretch);
          */
        }
      }
      for (float transposition : {-24.f, -1.f, 0.f, 1.f, 24.f}) {
        for (double stretch : {0.5, 0.83, 1.0, 1.25, 2.0}) {
          testWithParams(pro, channelCount, 44100, 44100, transposition,
                         stretch);
          testWithParams(pro, channelCount, 44100, 48000, transposition,
                         stretch);
          testWithParams(pro, channelCount, 48000, 44100, transposition,
                         stretch);
          testWithParams(pro, channelCount, 48000, 16000, transposition,
                         stretch);
          testWithParams(pro, channelCount, 16000, 16000, transposition,
                         stretch);
        }
      }
    }
  }
}