"""Tests for the provenance correlator (engine/correlate.py).

Synthetic: embed a known slice in a longer 'stem' and recover its offset; confirm
a foreign clip scores near zero. No audio files, no librosa needed for the math.
"""
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from engine import correlate as C


def _stem(n=200_000, seed=0):
    return np.random.default_rng(seed).standard_normal(n).astype(np.float32)


def test_ncc_recovers_exact_offset():
    stem = _stem()
    off, L = 73_210, 16_000
    loop = stem[off:off + L].copy()
    score, found = C.ncc_locate(loop, stem)
    assert found == off
    assert score > 0.99


def test_ncc_foreign_clip_scores_low():
    stem = _stem(seed=1)
    foreign = np.random.default_rng(99).standard_normal(16_000).astype(np.float32)
    score, _ = C.ncc_locate(foreign, stem)
    assert score < 0.2


def test_ncc_loop_longer_than_stem():
    assert C.ncc_locate(_stem(1000), _stem(500)) == (0.0, 0)


def test_ncc_silent_windows_dont_blow_up():
    # a stem that's mostly silence with one real region; a FOREIGN loop must score
    # low, not ∞ (silent windows underflow the normaliser → spurious perfect match).
    rng = np.random.default_rng(0)
    stem = np.zeros(200_000, dtype=np.float32)
    stem[80_000:120_000] = rng.standard_normal(40_000).astype(np.float32)
    foreign = rng.standard_normal(12_000).astype(np.float32)
    score, _ = C.ncc_locate(foreign, stem)
    assert score <= 1.0 and score < 0.3
    # the true slice in the loud region still resolves exactly
    loop = stem[90_000:102_000].copy()
    s2, off = C.ncc_locate(loop, stem)
    assert off == 90_000 and s2 > 0.99


def test_ncc_survives_gain_change():
    # the loop graded as-cut, but a global gain shouldn't break the match (NCC is
    # scale-invariant) — guards against a future "PLN nudged the level" case.
    stem = _stem(seed=2)
    off, L = 40_000, 12_000
    loop = stem[off:off + L].copy() * 0.5
    score, found = C.ncc_locate(loop, stem)
    assert found == off and score > 0.99


def test_locate_picks_right_stem(monkeypatch):
    stems = {"trackA/bass": _stem(seed=3), "trackB/vocals": _stem(seed=4)}
    off = 50_000
    true_loop = stems["trackB/vocals"][off:off + 10_000].copy()
    # bypass file IO: feed the array straight through load_mono
    monkeypatch.setattr(C, "load_mono", lambda p, sr=C.TARGET_SR: true_loop)
    hit = C.locate("ignored.wav", C.prepare(stems, sr=4000), sr=4000)
    assert hit["stem"] == "trackB/vocals"
    assert hit["score"] > 0.99
    assert hit["offset_s"] == round(off / 4000, 3)


def test_locate_returns_none_below_threshold(monkeypatch):
    stems = {"trackA/bass": _stem(seed=5)}
    foreign = np.random.default_rng(7).standard_normal(8_000).astype(np.float32)
    monkeypatch.setattr(C, "load_mono", lambda p, sr=C.TARGET_SR: foreign)
    assert C.locate("ignored.wav", C.prepare(stems)) is None


def test_envelope_matches_slice_offset():
    # the coarse prefilter: a slice's envelope should rank its own source highest
    rng = np.random.default_rng(11)
    stem = rng.standard_normal(120_000).astype(np.float32)
    loop = stem[30_000:42_000].copy()
    se, le = C.envelope(stem, 4000), C.envelope(loop, 4000)
    score, _ = C.ncc_locate(le, se)
    assert score > 0.9          # envelope NCC still discriminates the true source
