// pitch_shift — pure implementation. See pitch_shift.h for the contract and regression history.
//
// Algorithm: a delay ring of 2*window frames. The write head advances one frame per input
// sample (source rate, duration preserved). One active read tap advances by the shift
// `ratio_` per frame, so its delay behind the writer drifts at (1 - ratio) per frame. When
// that delay leaves the safe band [dLow, dHigh], the tap is relocated by a nominal jump of
// one window — clamped to the filled span so it never lands in unwritten silence — refined
// by a cross-correlation search over +/- maxLag plus a parabolic peak interpolation for a
// sub-sample lag (an integer-only lag left +/-0.5-sample errors: a sideband comb at the
// splice cadence on a repitched pure sine). Old and new taps then crossfade over fadeFrames
// with a raised-cosine, amplitude-complementary pair (in-phase content sums to unity gain).
// At unity ratio the delay is frozen mid-band and no splice ever fires.

#include "core/instrument/engine/pitch_shift.h"

#include <algorithm>
#include <cmath>
#include <limits>

namespace reasampler::instrument::engine {

namespace {

constexpr double kPi = 3.14159265358979323846;

}  // namespace

void PitchShifter::configure(std::int64_t windowFrames) {
    window_ = windowFrames;
    if (window_ <= 1) {
        // Pass-through: no ring, process() returns input unchanged.
        ring_.clear();
        ringLen_ = 0;
        writePos_ = 0;
        posA_ = posB_ = 0.0;
        fading_ = false;
        fadePos_ = 0;
        fadeFrames_ = fadeLen_ = maxLag_ = corrFrames_ = dLow_ = dHigh_ = 0;
        filled_ = 0;
        ratio_ = 1.0;
        tailFrozen_ = false;
        lastSplice_ = SpliceEvent{};
        return;
    }
    // 2x-window ring: one window of splice-jump span plus search + fade headroom on each side.
    ringLen_ = 2 * window_;
    ring_.assign(static_cast<std::size_t>(ringLen_), 0.0f);
    // Geometry (all quarters of the window):
    //  - fadeFrames_: nominal splice crossfade; only safe while the outgoing tap can't reach
    //    the writer before the fade ends. splice() scales fadeLen_ down by ratio for up-shifts
    //    past ~2x so ordinary transpositions (+24 st) never read stale data mid-fade.
    //  - maxLag_: alignment search half-range — one window/4 covers a full period of any tone
    //    down to 4/window cycles-per-frame (~80 Hz at the product's 50 ms window, 44.1k).
    //  - dLow_/dHigh_: safe delay band; unity parks the tap mid-band (window/2 delay).
    //  - corrFrames_: at an up-splice the reference segment reads forward from the tap at
    //    delay ~dLow_, so dLow_-1 is exactly what exists between tap and writer; 512 bounds
    //    the splice burst.
    fadeFrames_ = std::max<std::int64_t>(window_ / 4, 1);
    maxLag_ = window_ / 4;
    dLow_ = window_ / 4;
    dHigh_ = ringLen_ - window_ / 4;
    corrFrames_ = std::max<std::int64_t>(1, std::min<std::int64_t>(dLow_ - 1, 512));
    fadeLen_ = 0;
    reset();
}

void PitchShifter::reset() {
    if (window_ > 1) {
        // Seed the active tap one window behind the writer — the exact middle of the safe
        // band [dLow, dHigh], so unity holds it there forever with maximal drift room either
        // direction. No history declared (filled_ = 0): follow with prime() or warm().
        std::fill(ring_.begin(), ring_.end(), 0.0f);
        writePos_ = 0;
        posA_ = static_cast<double>(ringLen_ - window_);
        posB_ = posA_;
        fading_ = false;
        fadePos_ = 0;
        fadeLen_ = 0;
    } else {
        writePos_ = 0;
        posA_ = posB_ = 0.0;
        fading_ = false;
        fadePos_ = 0;
        fadeLen_ = 0;
    }
    filled_ = 0;
    ratio_ = 1.0;
    tailFrozen_ = false;
    lastSplice_ = SpliceEvent{};
}

void PitchShifter::freezeTail() {
    if (window_ <= 1 || tailFrozen_) return;
    tailFrozen_ = true;
    // An in-flight crossfade was sized for a retreating writer (outgoing tap drains at
    // ratio-1 per frame); frozen, it closes at the full ratio instead. Cap the live fade so
    // it completes before tap B reaches the parked writer and reads lapped content mid-fade.
    if (fading_) {
        // Preserve t = fadePos_/fadeLen_ across the shortening so gNew is continuous at the
        // freeze frame (no gain step). Compute tOld before overwriting fadeLen_.
        const double tOld =
            static_cast<double>(fadePos_) / static_cast<double>(fadeLen_);
        double dB = static_cast<double>(writePos_) - posB_;
        const double len = static_cast<double>(ringLen_);
        while (dB < 0.0) dB += len;
        while (dB >= len) dB -= len;
        // Clamp in double before the int64 cast (matches splice() pattern; guards against UB
        // when dB/ratio_ is very large, e.g. near-unity ratio at a high sample rate).
        double left = (dB - 2.0) / ratio_;
        if (left > static_cast<double>(fadeFrames_)) left = static_cast<double>(fadeFrames_);
        const std::int64_t leftFrames = left > 1.0 ? static_cast<std::int64_t>(left) : 1;
        const std::int64_t newFadeLen = std::min(fadeLen_, fadePos_ + leftFrames);
        // Re-anchor: tOld < 1 because we are mid-fade, so newFadePos < newFadeLen (still fading).
        fadePos_ = static_cast<std::int64_t>(tOld * static_cast<double>(newFadeLen));
        fadeLen_ = newFadeLen;
    }
}

void PitchShifter::prime(const AudioSample* src, std::int64_t count) {
    if (window_ <= 1) return;  // pass-through needs no priming
    // Clamp to one window: delay == count must stay inside the safe band so the seed itself
    // never triggers a splice.
    if (count < 0) count = 0;
    if (count > window_) count = window_;
    std::fill(ring_.begin(), ring_.end(), 0.0f);
    for (std::int64_t i = 0; i < count; ++i) ring_[static_cast<std::size_t>(i)] = src[i];
    // Writer continues after the primed span; the tap parks ON src[0] (delay == count), so
    // the very first process() output is src[0] — zero structural latency at every ratio.
    writePos_ = count % ringLen_;
    posA_ = posB_ = 0.0;
    fading_ = false;
    fadePos_ = 0;
    fadeLen_ = 0;
    filled_ = count;
    tailFrozen_ = false;  // a fresh note-on always starts with a live writer
    lastSplice_ = SpliceEvent{};
    // ratio_ deliberately untouched: the voice sets it per frame around the prime.
}

void PitchShifter::warm() {
    if (window_ <= 1) return;  // pass-through needs no warm-up
    // A prime() with one window of silence: same geometry (tap parked mid-band one window
    // behind the writer), the zeros declared as valid history. At unity this is a bit-exact
    // window() delay; an up-shift plays ~a window of silence before speaking (the pre-GA2
    // onset) — stream callers with access to the upcoming source should prime() instead.
    std::fill(ring_.begin(), ring_.end(), 0.0f);
    writePos_ = window_ % ringLen_;
    posA_ = posB_ = 0.0;
    fading_ = false;
    fadePos_ = 0;
    fadeLen_ = 0;
    filled_ = window_;
    tailFrozen_ = false;
    lastSplice_ = SpliceEvent{};
}

void PitchShifter::setShiftRatio(double ratio) {
    if (ratio > 0.0) ratio_ = ratio;  // ignore non-positive (never run the tap backward/stall)
}

double PitchShifter::readTap(double pos) const {
    // Fractional linear interpolation with ring wrap.
    double p = pos;
    const double len = static_cast<double>(ringLen_);
    while (p < 0.0) p += len;
    while (p >= len) p -= len;
    const std::int64_t i0 = static_cast<std::int64_t>(p);
    const double frac = p - static_cast<double>(i0);
    std::int64_t i1 = i0 + 1;
    if (i1 >= ringLen_) i1 = 0;
    const double s0 = static_cast<double>(ring_[static_cast<std::size_t>(i0)]);
    const double s1 = static_cast<double>(ring_[static_cast<std::size_t>(i1)]);
    return s0 + (s1 - s0) * frac;
}

void PitchShifter::splice(std::int64_t nominalJump, double delay) {
    // Relocate the active tap by `nominalJump` frames of added delay (+window_ = toward older
    // content, -window_ = toward the writer), refined by a correlation search so the relocated
    // read point is waveform-aligned with the outgoing tap's upcoming content. Search is coarse
    // (step 4 over +/- maxLag_) then fine (+/- 3 around the coarse best, then a parabolic
    // sub-sample peak): a bounded burst of ~ (maxLag_/2 + 9) * corrFrames_ multiply-adds, once
    // per splice.
    const std::int64_t d = static_cast<std::int64_t>(delay);

    // An up-jump may only relocate into valid history. The deepest slot the search (plus the
    // parabola's +/-1 probe and the interpolator's read-ahead) can touch is d + jump + maxLag + 2,
    // so the cap is filled_ - d - maxLag_ - 1 (one sample looser than that derived bound, not
    // extra margin — ring indexing wraps via modulo everywhere regardless). In steady state
    // (filled_ == ringLen_) this exceeds window_ and the nominal jump is untouched; near a primed
    // onset it shrinks the jump to what real history exists. The floor of 1 only fires on the
    // degenerate reset-without-prime path.
    std::int64_t jump = nominalJump;
    if (jump > 0) {
        const std::int64_t maxJump = filled_ - d - maxLag_ - 1;
        if (jump > maxJump) jump = maxJump;
        if (jump < 1) jump = 1;
    }
    // The correlation reference reads FORWARD from the tap; keep it strictly behind the
    // writer even when the trigger undershot dLow_ by a large per-frame drift (extreme
    // up-ratios): d - corr must stay >= 0.
    const std::int64_t corr = std::max<std::int64_t>(1, std::min<std::int64_t>(corrFrames_, d - 1));

    const std::int64_t iA =
        ((static_cast<std::int64_t>(posA_) % ringLen_) + ringLen_) % ringLen_;

    auto scoreAt = [&](std::int64_t lag) -> double {
        std::int64_t ia = iA;
        std::int64_t ic = ((iA - jump + lag) % ringLen_ + ringLen_) % ringLen_;
        double s = 0.0, ec = 0.0;
        for (std::int64_t k = 0; k < corr; ++k) {
            const double a = static_cast<double>(ring_[static_cast<std::size_t>(ia)]);
            const double c = static_cast<double>(ring_[static_cast<std::size_t>(ic)]);
            s += a * c;
            ec += c * c;
            if (++ia >= ringLen_) ia = 0;
            if (++ic >= ringLen_) ic = 0;
        }
        // Normalized cross-correlation: a raw dot product biases toward the higher-energy lag,
        // so on a decaying tail every up-splice would prefer the loudest candidate over the
        // best-aligned one. The reference segment's energy is constant across lags, so dividing
        // by sqrt(Ec) alone ranks identically to the full normalized form.
        return ec > 0.0 ? s / std::sqrt(ec) : 0.0;
    };

    std::int64_t bestLag = 0;
    double bestScore = -std::numeric_limits<double>::infinity();
    for (std::int64_t lag = -maxLag_; lag <= maxLag_; lag += 4) {
        const double s = scoreAt(lag);
        if (s > bestScore) {
            bestScore = s;
            bestLag = lag;
        }
    }
    const std::int64_t coarse = bestLag;
    for (std::int64_t lag = coarse - 3; lag <= coarse + 3; ++lag) {
        if (lag == coarse || lag < -maxLag_ || lag > maxLag_) continue;
        const double s = scoreAt(lag);
        if (s > bestScore) {
            bestScore = s;
            bestLag = lag;
        }
    }

    // Sub-sample peak: the integer-lag best leaves a residual misalignment of up to half a
    // sample, which at the splice cadence phase-modulates a pure tone into an audible sideband
    // comb. A parabola through the scores at bestLag-1/bestLag/bestLag+1 locates the peak to a
    // fraction of a sample; readTap()'s linear interpolation realizes it. The denominator is
    // negative at a genuine peak — flat correlation (DC/silence) keeps the integer lag, benign.
    double frac = 0.0;
    {
        const double sM = scoreAt(bestLag - 1);
        const double sP = scoreAt(bestLag + 1);
        const double den = sM - 2.0 * bestScore + sP;
        if (den < 0.0) {
            frac = 0.5 * (sM - sP) / den;
            if (frac > 0.5) frac = 0.5;
            if (frac < -0.5) frac = -0.5;
        }
    }

    // Hand the current position to the outgoing tap and relocate the active one.
    posB_ = posA_;
    double p = posA_ - static_cast<double>(jump) + static_cast<double>(bestLag) + frac;
    const double len = static_cast<double>(ringLen_);
    while (p < 0.0) p += len;
    while (p >= len) p -= len;
    posA_ = p;
    // Ratio-scaled fade length. At an up-splice the outgoing tap keeps draining toward the
    // writer at (ratio - 1) per frame; the nominal window/4 fade only keeps it behind the
    // writer for ratios up to 2 — beyond that (e.g. +24 st = ratio 4) it would cross mid-fade
    // and play stale read-ahead data. Cap the live fade at the drain headroom actually
    // available, minus 2 (trigger undershoot + interpolator read-ahead margin). Down-shifts
    // drain at (1 - ratio) < 1 per frame and can't reach the ring end within window/4 frames,
    // so they always keep the full fade.
    //
    // Tail-frozen: with the writer parked, the outgoing tap closes on it at the full ratio in
    // either shift direction, so the drain rate is ratio_ instead of (ratio_ - 1) and the cap
    // applies at every ratio (including unity, since delay now drains at unity too).
    fadeLen_ = fadeFrames_;
    const double drainRate = tailFrozen_ ? ratio_ : (ratio_ - 1.0);
    if (drainRate > 0.0) {
        const double headroom = static_cast<double>(dLow_) - drainRate - 2.0;
        // Clamp in double before the int64 cast to avoid UB at pathological near-unity ratios
        // at very high sample rates (where headroom/drainRate could overflow int64).
        const double safeDbl = headroom > 0.0
            ? std::min(headroom / drainRate, static_cast<double>(fadeFrames_))
            : 1.0;
        fadeLen_ = std::max<std::int64_t>(1, static_cast<std::int64_t>(safeDbl));
    }
    fading_ = true;
    fadePos_ = 0;
    // Record the decision for a linked follower channel — applied verbatim there so both
    // channels share one lag and one splice schedule.
    lastSplice_ = SpliceEvent{true, jump, bestLag, frac, fadeLen_};
}

void PitchShifter::applySplice(const SpliceEvent& ev) {
    // Follower half of the linked lag: relocate + fade with the master's decision, no
    // correlation search of our own — the master's jump/fade derive from shared geometry +
    // ratio, which match ours by the lockstep contract (identical configure/prime/ratio history).
    posB_ = posA_;
    double p = posA_ - static_cast<double>(ev.jump) + static_cast<double>(ev.lag) + ev.frac;
    const double len = static_cast<double>(ringLen_);
    while (p < 0.0) p += len;
    while (p >= len) p -= len;
    posA_ = p;
    fadeLen_ = std::max<std::int64_t>(1, ev.fadeLen);
    fading_ = true;
    fadePos_ = 0;
    lastSplice_ = ev;  // observable mirror (tests assert follower == master per frame)
}

AudioSample PitchShifter::process(AudioSample in) { return processImpl(in, nullptr); }

AudioSample PitchShifter::processLinked(AudioSample in, const SpliceEvent& master) {
    return processImpl(in, &master);
}

AudioSample PitchShifter::processImpl(AudioSample in, const SpliceEvent* linked) {
    if (window_ <= 1) return in;  // pass-through (unconfigured / degenerate)

    // Copy the linked decision before clearing lastSplice_ (guards a self-aliased pointer).
    const SpliceEvent linkedEv = linked != nullptr ? *linked : SpliceEvent{};
    lastSplice_ = SpliceEvent{};  // cleared every frame; set again if this frame splices

    // Tail-frozen: the source is exhausted, `in` is padding, not stream — write nothing (the
    // ring keeps its all-real final two windows) and hold the write head; read/splice/fade
    // below run unchanged over the frozen content.
    if (!tailFrozen_) {
        ring_[static_cast<std::size_t>(writePos_)] = in;
        if (filled_ < ringLen_) ++filled_;
    }

    // Read the active tap; while a splice fade is live, crossfade against the outgoing tap.
    // Raised-cosine complementary gains (gNew + gOld == 1): correlation-aligned content is in
    // phase, so the sum holds unity amplitude through the fade (equal-power would bulge).
    double out = readTap(posA_);
    if (fading_) {
        const double t = static_cast<double>(fadePos_) / static_cast<double>(fadeLen_);
        const double gNew = 0.5 * (1.0 - std::cos(kPi * t));
        out = gNew * out + (1.0 - gNew) * readTap(posB_);
        if (++fadePos_ >= fadeLen_) fading_ = false;
    } else if (linked != nullptr) {
        // Follower: no trigger test, no search — splice exactly when and how the master did
        // this frame (lockstep means our own trigger would have fired the same frame anyway).
        if (linkedEv.fired) {
            applySplice(linkedEv);
        } else {
            // Self-healing fallback: if the processor ever renders a mono block mid-note, this
            // follower is skipped for that block while the master keeps advancing, and could
            // never resync via the `linkedEv.fired` path alone. So also check this follower's
            // own tap distance against the safe band and splice via its own search when it has
            // left [dLow_, dHigh_] — never triggers in the normal (non-mono-block) case, since
            // the master's trigger always fires first.
            double d = static_cast<double>(writePos_) - posA_;
            const double len = static_cast<double>(ringLen_);
            while (d < 0.0) d += len;
            while (d >= len) d -= len;
            if (d <= static_cast<double>(dLow_)) {
                splice(+window_, d);
            } else if (d >= static_cast<double>(dHigh_)) {
                splice(-window_, d);
            }
        }
    } else {
        // Splice scheduling: relocate when the active tap's delay leaves the safe band.
        // Up-shifts drain the delay toward 0 -> jump one window older; down-shifts grow it
        // toward the ring length -> jump one window toward the writer. At unity the delay is
        // frozen at window/2 and neither trigger ever fires.
        double d = static_cast<double>(writePos_) - posA_;
        const double len = static_cast<double>(ringLen_);
        while (d < 0.0) d += len;
        while (d >= len) d -= len;
        if (d <= static_cast<double>(dLow_)) {
            splice(+window_, d);
        } else if (d >= static_cast<double>(dHigh_)) {
            splice(-window_, d);
        }
    }

    // Advance heads: write head one frame (parked while tail-frozen), tap(s) by the shift ratio.
    if (!tailFrozen_) {
        ++writePos_;
        if (writePos_ >= ringLen_) writePos_ = 0;
    }
    const double len = static_cast<double>(ringLen_);
    posA_ += ratio_;
    while (posA_ >= len) posA_ -= len;
    if (fading_) {
        posB_ += ratio_;
        while (posB_ >= len) posB_ -= len;
    }

    return static_cast<AudioSample>(out);
}

}  // namespace reasampler::instrument::engine