"""correlate — recover a loop's provenance by finding the stem it was cut from.

A hand-cut ParVagues loop is *literally a slice of a separated stem* (PLN keeps
stems as-cut, no normalize — PHASE2_FINDINGS §0/§5·0). So normalized cross-
correlation of the loop's waveform against a stem's waveform has a razor-sharp
peak at the true offset: a real slice scores ~1.0, a foreign clip ~0.03. That
peak gives us, for free and at scale: which source track, which stem, the start
offset, and the length — dozens of ground-truth (sample → source) pairs the
finder (#7) and the tierlist cut-critique (#11) can lean on.

Only the ~recent demucs-derived kits will match (separated/ holds far fewer
tracks than the kit count); the driver QUANTIFIES coverage rather than pretending
every loop has a recoverable source.

`ncc_locate` is pure (numpy + scipy, tested on synthetic slices). Runs under
system python3; no GPU, no network.

    python3 -m engine.correlate <loop.wav> <stem.wav> [more-stems…]   # locate one loop
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional

import numpy as np
from scipy.signal import fftconvolve

TARGET_SR = 4000          # plenty of waveform shape for a discriminative match; fast
MIN_SCORE = 0.55          # below this, treat as "no source found" (foreign ≈ 0.03)
_EPS = 1e-9


def load_mono(path: str | Path, sr: int = TARGET_SR) -> np.ndarray:
    """Load any audio → mono float32 at `sr` (soxr-backed resample via librosa)."""
    import librosa
    y, _ = librosa.load(str(path), sr=sr, mono=True)
    return y.astype(np.float32)


def ncc_locate(loop: np.ndarray, stem: np.ndarray) -> tuple[float, int]:
    """Normalized cross-correlation peak of `loop` within `stem`.

    Returns (score 0..1, offset in samples). score≈1 ⇒ loop is that slice of stem.
    If the loop is longer than the stem, returns (0, 0).
    """
    L = loop.size
    if L == 0 or stem.size < L:
        return 0.0, 0
    loop = loop - loop.mean()
    ln = float(np.sqrt((loop.astype(np.float64) ** 2).sum())) + _EPS
    # numerator: sliding dot product (cross-correlation) via FFT
    num = fftconvolve(stem, loop[::-1], mode="valid")
    # denominator: sliding L2 norm of the (mean-removed) stem window
    st = stem.astype(np.float64)
    c1 = np.cumsum(np.r_[0.0, st])
    c2 = np.cumsum(np.r_[0.0, st * st])
    wsum = c1[L:] - c1[:-L]
    wsq = c2[L:] - c2[:-L]
    wvar = np.maximum(wsq - wsum * wsum / L, 0.0)      # window L2² (mean-removed)
    # Floor the window norm relative to the stem's overall energy so near-silent
    # windows (where catastrophic cancellation drives wvar→0) can't blow the ratio
    # up — otherwise a quiet gap scores ∞ and masquerades as a perfect match.
    gstd = float(st.std()) + _EPS
    floor = np.sqrt(L) * 0.05 * gstd
    denom = np.maximum(np.sqrt(wvar), floor) * ln
    ncc = np.clip(num / denom, -1.0, 1.0)
    i = int(np.argmax(ncc))
    return float(ncc[i]), i


def envelope(y: np.ndarray, sr: int, fps: int = 100) -> np.ndarray:
    """Frame-RMS energy envelope — a cheap, phase-robust fingerprint for prefiltering.

    A true slice's envelope matches a contiguous span of the source's envelope, so
    NCC on envelopes ranks candidate stems at ~1/40th the sample-domain cost (and
    unlike naive decimation, it doesn't depend on the offset being phase-aligned).
    """
    hop = max(1, sr // fps)
    n = y.size // hop
    if n == 0:
        return np.sqrt(np.mean(y ** 2, keepdims=True)) if y.size else np.zeros(1)
    frames = y[: n * hop].reshape(n, hop)
    return np.sqrt((frames.astype(np.float64) ** 2).mean(axis=1) + _EPS)


def prepare(stems: dict[str, np.ndarray], *, sr: int = TARGET_SR) -> dict[str, tuple]:
    """Precompute (full, envelope) per stem once, so locate() can coarse-rank cheaply."""
    return {sid: (arr, envelope(arr, sr)) for sid, arr in stems.items()}


def locate(loop_path: str | Path, prepared: dict[str, tuple],
           *, sr: int = TARGET_SR, min_score: float = MIN_SCORE,
           topk: int = 3) -> Optional[dict]:
    """Find the stem a loop was cut from. Coarse-to-fine:

    1. rank every stem by NCC on the energy envelope (cheap);
    2. run the precise sample-domain NCC only on the top-`topk` candidates.

    `prepared` comes from prepare(); returns {stem, score, offset_s, length_s} or None.
    """
    loop = load_mono(loop_path, sr)
    lenv = envelope(loop, sr)
    coarse = []
    for sid, (arr, senv) in prepared.items():
        if senv.size < lenv.size or arr.size < loop.size:
            continue
        sc, _ = ncc_locate(lenv, senv)
        coarse.append((sc, sid))
    if not coarse:
        return None
    coarse.sort(reverse=True)
    best = None
    for _, sid in coarse[:topk]:
        arr = prepared[sid][0]
        score, off = ncc_locate(loop, arr)
        if best is None or score > best[1]:
            best = (sid, score, off)
    if best is None or best[1] < min_score:
        return None
    sid, score, off = best
    return {"stem": sid, "score": round(score, 4),
            "offset_s": round(off / sr, 3), "length_s": round(loop.size / sr, 3)}


def _main(argv=None) -> int:
    import argparse
    import json

    ap = argparse.ArgumentParser(description="locate the stem a loop was cut from")
    ap.add_argument("loop")
    ap.add_argument("stems", nargs="+", help="candidate stem wavs")
    ap.add_argument("--min-score", type=float, default=MIN_SCORE)
    a = ap.parse_args(argv)

    prepared = prepare({f"{Path(s).parent.name}/{Path(s).stem}": load_mono(s) for s in a.stems})
    hit = locate(a.loop, prepared, min_score=a.min_score)
    print(json.dumps(hit, indent=2) if hit else "no source found above threshold")
    return 0


if __name__ == "__main__":
    raise SystemExit(_main())
