"""Tests for the loop finder (engine/loops.py).

Pure DSP/logic on synthetic inputs, plus one light integration (beat_grid tempo
recovery, which guards the §7 octave fix). Composite scoring weights are v1 and
calibrated in #7, so assertions are structural/relative, not pinned to thresholds.
"""
import sys
from pathlib import Path

import numpy as np
import soundfile as sf

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

from engine import loops as L


def test_checkerboard_shape_and_signs():
    k = 4
    ker = L._checkerboard(k)
    assert ker.shape == (2 * k, 2 * k)
    # top-left & bottom-right quadrants positive; off-diagonal quadrants negative
    assert ker[:k, :k].sum() > 0 and ker[k:, k:].sum() > 0
    assert ker[:k, k:].sum() < 0 and ker[k:, :k].sum() < 0


def test_foote_novelty_peaks_at_block_boundary():
    # two self-similar blocks → a checkerboard novelty spike at their seam (idx 10)
    n = 20
    ssm = np.zeros((n, n))
    ssm[:10, :10] = 1.0
    ssm[10:, 10:] = 1.0
    nov = L.foote_novelty(ssm, 4)
    assert 8 <= int(np.argmax(nov)) <= 12


def test_structural_high_when_window_repeats():
    n = 8
    ssm = np.eye(n)
    # window [0:2] strongly resembles [4:6] elsewhere
    ssm[0, 4] = ssm[1, 5] = ssm[4, 0] = ssm[5, 1] = 1.0
    rep = L._structural(ssm, 0, 2)
    uniq = L._structural(ssm, 2, 4)
    assert rep > uniq


def _cand(start, end, score, bars=2):
    return L.LoopCandidate(start_s=start, end_s=end, bars=bars, bpm=120,
                           tempo_unstable=False, score=score, structural=0.5,
                           novelty=0.5, seam=0.5, zc=0.5)


def test_overlap_fraction():
    a, b = _cand(0, 4, .9), _cand(2, 6, .8)
    assert 0.4 < L._overlap(a, b) < 0.6           # 2s overlap of 4s windows
    assert L._overlap(_cand(0, 4, .9), _cand(10, 14, .8)) == 0.0


def test_dedup_keeps_best_drops_overlap():
    cands = [_cand(0, 4, .9), _cand(2, 6, .8), _cand(20, 24, .7)]
    out = L._dedup(cands, top_n=8)
    starts = sorted(c.start_s for c in out)
    assert starts == [0.0, 20.0]                  # the .8 overlapping the .9 is dropped


def test_snap_zc_moves_to_crossing():
    sr = 44100
    t = np.arange(sr) / sr
    y = np.sin(2 * np.pi * 100 * t).astype(np.float32)  # zero crossings every 220.5 samp
    idx = 1000
    snapped = L._snap_zc(y, idx, sr)
    assert abs(y[snapped]) < abs(y[idx]) + 1e-3   # lands on/near a crossing
    assert abs(snapped - idx) <= int(sr * L.ZC_TOL_MS / 1000)


def test_snap_length_preserving_keeps_exact_length_on_near_zero_edges():
    # the whole window slides by one common offset → length is rigid, and both edges
    # land near zero. Independent per-edge snap would let the two edges drift apart.
    sr = 44100
    t = np.arange(sr * 3) / sr
    y = np.sin(2 * np.pi * 100 * t).astype(np.float32)   # zero crossings every ~220.5 samp
    a0, length = 5000, 88200                              # 2 s window, arbitrary offset
    a, b = L._snap_length_preserving(y, a0, sr, length)
    assert b - a == length                               # LENGTH exactly preserved
    assert abs(a - a0) <= int(sr * L.ZC_TOL_MS / 1000)   # start within B1 tolerance
    assert abs(y[a]) < 0.15 and abs(y[b]) < 0.15         # both edges near zero


def test_bar_len_samples_matches_bpm():
    sr = 44100
    # 132.51 bpm, beat = 60/132.51 s, bar = 4 beats; a 2-bar loop:
    n = L.bar_len_samples(2, 132.51, sr)
    expected = 2 * 4 * (60.0 / 132.51) * sr
    assert abs(n - expected) < 1.0                       # rounds to the nearest sample


def test_exported_duration_is_exact_bar_multiple(tmp_path):
    # END-TO-END regression: synthesize a tempo'd stem, export a 2-bar take, and assert the
    # written file's duration is within <1 ms of exactly 2 bars — the guarantee the
    # length-preserving snap must give (independent per-edge snap drifted −4 to −11 ms).
    from engine.model import Catch, Stem
    sr = 44100
    bpm = 132.51
    bar_len_s = 4 * 60.0 / bpm
    rng = np.random.default_rng(7)
    T = np.arange(int(20 * sr)) / sr
    # tonal bed + light noise so zero crossings exist everywhere for the snap to grab
    y = (0.3 * np.sin(2 * np.pi * 110 * T) + 0.05 * rng.standard_normal(len(T))).astype(np.float32)

    catch_dir = tmp_path / "ws" / "catchslug"
    (catch_dir / "stems").mkdir(parents=True)
    stem_path = catch_dir / "stems" / "drums.wav"
    sf.write(str(stem_path), y, sr, subtype="PCM_24")
    catch = Catch(slug="catchslug", title="t", source_id="x", source_url="u",
                  source_path="source.wav", duration=20.0,
                  stems=[Stem(name="drums", path="stems/drums.wav")])

    take = L.Take(start_s=3.3, end_s=3.3 + 2 * bar_len_s, bars=2, bpm=bpm,
                  tempo_unstable=False, score=0.9, n_stems=1,
                  stems=[L.StemSlice(name="drums", start_s=3.3, end_s=3.3 + 2 * bar_len_s,
                                     score=0.9, rms_dbfs=-12.0, is_one_shot=False)])
    res = L.export_take_report(take, "test_kit", tmp_path / "ws", catch,
                               keep=("drums",), samples_root=tmp_path / "Samples")
    assert res["written"], "should have written the drums slice"
    info = sf.info(str(res["written"][0]))
    dur = info.frames / info.samplerate
    assert abs(dur - 2 * bar_len_s) < 1e-3, f"exported {dur:.6f}s vs {2*bar_len_s:.6f}s"
    # and the explicit multiples check must agree
    assert res["multiples"]["ok"]
    assert res["multiples"]["files"][0]["dev_pct"] < 0.5


def _loop_with_downbeat_at(slot, n_beats=8, sr=22050, beat_s=0.5):
    """Synthetic bar-exact loop: soft high ticks on every beat, one BIG low thump at `slot`."""
    n = int(n_beats * beat_s * sr)
    y = 0.02 * np.random.default_rng(9).standard_normal(n).astype(np.float32)
    for k in range(n_beats):
        i = int(k * n / n_beats)
        L_hit = int(0.05 * sr)
        env = np.exp(-np.linspace(0, 10, L_hit)).astype(np.float32)
        tick = env * np.sin(2 * np.pi * 3000 * np.arange(L_hit) / sr).astype(np.float32)
        y[i:i + L_hit] += 0.1 * tick                                # every beat: soft tick
        if k == slot:                                               # the downbeat: low thump
            kick = env * np.sin(2 * np.pi * 60 * np.arange(L_hit) / sr).astype(np.float32)
            y[i:i + L_hit] += 0.9 * kick
    return y, sr


def test_rotate_to_downbeat_moves_kick_to_zero():
    # BCDA regression: a loop whose downbeat (low thump) sits at slot 5 must come out
    # rotated so the thump is at position 0 (within the zc-snap tolerance of the slot).
    y, sr = _loop_with_downbeat_at(5)
    rotated, slot, roll = L._rotate_to_downbeat(y, sr, 8)
    assert slot == 5
    assert rotated.shape == y.shape                                 # length untouched
    # the low thump now opens the loop: first 60 ms carries (near-)max energy
    w = int(0.06 * sr)
    first = float(np.sqrt(np.mean(rotated[:w] ** 2)))
    windows = [float(np.sqrt(np.mean(rotated[i:i + w] ** 2)))
               for i in range(0, len(rotated) - w, w)]
    assert first >= 0.9 * max(windows), "downbeat thump should now sit at position 0"


def test_rotate_noop_when_downbeat_already_at_zero():
    y, sr = _loop_with_downbeat_at(0)
    rotated, slot, roll = L._rotate_to_downbeat(y, sr, 8)
    assert slot == 0 and roll == 0
    assert np.array_equal(rotated, y)


def test_rotation_never_breaks_a_clean_seam():
    # a seamless loop (integer periods, continuous everywhere) stays seamless after
    # rotation — the new wrap joins samples that were contiguous in the source.
    from engine import grade as G
    sr = 22050
    n = sr * 2
    y = (0.4 * np.sin(2 * np.pi * 110 * np.arange(n) / sr)).astype(np.float32)  # 220 periods
    rotated, slot, roll = L._rotate_to_downbeat(y, sr, 8)
    assert rotated.shape == y.shape
    seam, _ = G.seam_score(rotated.astype(np.float64))
    assert seam > 0.6, "rotating a continuous loop must not introduce a wrap discontinuity"


def test_export_rotation_locks_stems_sharing_a_window(tmp_path):
    # stems exported from the SAME window must rotate by the SAME slot (anchored on
    # drums), or their relative groove shifts by whole beats. drums' downbeat is at
    # slot 3; the other stem's own loudest attack is at slot 6 — both must come out
    # rotated by drums' slot 3, so the other stem's marker lands at slot 3 (6−3), NOT 0.
    from engine.model import Catch, Stem
    sr = 22050
    bpm = 120.0
    n_beats, beat_s = 8, 0.5
    total = int(20 * sr)
    drums = 0.01 * np.random.default_rng(2).standard_normal(total).astype(np.float32)
    other = 0.01 * np.random.default_rng(3).standard_normal(total).astype(np.float32)
    win_start = 5.0
    L_hit = int(0.05 * sr)
    env = np.exp(-np.linspace(0, 10, L_hit)).astype(np.float32)
    kick = 0.9 * env * np.sin(2 * np.pi * 60 * np.arange(L_hit) / sr).astype(np.float32)
    stab = 0.9 * env * np.sin(2 * np.pi * 800 * np.arange(L_hit) / sr).astype(np.float32)
    drums[int((win_start + 3 * beat_s) * sr):int((win_start + 3 * beat_s) * sr) + L_hit] += kick
    other[int((win_start + 6 * beat_s) * sr):int((win_start + 6 * beat_s) * sr) + L_hit] += stab

    ws = tmp_path / "ws"
    cd = ws / "c" / "stems"; cd.mkdir(parents=True)
    sf.write(str(cd / "drums.wav"), drums, sr, subtype="PCM_24")
    sf.write(str(cd / "other.wav"), other, sr, subtype="PCM_24")
    catch = Catch(slug="c", title="t", source_path="s.wav",
                  stems=[Stem(name="drums", path="stems/drums.wav"),
                         Stem(name="other", path="stems/other.wav")])
    take = L.Take(start_s=win_start, end_s=win_start + n_beats * beat_s, bars=2, bpm=bpm,
                  tempo_unstable=False, score=0.9, n_stems=2,
                  stems=[L.StemSlice(name="drums", start_s=win_start,
                                     end_s=win_start + n_beats * beat_s, score=0.9,
                                     rms_dbfs=-12.0, is_one_shot=False),
                         L.StemSlice(name="other", start_s=win_start,
                                     end_s=win_start + n_beats * beat_s, score=0.8,
                                     rms_dbfs=-12.0, is_one_shot=False)])
    written = L.export_take(take, "rotkit", ws, catch, keep=("drums", "other"),
                            samples_root=tmp_path / "Samples")
    assert len(written) == 2
    outs = {p.stem.split("_")[1]: p for p in written}

    def marker_slot(path):
        y, srf = sf.read(str(path), always_2d=True); m = np.abs(y[:, 0])
        n = len(m)
        return int(round(int(np.argmax(m)) / (n / n_beats))) % n_beats

    assert marker_slot(outs["drums"]) == 0, "drums downbeat rotated to the loop start"
    assert marker_slot(outs["other"]) == 3, "other follows drums' rotation (6−3), not its own"


def test_kit_multiples_check_flags_off_grid(tmp_path):
    sr = 44100
    bpm = 120.0
    bar_len_s = 4 * 60.0 / bpm                           # 2.0 s
    good = tmp_path / "good.wav"; bad = tmp_path / "bad.wav"
    sf.write(str(good), np.zeros(int(2 * bar_len_s * sr), dtype=np.float32), sr)  # exactly 2 bars
    sf.write(str(bad), np.zeros(int(2.15 * bar_len_s * sr), dtype=np.float32), sr)  # 7.5% off
    rep = L.kit_multiples_check([good, bad], bpm)
    assert not rep["ok"]
    rows = {r["name"]: r for r in rep["files"]}
    assert rows["good.wav"]["ok"] and rows["good.wav"]["dev_pct"] < 0.5
    assert not rows["bad.wav"]["ok"] and rows["bad.wav"]["dev_pct"] > 0.5


def test_analyze_stem_honors_shared_grid():
    # a sparse stem (few onsets) gets no usable per-stem grid, but a shared grid
    # from a rhythmic source lets it produce candidates (#18 grid-from-drums).
    sr = 22050
    sparse = (0.05 * np.random.default_rng(3).standard_normal(sr * 12)).astype(np.float32)
    # build a 120-bpm grid (beats every 0.5 s) as (peaks_frames, times, bpm)
    times = np.arange(0, 12, 0.5)
    peaks = (times * sr / 512).astype(int)
    bpm = np.full(len(times), 120.0)
    cands = L.analyze_stem(sparse, sr, bars=(2,), top_n=4, grid=(peaks, times, bpm))
    assert isinstance(cands, list)            # runs on the shared grid without recomputing
    for c in cands:
        assert 100 <= c.bpm <= 140            # uses the supplied grid's tempo


def test_analyze_chops_finds_subbar_slices_between_onsets():
    # three short bursts separated by silence → onset-driven chops between them
    sr = 22050
    y = np.zeros(sr * 3, dtype=np.float32)
    rng = np.random.default_rng(5)
    for t in (0.4, 1.2, 2.0):                  # bursts at 0.4/1.2/2.0 s
        i = int(t * sr)
        y[i:i + int(0.25 * sr)] = (rng.standard_normal(int(0.25 * sr)) * 0.4).astype(np.float32)
    chops = L.analyze_chops(y, sr, top_n=8)
    assert chops, "should find onset-to-onset chops"
    for c in chops:
        assert c.bars == 0                     # bars=0 marks a chop, not a bar-loop
        assert 0.05 <= (c.end_s - c.start_s) <= 2.0


def test_verify_rerank_demotes_clicking_wrap_below_seamless():
    # #19 regression: two candidates with EQUALLY good proxy features (structural,
    # novelty, proxy seam), but one wraps seamlessly and the other wraps on a
    # discontinuity that SURVIVES the zero-crossing snap. The verify pass measures the
    # true post-snap seam and must rank the seamless one above the clicking one — even
    # though before verify they scored identically. This is exactly the divergence seen
    # on the Super freak dub kit (proxy-seam≈1.0 windows that clicked after B1 snap).
    sr = 22050
    rng = np.random.default_rng(1)
    T = np.arange(int(8.0 * sr)) / sr
    # continuous percussive texture so the zc-snap has real neighbours to snap to
    y = (0.3 * np.sin(2 * np.pi * 110 * T)
         + 0.1 * rng.standard_normal(len(T))).astype(np.float32)
    # window B ([4,6) s) gets a sustained cliff over its last 40 ms — wider than the
    # ±10 ms snap window, so zc-snap cannot escape it: its wrap genuinely clicks.
    cliff = int(0.04 * sr)
    y[int(6.0 * sr) - cliff:int(6.0 * sr)] += 3.0

    cA = L.LoopCandidate(start_s=1.0, end_s=3.0, bars=2, bpm=120,          # seamless
                         tempo_unstable=False, score=0.700, structural=0.6,
                         novelty=0.6, seam=1.0, zc=0.8)
    cB = L.LoopCandidate(start_s=4.0, end_s=6.0, bars=2, bpm=120,          # clicks at wrap
                         tempo_unstable=False, score=0.700, structural=0.6,
                         novelty=0.6, seam=1.0, zc=0.8)                    # SAME proxy seam

    out = L._verify_rerank([cA, cB], y, sr, L.W, top_n=8)
    a = next(c for c in out if c.start_s == 1.0)
    b = next(c for c in out if c.start_s == 4.0)
    assert a.seam > b.seam, "clicking window's verified seam must drop below the clean one"
    assert a.score > b.score, "seamless window must outrank the clicking one after verify"
    assert b.seam < L.VERIFY_SEAM_FLOOR, "the click should fall below the demote floor"
    assert b.score < a.score - 0.5, "a sub-floor click is hard-demoted well below clean"


def test_verify_off_reproduces_raw_proxy_ranking():
    # verify=False must leave the proxy seam untouched (autotune-baseline reproducibility)
    sr = 22050
    t = np.arange(sr * 8) / sr
    y = (0.3 * np.sin(2 * np.pi * 110 * t)).astype(np.float32)
    times = np.arange(0, 8, 0.5)
    peaks = (times * sr / 512).astype(int)
    bpm = np.full(len(times), 120.0)
    raw = L.analyze_stem(y, sr, bars=(2,), top_n=8, grid=(peaks, times, bpm), verify=False)
    # proxy seam is measured on the un-snapped window; just assert it ran and produced
    # candidates with seam populated (not mutated by a verify pass)
    assert raw and all(0.0 <= c.seam <= 1.0 for c in raw)


def test_sparse_percussion_clean_loop_not_flagged_clicking():
    # finding-2 regression: a clean, SPARSE percussion loop (mostly silence + a few
    # hits, wrapping at zero) must NOT be flagged as clicking. The old mean-|2nd-diff|
    # denominator collapsed to ≈0 on sparse material and blew the click ratio up
    # (129× on a genuinely clean loop); the 90th-percentile denominator keeps it sane.
    from engine import grade as G
    sr = 22050
    y = np.zeros(sr, dtype=np.float32)            # 1 s, silent baseline
    # four short decaying hits (sparse dub kick/rim feel), all starting & ending at 0
    for onset in (0.0, 0.25, 0.5, 0.75):
        i = int(onset * sr)
        L_hit = int(0.06 * sr)
        env = np.exp(-np.linspace(0, 8, L_hit)).astype(np.float32)
        y[i:i + L_hit] += (env * np.sin(2 * np.pi * 90 * np.arange(L_hit) / sr)).astype(np.float32)
    # window wraps in silence (x[0]=x[-1]=x[-2]=0) → curvature at wrap is ~0 = seamless
    score, click_db = G.seam_score(y)
    assert score > 0.6, f"clean sparse loop wrongly flagged: seam={score:.2f} click_db={click_db:.1f}"
    assert click_db < G.THRESH["seam_db_click"], "click_db should sit below the click threshold"
    g = G.grade_array(y, sr)
    assert "seam-click" not in g.flags


def test_beat_grid_recovers_tempo_not_its_octave():
    # 120 bpm click train; the fix must report ~120, NOT ~240
    sr = 22050
    spb = 0.5
    n = int(sr * spb * 16)
    y = np.zeros(n, dtype=np.float32)
    for b in range(16):
        i = int(b * spb * sr)
        y[i:i + 200] += np.hanning(200).astype(np.float32)
    _, times, bpm = L.beat_grid(y, sr)
    assert len(times) >= 8
    med = float(np.median(bpm))
    assert 100 <= med <= 140                       # not the 240 octave


def test_rotation_checks_its_own_work_and_repeats():
    """One rotation pass is not enough on real music. The roll point is zero-crossing
    snapped by up to ~10 ms against a 20 ms attack window, which is enough to flip the
    winner between two near-tied slots — so the loop can come out starting on a weak
    beat, the "timed ok but cut BCDA/DABC" failure this function exists to prevent.
    Measured over 60 shipped bar-loops: slot 0 was the best slot in 29 and in the worse
    half in 10 after one pass; 46 and 1 with the check-and-repeat.
    """
    import numpy as np
    from engine import loops as L

    sr, bpm, bars = 44100, 120.0, 4
    nb = bars * 4
    n = int(round(bars * 4 * 60.0 / bpm * sr))
    beat = n // nb
    y = np.zeros(n)
    for b in range(nb):                       # strongest kick deliberately on slot 4
        i = b * beat
        env = np.exp(-np.arange(beat) / (sr * 0.04))
        y[i:i + beat] += (np.sin(2 * np.pi * 55 * np.arange(beat) / sr) * env
                          * (1.0 if b == 4 else 0.35))
    clip = y.astype(np.float32)[:, None]

    rot, slot, _ = L._rotate_to_downbeat(clip, sr, nb)
    assert slot == 4                                       # found the real downbeat
    assert L._downbeat_miss(rot, sr, nb) <= 0.02           # and it now sits at 0
    # rotation is free: a bar-exact loop keeps its length and its wrap
    assert rot.shape == clip.shape
    # idempotent — a rotated loop is already on its downbeat
    _, slot2, _ = L._rotate_to_downbeat(rot, sr, nb)
    assert slot2 == 0


def test_the_downbeat_scorer_is_shared_so_a_check_cannot_disagree():
    """The rotation and its verification must use the same scorer; a lookalike with a
    different window produces disagreements that are the measurement's fault."""
    import numpy as np
    from engine import loops as L
    sr, nb = 44100, 8
    rs = np.random.RandomState(0)
    y = (rs.randn(sr * 2) * 0.1).astype(np.float32)
    s = L._downbeat_slot_scores(y, sr, nb)
    assert s.shape == (nb,)
    assert 0.0 <= L._downbeat_miss(y, sr, nb) <= 1.0
