"""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


def test_the_finder_will_not_spend_its_candidate_slots_on_silence():
    """A stem that plays for a few bars and then rests has a degenerate optimum: every
    term in the composite — structural self-similarity, seam, zero crossings — is
    vacuously perfect across the rest. Measured on BIGHEN's strings stem, all six top
    windows sat inside a gap, so filtering at export left the kit with one loop where it
    should have had four.

    The grid has to come from ANOTHER stem for this to reproduce, which is also how
    `stempack` calls it: a per-stem grid finds no beats in the silence, so no window is
    ever generated there and the bug cannot appear. Share a rhythmic track's grid — the
    whole point of doing so — and every stem gets windows across the full duration,
    including the parts where it is not playing.
    """
    import numpy as np
    from engine import loops as L
    sr = 22050
    bpm, bars = 120.0, 4
    bar_s = 4 * 60.0 / bpm
    def ticks(n_bars, seed):
        r = np.random.RandomState(seed)
        out = np.zeros(int(round(n_bars * bar_s * sr)), dtype=np.float32)
        step = int(round(bar_s / 4 * sr))
        env = np.exp(-np.arange(step) / (sr * 0.03))
        for i in range(0, len(out) - step, step):
            out[i:i + step] += (r.randn(step) * env * 0.4).astype(np.float32)
        return out

    drums = ticks(32, 1)                      # plays throughout: the grid source
    strings = np.concatenate([ticks(8, 2),    # plays for 8 bars …
                              np.zeros(len(drums) - int(round(8 * bar_s * sr)),
                                       dtype=np.float32)])   # … then rests
    assert len(strings) == len(drums)
    grid = L.beat_grid(drums, sr)

    cands = L.analyze_stem(strings, sr, bars=(bars,), top_n=6, grid=grid)
    assert cands, "the played section should still yield candidates"
    played_s = 8 * bar_s
    assert all(c.start_s < played_s for c in cands), \
        [(c.start_s, c.score) for c in cands]
