// period_detect — pure implementation. See period_detect.h for the contract.
//
// YIN (de Cheveigne & Kawahara 2002), two-pass: a cumulative-mean-normalized difference
// function on a 4x box-decimated copy picks the period, then the raw difference function at
// full rate refines it to a fraction of a frame. The decimated pass is what makes the cost
// bounded; the full-rate pass is what makes the estimate precise enough to multiply — the
// splice jump is n periods, so an error of e frames lands as n*e frames of misalignment.
//
// Hand-rolled rather than autocorrelation-with-an-FFT: no third-party dependency, and the
// difference function's absolute threshold is what lets "no period here" be a real answer.

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

#include <algorithm>
#include <cmath>
#include <cstddef>

namespace reasampler::instrument::engine {

namespace {

constexpr int kDecimate = 4;
// Below this RMS a block carries no signal to find a period in; its difference function is
// numerically degenerate rather than merely inconclusive.
constexpr double kSilenceRms = 1e-5;

// Box-decimate `src[from, from+count)` by kDecimate. The averaging is the anti-alias filter:
// a plain stride would fold high partials onto the low lags the coarse pass searches.
std::vector<double> decimate(const std::vector<AudioSample>& src, std::size_t from,
                             std::size_t count) {
    std::vector<double> out(count / kDecimate);
    for (std::size_t i = 0; i < out.size(); ++i) {
        double s = 0.0;
        for (int k = 0; k < kDecimate; ++k) {
            s += static_cast<double>(src[from + i * kDecimate + static_cast<std::size_t>(k)]);
        }
        out[i] = s / kDecimate;
    }
    return out;
}

// The cumulative-mean-normalized difference d'(tau) over lags [1, lagHi], analysis window W:
//   d(tau)  = sum_{j<W} (x[j] - x[j+tau])^2
//   d'(tau) = d(tau) / ((1/tau) * sum_{t=1..tau} d(t))
// Index 0 is unused (set to 1.0, YIN's convention). The normalization is what makes the
// threshold below an absolute one rather than a signal-dependent one.
std::vector<double> cmndf(const std::vector<double>& x, std::size_t W, std::size_t lagHi) {
    std::vector<double> dp(lagHi + 1, 1.0);
    double running = 0.0;
    for (std::size_t tau = 1; tau <= lagHi; ++tau) {
        double d = 0.0;
        for (std::size_t j = 0; j < W; ++j) {
            const double diff = x[j] - x[j + tau];
            d += diff * diff;
        }
        running += d;
        dp[tau] = running > 0.0 ? d * static_cast<double>(tau) / running : 1.0;
    }
    return dp;
}

// Parabolic vertex through (i-1, i, i+1) as an offset in [-0.5, 0.5] from i. Zero at an end
// point or a non-minimum, which leaves the integer lag — benign, and the full-rate pass
// refines it again anyway.
double parabolicOffset(const std::vector<double>& y, std::size_t i) {
    if (i == 0 || i + 1 >= y.size()) return 0.0;
    const double den = y[i - 1] - 2.0 * y[i] + y[i + 1];
    if (!(den > 0.0)) return 0.0;  // a minimum has positive curvature
    double f = 0.5 * (y[i - 1] - y[i + 1]) / den;
    if (f > 0.5) f = 0.5;
    if (f < -0.5) f = -0.5;
    return f;
}

// YIN's absolute-threshold rule: take the FIRST dip below the threshold, walked down to its
// local bottom — not the global minimum. A periodic signal dips at every multiple of its
// period, so the global minimum is as likely to be 2P or 3P; taking the first dip is what
// makes the answer the fundamental period rather than some harmonic of it.
bool pickPeriod(const std::vector<double>& dp, std::size_t lagLo, double& tauOut,
                double& dissimilarity) {
    for (std::size_t tau = lagLo; tau + 1 < dp.size(); ++tau) {
        if (dp[tau] >= kPeriodDetectThreshold) continue;
        std::size_t t = tau;
        while (t + 1 < dp.size() && dp[t + 1] < dp[t]) ++t;
        tauOut = static_cast<double>(t) + parabolicOffset(dp, t);
        dissimilarity = dp[t];
        return true;
    }
    return false;
}

// The raw difference function over [lo, hi] at FULL rate, minimized parabolically. The coarse
// pass already chose which dip; this only says exactly where its bottom is. Amplitude drift
// over the few frames spanned here is negligible, so the unnormalized d() suffices.
double refineFullRate(const std::vector<AudioSample>& pcm, std::size_t from, std::size_t W,
                      std::size_t lo, std::size_t hi) {
    std::vector<double> d(hi - lo + 1, 0.0);
    for (std::size_t tau = lo; tau <= hi; ++tau) {
        double s = 0.0;
        for (std::size_t j = 0; j < W; ++j) {
            const double diff = static_cast<double>(pcm[from + j]) -
                                static_cast<double>(pcm[from + j + tau]);
            s += diff * diff;
        }
        d[tau - lo] = s;
    }
    const std::size_t best =
        static_cast<std::size_t>(std::min_element(d.begin(), d.end()) - d.begin());
    return static_cast<double>(lo + best) + parabolicOffset(d, best);
}

double blockRms(const std::vector<AudioSample>& pcm, std::size_t from, std::size_t count) {
    double e = 0.0;
    for (std::size_t i = 0; i < count; ++i) {
        const double x = static_cast<double>(pcm[from + i]);
        e += x * x;
    }
    return std::sqrt(e / static_cast<double>(count));
}

}  // namespace

PeriodEstimate detectPeriod(const std::vector<AudioSample>& pcm, int sampleRate,
                            std::size_t spanFrom, std::size_t spanCount) {
    if (sampleRate <= 0 || spanCount == 0) return {};
    if (spanFrom > pcm.size() || spanCount > pcm.size() - spanFrom) return {};
    const double rate = static_cast<double>(sampleRate);
    std::size_t lagHi = longestLagFrames(sampleRate);
    const std::size_t lagLo = static_cast<std::size_t>(rate / kPeriodDetectMaxHz);
    if (lagLo < 2) return {};  // a rate so low the whole search band collapses

    // One probe block is W + lagHi frames with W == lagHi (YIN's usual sizing: the analysis
    // window must cover the longest lag being tested). A short span shortens the search
    // rather than refusing outright — a 200 ms one-shot still has a period worth finding.
    if (spanCount < 2 * lagHi) lagHi = spanCount / 2;
    if (lagHi <= lagLo + 2) return {};
    const std::size_t block = 2 * lagHi;
    // Probe POSITIONS, not disjoint blocks — see kPeriodDetectProbes in the header for why
    // lagHi is the separation that makes two overlapping probes independent evidence.
    const std::size_t room = spanCount - block;
    const std::size_t probes = std::min<std::size_t>(kPeriodDetectProbes, 1 + room / lagHi);
    // Room to spare after the last probe's block is spread between them, so the probes sample
    // the whole span rather than only its opening.
    const std::size_t stride = probes > 1 ? room / (probes - 1) : 0;

    std::vector<double> periods;
    std::vector<double> confidences;
    // Probes that carried signal AND ran the real dip search — the agreement denominator. A
    // silent block is no evidence either way; a block whose decimated search band or full-rate
    // refine bracket collapsed to nothing (the two geometry continues below) never ran that
    // search either, so it is excluded on the same footing as silence, not counted as if it had.
    std::size_t evidence = 0;
    for (std::size_t p = 0; p < probes; ++p) {
        // (probes - 1) * stride <= room by construction, so the last block always fits.
        const std::size_t from = spanFrom + p * stride;
        if (blockRms(pcm, from, block) < kSilenceRms) continue;

        const std::vector<double> small = decimate(pcm, from, block);
        const std::size_t smallHi = lagHi / kDecimate;
        const std::size_t smallW = small.size() - smallHi;
        if (smallHi <= lagLo / kDecimate + 2 || smallW == 0) continue;  // degenerate geometry
        const std::vector<double> dp = cmndf(small, smallW, smallHi);

        double coarseTau = 0.0, dissimilarity = 1.0;
        if (!pickPeriod(dp, std::max<std::size_t>(2, lagLo / kDecimate), coarseTau,
                        dissimilarity)) {
            ++evidence;  // the search ran and found no dip: real evidence against a period
            continue;
        }

        // Bracket the full-rate refinement at +/- 2 decimated samples around the coarse pick:
        // the decimated parabola is already sub-decimated-sample accurate, so this is margin,
        // not a second search.
        const double centre = coarseTau * kDecimate;
        const std::size_t lo = static_cast<std::size_t>(
            std::max(static_cast<double>(lagLo), centre - 2.0 * kDecimate));
        const std::size_t hi = static_cast<std::size_t>(
            std::min(static_cast<double>(lagHi), centre + 2.0 * kDecimate));
        if (hi <= lo) continue;  // degenerate refine bracket
        ++evidence;  // the search ran and found a period: real evidence for one
        periods.push_back(refineFullRate(pcm, from, block - hi, lo, hi));
        confidences.push_back(1.0 - dissimilarity);
    }

    if (periods.empty()) return {};

    // ONE piece of evidence in the whole span — either it hosted a single probe position, or
    // every other probe was silent. Nothing can rule against this estimate, so the accept rests
    // on pickPeriod's absolute threshold, which is a real test and not an absence of one: the
    // block genuinely repeats at this lag across its whole analysis window. Refusing instead
    // would deny every short one-shot a period, and a period that turns out wrong costs a
    // mis-centred correlation search at the splice, not an unrefined one (pitch_shift.cpp's
    // splice searches +/- maxLag around whichever jump it is handed). Do not "unify" this back
    // into the majority test — at one piece of evidence that test accepts unconditionally, which
    // is the same behaviour with none of the reasoning. Nor key it on how many probes SURVIVED:
    // one survivor out of four that all carried signal is not this case at all.
    if (evidence == 1) {
        PeriodEstimate lone;
        lone.frames = periods[0];
        lone.confidence = confidences[0];
        return lone;
    }

    std::vector<double> sorted = periods;
    std::sort(sorted.begin(), sorted.end());
    const double median = sorted[sorted.size() / 2];

    // Average the probes that agree with the median rather than taking the median outright:
    // averaging cancels each probe's own estimation jitter, and the jump multiplies whatever
    // error survives by n.
    double sum = 0.0, confSum = 0.0;
    std::size_t agree = 0;
    for (std::size_t i = 0; i < periods.size(); ++i) {
        if (std::fabs(periods[i] - median) > kPeriodDetectAgreeTolerance * median) continue;
        sum += periods[i];
        confSum += confidences[i];
        ++agree;
    }
    // A STRICT MAJORITY of the probes that carried signal must agree, not merely two of them: a
    // source whose first half is one period and second half another gives two probes each way,
    // and taking either as "the" period would misalign every splice in the other half. Refusing
    // is the right answer there — the fixed-window fallback is what a source with no ONE period
    // gets. The denominator is `evidence` and not `periods.size()` because once probes overlap
    // a straddling block finds no period at all rather than a third one, and counting only the
    // survivors turned that two-and-two split into a two-of-three accept.
    // Reached only with two or more pieces of evidence; the lone case returned above.
    if (agree * 2 <= evidence) return {};

    PeriodEstimate est;
    est.frames = sum / static_cast<double>(agree);
    est.confidence = confSum / static_cast<double>(agree);
    return est;
}

PeriodEstimate detectPeriod(const std::vector<AudioSample>& pcm, int sampleRate) {
    return detectPeriod(pcm, sampleRate, 0, pcm.size());
}

AnalysisSpan periodAnalysisSpan(std::size_t frameCount, std::int64_t loopStart,
                                std::int64_t loopEnd, bool hasLoop, int sampleRate) {
    const AnalysisSpan whole{0, frameCount};
    if (!hasLoop || sampleRate <= 0) return whole;
    if (loopStart < 0 || loopEnd <= loopStart) return whole;
    if (static_cast<std::uint64_t>(loopEnd) > frameCount) return whole;

    const std::size_t length = static_cast<std::size_t>(loopEnd - loopStart);
    // One full probe block. Below it detectPeriod shortens lagHi to fit, which raises the
    // lowest findable fundamental — the one thing the narrower span may never cost.
    const std::size_t minimum = 2 * longestLagFrames(sampleRate);
    if (length < minimum) return whole;
    return AnalysisSpan{static_cast<std::size_t>(loopStart), length};
}

}  // namespace reasampler::instrument::engine
