"""loops — the loop finder (task #5).

Generate ranked loop candidates from a separated stem, per the phase-2 design
(research/PHASE2_FINDINGS §1, §3, §4). The pipeline, all librosa/scipy/numpy:

  1. beat grid — drift-tolerant PLP (predominant local pulse), NOT beat_track.
  2. beat-synchronous SSM — MFCC(13) ⊕ chroma(12) synced to beats →
     recurrence affinity. The beat axis absorbs setcps drift for free.
  3. Foote checkerboard novelty on the beat-sync SSM (1-bar and 4-bar kernels).
  4. candidate windows at {1,2,4,8} bars between novelty peaks.
  5. score each window (§3 composite) and dedup.
  6. zero-crossing snap at export.

The finder works PER STEM (§0): 1..x ranked candidates each, then loose cross-
stem grouping into Takes (a time window the human auditions together). The
loopability score REUSES the grader's seam/zc sub-scorers (engine.grade, DRY) and
adds the two terms only the finder has context for: structural (does the window
repeat in the SSM?) and boundary novelty (clean segment edges).

Suggestions are never auto-applied; a wrong rank costs one click to ignore.

    python3 -m engine.loops <stem.wav> [--bars 2 4 8] [--top 8]   # rank one stem
"""
from __future__ import annotations

import math
from pathlib import Path
from typing import Callable, Optional

import numpy as np
import soundfile as sf
from pydantic import BaseModel, Field

from . import grade as G
from .model import Catch

BEATS_PER_BAR = 4
EXPORT_SR = 44100                 # B8: SuperDirt/SC rate
ZC_TOL_MS = 10.0                  # B1
XFADE_MS = 20.0                   # B3 (3 ms for drums, applied at export)

# §3 composite weights (provisional; calibrate with #7)
W = {"struct": 0.4, "novel": 0.2, "seam": 0.3, "zc": 0.1, "tempo_penalty": 0.2}

# verify-rerank (task #19): the finder scores the seam on the RAW window, but the
# exported loop is zero-crossing-snapped (B1) — and the snap changes the wrap. On the
# Super freak dub kit several proxy-seam≈1.0 windows CLICKED after the snap (post-snap
# grade C, click_db 15–20) while genuinely seamless windows ranked below them. The
# verify pass re-measures the TRUE post-snap seam on surviving candidates and folds it
# into the rank so the finder's order matches the grader (grade.py) ground truth.
VERIFY_SEAM_FLOOR = 0.15          # true post-snap seam below this ⇒ hard-demote a click
VERIFY_SURVIVOR_MULT = 3          # verify ~3× top_n survivors, then cut to top_n


class LoopCandidate(BaseModel):
    """One scored window within a single stem."""
    start_s: float
    end_s: float
    bars: int
    bpm: float
    tempo_unstable: bool
    score: float
    structural: float
    novelty: float
    seam: float
    zc: float


class StemSlice(BaseModel):
    name: str                     # drums|bass|other|vocals
    start_s: float
    end_s: float
    score: float
    rms_dbfs: float
    is_one_shot: bool


class Take(BaseModel):
    """A bar-window = the unit PLN selects; gathers each stem's slice in it."""
    start_s: float
    end_s: float
    bars: int
    bpm: float
    tempo_unstable: bool
    score: float                  # best stem score in the window
    n_stems: int
    stems: list[StemSlice] = Field(default_factory=list)


# ── pipeline ──────────────────────────────────────────────────────────────────
def beat_grid(y: np.ndarray, sr: int, *, tempo_min=70.0, tempo_max=160.0, hop=512):
    """Drift-tolerant beat times + per-beat local BPM via PLP. (frames, times, bpm).

    PLP's tempo range is pinned to one musical octave and peak-picking enforces a
    minimum inter-beat spacing — without this PLP latches onto the eighth-note
    pulse and reports ~2× tempo (the §7 octave hazard; Loituma read as 250 not 125).
    """
    import librosa
    oenv = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop)
    pulse = librosa.beat.plp(onset_envelope=oenv, sr=sr, hop_length=hop,
                             tempo_min=tempo_min, tempo_max=tempo_max)
    from scipy.signal import find_peaks
    min_dist = max(1, int((60.0 / tempo_max) * sr / hop))
    peaks, _ = find_peaks(pulse, distance=min_dist)
    times = librosa.frames_to_time(peaks, sr=sr, hop_length=hop)
    if len(times) >= 2:
        ibi = np.diff(times)
        bpm = 60.0 / np.clip(ibi, 1e-3, None)
        bpm = np.r_[bpm, bpm[-1]]
    else:
        bpm = np.array([120.0] * len(times))
    return peaks, times, bpm


def beat_sync_ssm(y: np.ndarray, sr: int, beat_frames: np.ndarray) -> np.ndarray:
    """MFCC(13) ⊕ chroma_cqt(12) synced onto the beat grid → recurrence affinity."""
    import librosa
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    feat = np.vstack([librosa.util.normalize(mfcc, axis=1),
                      librosa.util.normalize(chroma, axis=1)])
    synced = librosa.util.sync(feat, beat_frames, aggregate=np.mean)
    if synced.shape[1] < 3:
        return np.zeros((synced.shape[1], synced.shape[1]))
    return librosa.segment.recurrence_matrix(synced, mode="affinity", sym=True)


def _checkerboard(k: int) -> np.ndarray:
    """Gaussian-tapered checkerboard kernel (2k×2k) for Foote novelty."""
    g = np.outer(*[np.exp(-0.5 * (np.linspace(-2, 2, 2 * k)) ** 2)] * 2)
    sign = np.ones((2 * k, 2 * k))
    sign[:k, k:] = sign[k:, :k] = -1
    return g * sign


def foote_novelty(ssm: np.ndarray, k: int) -> np.ndarray:
    """Checkerboard-kernel novelty along the SSM diagonal; peaks = segment edges."""
    N = ssm.shape[0]
    if N < 2 * k + 1:
        return np.zeros(N)
    ker = _checkerboard(k)
    nov = np.zeros(N)
    pad = np.pad(ssm, k, mode="edge")
    for i in range(N):
        nov[i] = float(np.sum(pad[i:i + 2 * k, i:i + 2 * k] * ker))
    nov = np.maximum(nov, 0)
    return nov / (nov.max() + 1e-9)


def _structural(ssm: np.ndarray, s: int, e: int) -> float:
    """Does the window [s,e) recur elsewhere? Mean best off-segment affinity."""
    if e - s < 1 or ssm.shape[0] < 3:
        return 0.0
    block = ssm[s:e].copy()
    block[:, s:e] = 0.0                       # mask its own diagonal band
    return float(block.max(axis=1).mean()) if block.size else 0.0


def _snapped_window(y: np.ndarray, s_smp: int, e_smp: int, sr: int) -> np.ndarray:
    """The exact slice export_take would write: length-preserving zc-snap (B1).

    Mirrors export_take's snap so the verify-rerank measures the window that actually
    ships. The length is held rigid at (e_smp − s_smp) and the whole window slides by one
    common offset onto near-zero edges — identical scheme to `_snap_length_preserving`, so
    duration is preserved and the seam read is export-faithful.
    """
    length = e_smp - s_smp
    a, b = _snap_length_preserving(y, s_smp, sr, length)
    if b <= a:                                # degenerate snap → fall back to raw bounds
        a, b = s_smp, e_smp
    return y[a:b]


def _true_seam(y: np.ndarray, s_smp: int, e_smp: int, sr: int, n_beats: int = 0) -> float:
    """TRUE post-snap seam quality (0..1) of a window, via the grader's seam scorer.

    Reuses grade.seam_score (DRY — no duplicated DSP) on the zc-snapped slice that will
    actually be exported, so the finder's seam term is the same quantity the grader
    reports rather than a raw-window proxy that the snap invalidates.

    Export-faithful under rotation: with the downbeat rotation (`_rotate_to_downbeat`,
    n_beats > 0) the shipped file's looped playback has TWO junctions — the rotation
    point (its new wrap) and the old window wrap (now interior, but it still plays every
    cycle). We return the MIN of both seams so the binding, audible junction drives the
    rank; a rotation can never hide a bad wrap by moving it inside the file.
    """
    snapped = _snapped_window(y, s_smp, e_smp, sr)
    if snapped.size < 3:
        return 1.0
    seam, _ = G.seam_score(snapped)                    # old wrap (interior post-rotation)
    if n_beats > 1:
        rotated, slot, _ = _rotate_to_downbeat(snapped, sr, n_beats)
        if slot != 0:
            seam_rot, _ = G.seam_score(rotated)        # rotation-point wrap of shipped bytes
            seam = min(seam, seam_rot)
    return seam


# Amplitude below which a window has nothing in it. One number with `grade`, by import,
# so the finder and the grader cannot disagree about what silence is.
_EMPTY_AMP = 10.0 ** (G.THRESH["empty_dbfs"] / 20.0)


def analyze_stem(y: np.ndarray, sr: int, *, bars=(2, 4, 8), top_n=8,
                 grid=None, weights=None, verify=True) -> list[LoopCandidate]:
    """Rank loop candidates within one stem.

    `grid` = a shared (peaks, times, bpm) beat grid (e.g. from the drums stem). Passing
    it fixes the §7 finding that PLP on sparse-onset stems (vocals) yields a weak grid
    and poor windows: derive the grid once from a rhythmic source and reuse it here so
    every stem is cut on the same musical bar lines. None ⇒ compute per-stem (CLI case).

    `weights` = optional override of the §3 composite weights W (struct/novel/seam/zc/
    tempo_penalty); None ⇒ module defaults. The autotune harness sweeps these.

    `verify` (default True) = run the verify-rerank pass: after candidate generation the
    seam term of the ~3× top_n survivors is REPLACED by the true post-snap seam (measured
    on the zero-crossing-snapped window that export actually writes), and candidates whose
    true seam falls below VERIFY_SEAM_FLOOR are hard-demoted. This closes the divergence
    where a proxy-seam≈1.0 window clicked after the B1 snap (see module note). Set False
    to reproduce the pre-#19 raw-proxy ranking (autotune baselines).
    """
    w = W if weights is None else {**W, **weights}
    peaks, times, bpm = grid if grid is not None else beat_grid(y, sr)
    nb = len(times)
    if nb < BEATS_PER_BAR * min(bars) + 1:
        return []
    ssm = beat_sync_ssm(y, sr, peaks)         # peaks are onset-envelope frames
    n = ssm.shape[0]
    nov = np.maximum(foote_novelty(ssm, 2), foote_novelty(ssm, 8)) if n else np.zeros(nb)

    cands: list[LoopCandidate] = []
    for B in bars:
        W_beats = B * BEATS_PER_BAR
        for i in range(0, nb - W_beats):
            e = i + W_beats
            if e >= n:
                continue
            s_smp, e_smp = int(times[i] * sr), int(times[e] * sr)
            if e_smp - s_smp < int(0.5 * sr):
                continue
            sl = y[s_smp:e_smp]
            if sl.size < 8:
                continue
            # Presence, BEFORE ranking. Every term below is vacuously perfect on silence
            # — a seam between two silences is immaculate, its zero crossings are trivial,
            # a flat SSM block is maximally self-similar — so the composite's optimum is
            # "no audio" and a stem with long gaps spends all its candidate slots there.
            # Rejecting at export (via `grade`) is not enough: BIGHEN's strings stem had
            # all six of its top windows inside a rest, so the kit went from four silent
            # loops to one real one instead of to four real ones.
            if float(np.max(np.abs(sl))) < _EMPTY_AMP:
                continue
            seam, _ = G.seam_score(sl)
            zc, _, _ = G.zc_score(sl, sr)
            struct = _structural(ssm, i, e)
            novel = float((nov[i] + nov[min(e, n - 1)]) / 2)
            win_bpm = bpm[i:e]
            bmean = float(np.mean(win_bpm)); bspread = float(np.std(win_bpm))
            unstable = bspread > 5.0
            tempo_pen = min(1.0, bspread / 20.0)
            score = (w["struct"] * struct + w["novel"] * novel +
                     w["seam"] * seam + w["zc"] * zc - w["tempo_penalty"] * tempo_pen)
            cands.append(LoopCandidate(
                start_s=round(times[i], 3), end_s=round(times[e], 3), bars=B,
                bpm=round(bmean, 1), tempo_unstable=unstable,
                score=round(max(0.0, score), 4), structural=round(struct, 4),
                novelty=round(novel, 4), seam=round(seam, 4), zc=round(zc, 4)))

    if verify:
        cands = _verify_rerank(cands, y, sr, w, top_n)
    return _dedup(cands, top_n)


def _verify_rerank(cands: list[LoopCandidate], y: np.ndarray, sr: int,
                   w: dict, top_n: int) -> list[LoopCandidate]:
    """Re-score the top survivors' seam on the true post-snap window (see analyze_stem).

    Cheap: only ~VERIFY_SURVIVOR_MULT×top_n highest-proxy candidates are re-measured, each
    on a small already-in-memory slice. For each, the proxy seam term is swapped for the
    true post-snap seam (rebuilding the composite with the same weights), and a candidate
    whose true seam < VERIFY_SEAM_FLOOR is hard-demoted below all verified-clean ones so a
    click can never top the list. Unverified candidates keep their proxy score.

    Bounds come from the candidate's ROUNDED start_s/end_s — the exact values export_take
    slices at — so verify measures the window that ships, not the sub-sample-exact one
    (the 3-dp rounding alone can shift the wrap onto a click: a beat-exact 0.79 seam read
    0.07 once rounded, and that rounded slice is what gets written).
    """
    survivors = sorted(cands, key=lambda x: -x.score)[:VERIFY_SURVIVOR_MULT * top_n]
    for c in survivors:
        s_smp, e_smp = int(c.start_s * sr), int(c.end_s * sr)
        true_seam = _true_seam(y, s_smp, e_smp, sr, n_beats=c.bars * BEATS_PER_BAR)
        # rebuild the composite swapping proxy seam → true seam (same weights)
        base = c.score - w["seam"] * c.seam + w["seam"] * true_seam
        if true_seam < VERIFY_SEAM_FLOOR:     # audible wrap click after snap
            base -= 1.0                       # hard-demote below any clean candidate
        c.seam = round(true_seam, 4)
        c.score = round(max(0.0, base) if true_seam >= VERIFY_SEAM_FLOOR else base, 4)
    return cands


def analyze_chops(y: np.ndarray, sr: int, *, top_n=12, grid=None,
                  min_len_s=0.25, max_len_s=2.0) -> list[LoopCandidate]:
    """Onset-driven SUB-BAR candidates — how PLN actually chops vocal/sample sources.

    The bar-aligned finder (analyze_stem) targets 2/4/8-bar loops, but the corpus GT is
    chop-heavy: short, often sub-bar pieces (word chops, stabs, one-shots) cut between
    transients (#18 pt.2, the rest of the recall gap). This segments by onset and offers
    onset→onset slices (≤ max_len_s), scored on the seam/zc rules a short chop must still
    honour. `grid` is accepted for signature parity with analyze_stem (chops are
    transient-driven, so it's unused) — keeps find_takes' call site uniform.
    """
    import librosa
    onsets = librosa.onset.onset_detect(y=y, sr=sr, units="time", backtrack=True)
    if len(onsets) < 2:
        return []
    cands: list[LoopCandidate] = []
    for i in range(len(onsets) - 1):
        for j in (i + 1, i + 2, i + 3):               # onset→next … skip-two (phrase-length)
            if j >= len(onsets):
                continue
            s, e = float(onsets[i]), float(onsets[j])
            if not (min_len_s <= e - s <= max_len_s):  # below min = a click, above = a loop
                continue
            s_smp, e_smp = int(s * sr), int(e * sr)
            sl = y[s_smp:e_smp]
            if sl.size < 16:
                continue
            seam, _ = G.seam_score(sl)
            zc, _, _ = G.zc_score(sl, sr)
            lvl, _, _, _ = G.level_score(sl)
            # transient strength at the loop-in onset → favour clean attacks
            score = 0.45 * seam + 0.2 * zc + 0.2 * lvl + 0.15
            cands.append(LoopCandidate(
                start_s=round(s, 3), end_s=round(e, 3), bars=0,   # bars=0 ⇒ a chop, not a bar-loop
                bpm=0.0, tempo_unstable=False, score=round(score, 4),
                structural=0.0, novelty=0.0, seam=round(seam, 4), zc=round(zc, 4)))
    return _dedup(cands, top_n)


def _dedup(cands: list[LoopCandidate], top_n: int) -> list[LoopCandidate]:
    """Greedy NMS: keep highest-scoring, drop windows overlapping >50%."""
    out: list[LoopCandidate] = []
    for c in sorted(cands, key=lambda x: -x.score):
        if all(_overlap(c, k) < 0.5 for k in out):
            out.append(c)
        if len(out) >= top_n:
            break
    return out


def _overlap(a: LoopCandidate, b: LoopCandidate) -> float:
    lo, hi = max(a.start_s, b.start_s), min(a.end_s, b.end_s)
    inter = max(0.0, hi - lo)
    return inter / max(1e-6, min(a.end_s - a.start_s, b.end_s - b.start_s))


# ── takes (cross-stem grouping) ───────────────────────────────────────────────
def find_takes(catch: Catch, workspace: Path, *, bars=(2, 4, 8), top_n=8, mode="loops",
               progress: Optional[Callable[[str], None]] = None) -> list[Take]:
    """Per-stem candidates → loosely grouped Takes (a shared time window).

    mode="loops" (default): bar-aligned 2/4/8-bar windows, one beat grid computed from
    the most rhythmic stem (drums) and reused for every stem so vocals/other are cut on
    the same bar lines instead of each stem's own weak PLP grid (the §7/#18 pt.1 fix).
    mode="chops": onset-driven sub-bar slices (#18 pt.2) — how PLN chops vocal/sample
    sources into stabs and word-chops; grid is unused there.
    """
    analyze = analyze_chops if mode == "chops" else analyze_stem
    catch_dir = catch.dir(workspace)
    loaded: dict[str, tuple] = {}
    for st in catch.stems:
        p = catch_dir / st.path
        if p.exists():
            loaded[st.name] = (G._mono(G.load_audio(p)[0]).astype(np.float32), _sr_of(p))

    # shared grid from the rhythmic source (drums → bass → first available); loops only
    grid = None
    if mode != "chops":
        for pick in ("drums", "bass", *loaded):
            if pick in loaded:
                if progress:
                    progress(f"beat grid from {pick}")
                grid = beat_grid(*loaded[pick])
                break

    per_stem: dict[str, list[LoopCandidate]] = {}
    for name, (y, sr) in loaded.items():
        if progress:
            progress(f"analyzing {name}")
        per_stem[name] = analyze(y, sr, top_n=top_n, grid=grid) if mode == "chops" \
            else analyze(y, sr, bars=bars, top_n=top_n, grid=grid)

    # cluster candidates from all stems by start time (within ~1 bar slack)
    flat = [(name, c) for name, cs in per_stem.items() for c in cs]
    flat.sort(key=lambda nc: nc[1].start_s)
    takes: list[Take] = []
    used = [False] * len(flat)
    for i, (name, c) in enumerate(flat):
        if used[i]:
            continue
        slack = c.bars * BEATS_PER_BAR * 60.0 / max(c.bpm, 1) / BEATS_PER_BAR  # ~1 bar
        members = [(j, flat[j]) for j in range(len(flat))
                   if not used[j] and abs(flat[j][1].start_s - c.start_s) <= slack]
        best_per_stem: dict[str, tuple[int, LoopCandidate]] = {}
        for j, (nm, cc) in members:
            if nm not in best_per_stem or cc.score > best_per_stem[nm][1].score:
                best_per_stem[nm] = (j, cc)
        for j, _ in best_per_stem.values():
            used[j] = True
        lead = max(best_per_stem.values(), key=lambda jc: jc[1].score)[1]
        stems = [StemSlice(name=nm, start_s=cc.start_s, end_s=cc.end_s, score=cc.score,
                           rms_dbfs=0.0, is_one_shot=(cc.end_s - cc.start_s) < 0.75)
                 for nm, (j, cc) in best_per_stem.items()]
        takes.append(Take(start_s=lead.start_s, end_s=lead.end_s, bars=lead.bars,
                          bpm=lead.bpm, tempo_unstable=lead.tempo_unstable,
                          score=round(lead.score, 4), n_stems=len(stems), stems=stems))
    takes.sort(key=lambda t: -t.score)
    return takes[:top_n]


def _sr_of(path: Path) -> int:
    return sf.info(str(path)).samplerate


# ── export (B-rules) ──────────────────────────────────────────────────────────
def _snap_zc(y: np.ndarray, idx: int, sr: int) -> int:
    """Snap a boundary to the nearest zero crossing within B1 tolerance."""
    tol = int(sr * ZC_TOL_MS / 1000)
    lo, hi = max(0, idx - tol), min(len(y) - 1, idx + tol)
    seg = y[lo:hi]
    if seg.size < 2:
        return idx
    zc = np.where(np.diff(np.signbit(seg)))[0]
    return int(lo + zc[np.argmin(np.abs(zc - (idx - lo)))]) if zc.size else idx


def bar_len_samples(bars: int, bpm: float, sr: int) -> int:
    """Exact sample length of `bars` bars at `bpm` (BEATS_PER_BAR beats/bar)."""
    if bars <= 0 or bpm <= 0:
        return 0
    return int(round(bars * BEATS_PER_BAR * (60.0 / bpm) * sr))


def _snap_length_preserving(y: np.ndarray, a: int, sr: int, length: int) -> tuple[int, int]:
    """Zero-crossing snap that PRESERVES the loop LENGTH exactly (task: bar-multiple guarantee).

    The naïve export snapped each edge to its own nearest zero crossing independently, so
    the two edges drifted apart by up to ±ZC_TOL_MS each and the exported DURATION deviated
    from an exact bar multiple (measured −4.2 ms on a 4-bar, −11.4 ms on 2-bar loops — both
    edges snapped inward). loopAt absorbs it, but the guarantee we want is
    |exported_dur − bars·bar_len| < 1 ms.

    Fix: keep the length RIGID at `length` (the exact bar-multiple sample count) and slide the
    WHOLE window by a single common offset so both edges land near zero simultaneously. The
    window [a+δ, a+δ+length) is evaluated for δ ∈ ±ZC_TOL_MS and we pick the δ minimizing the
    combined edge amplitude |y[start]| + |y[end]|. Because both edges move by the same δ, the
    length — hence the duration — is exactly preserved (start and end are that many samples
    apart by construction). Returns (start, end) with end - start == length (clamped to y).
    """
    n = len(y)
    if length <= 0 or n < 2:
        b = min(n, a + max(2, length))
        return max(0, min(a, n - 2)), b
    tol = int(sr * ZC_TOL_MS / 1000)
    # candidate common offsets, clamped so both edges stay inside the signal
    lo = max(-tol, -a)
    hi = min(tol, n - length - a)
    if hi < lo:                                   # window barely fits — no room to slide
        a = max(0, min(a, n - length))
        return a, a + length
    best_delta, best_cost = 0, float("inf")
    for delta in range(lo, hi + 1):
        s = a + delta
        e = s + length
        cost = abs(float(y[s])) + abs(float(y[min(e, n - 1)]))
        # prefer the smaller |δ| on ties so we don't wander from the intended bounds
        cost += 1e-9 * abs(delta)
        if cost < best_cost:
            best_cost, best_delta = cost, delta
    s = a + best_delta
    return s, s + length


ROTATE_ATTACK_WIN_S = 0.030       # attack window each side of a beat slot
ROTATE_LOWBAND_HZ = 200.0         # downbeat = strongest LOW-band attack (kick-weighted)


def _downbeat_slot_scores(clip: np.ndarray, sr: int, n_beats: int) -> np.ndarray:
    """Per-slot downbeat score: low-band circular attack + a small full-band term.

    Factored out of `_rotate_to_downbeat` so the rotation can CHECK ITS OWN RESULT with
    the identical scorer instead of a lookalike — measuring a rotation with a slightly
    different window or weighting produces a disagreement that is the measurement's fault
    (`feedback_check_the_instrument_first`).
    """
    from scipy.signal import butter, sosfiltfilt
    c2 = clip if clip.ndim == 2 else clip[:, None]
    n = c2.shape[0]
    mono = c2.mean(axis=1).astype(np.float64)
    sos = butter(4, ROTATE_LOWBAND_HZ / (sr / 2), btype="low", output="sos")
    low = sosfiltfilt(sos, mono)
    w = max(8, int(ROTATE_ATTACK_WIN_S * sr))

    def _rms(x: np.ndarray, a: int, b: int) -> float:      # circular slice RMS
        idx = np.arange(a, b) % n
        return float(np.sqrt(np.mean(x[idx] ** 2) + 1e-12))

    out = np.empty(n_beats)
    for k in range(n_beats):
        p = int(round(k * n / n_beats))
        out[k] = ((_rms(low, p, p + w) - _rms(low, p - w, p))
                  + 0.25 * (_rms(mono, p, p + w) - _rms(mono, p - w, p)))
    return out


def _downbeat_miss(clip: np.ndarray, sr: int, n_beats: int) -> float:
    """How far slot 0 is from being the best slot: 0.0 = it IS the downbeat, 1.0 = worst."""
    s = _downbeat_slot_scores(clip, sr, n_beats)
    rng = float(s.max() - s.min())
    return 0.0 if rng <= 0 else float((s.max() - s[0]) / rng)


def _rotate_to_downbeat(clip: np.ndarray, sr: int, n_beats: int,
                        force_slot: Optional[int] = None,
                        passes: int = 3) -> tuple[np.ndarray, int, int]:
    """Rotate a bar-exact, clean-seam loop so the musical downbeat sits at position 0.

    PLN's ear-finding (2026-07-11): loops came out 'timed ok but cut BCDA/DABC' — the PLP
    beat grid has no downbeat anchor, so a candidate window can start on any beat and the
    strongest hit lands mid-loop (measured: other_skank opened −23 dB below its own peak
    30 ms; bass_deep's strongest onset sat on slot 4/16).

    Rotation is FREE for a verified loop: looped playback of a rotated loop is the same
    audio cycle — the old wrap junction plays interiorly and the new wrap (the rotation
    point) joins samples that were contiguous in the source, i.e. continuous by
    construction. So this only ever changes WHERE the loop starts, never how it wraps.

    Scoring: at each of the n_beats slot positions, the low-passed (ROTATE_LOWBAND_HZ)
    circular attack — RMS just after the slot minus RMS just before — plus a small
    full-band term; the kick-weighted low band anchors "downbeat" to the drum's 1. The
    winning slot becomes position 0; the exact roll offset is zc-snapped (B1) on the mono
    reference so a re-trigger starts at a zero crossing. Length untouched (np.roll).

    `force_slot` skips the scoring and rotates to that beat slot — used to keep stems
    that share one time window LOCKED to a common rotation (the downbeat is derived once
    from the most rhythmic stem of the window and applied to all of them; independent
    per-stem rotation could shift their relative groove by whole beats).

    `passes` (default 3): rotate, then CHECK, and rotate again if slot 0 still is not the
    strongest attack. One pass is not enough on real music, because the roll point is
    zero-crossing snapped (up to ~10 ms, against a 20 ms attack window) and that is enough
    to flip the winner between two near-tied slots. Measured over 60 shipped bar-loops from
    the Fred pack: after one pass, slot 0 was the best slot in 29 and in the WORSE HALF of
    slots in 10 — one loop in six starting on a weak beat, which is the "timed ok but cut
    BCDA/DABC" complaint this function exists to fix. With up to three passes and a
    strictly-improving guard: best slot in 46, worse half in 1, median miss 0.032 → 0.000.

    Iterating is safe because rotation is free (see above): two rolls compose into one, and
    the guard means a pass that does not improve the measured miss is discarded, so this
    cannot oscillate.

    `clip` is (n,) or (n, ch). Returns (rotated, best_slot, roll_samples); slot 0 ⇒ no-op.
    `best_slot`/`roll_samples` describe the FIRST pass, so a caller re-deriving a shared
    rotation with `force_slot` sees what it saw before.
    """
    c2 = clip if clip.ndim == 2 else clip[:, None]
    n = c2.shape[0]
    if n_beats <= 1 or n < n_beats * 8:
        return clip, 0, 0
    mono = c2.mean(axis=1).astype(np.float64)
    if force_slot is None:
        best_k = int(np.argmax(_downbeat_slot_scores(clip, sr, n_beats)))
    else:
        best_k = int(force_slot) % n_beats
    if best_k == 0:
        return clip, 0, 0
    r = _pocket_zc(mono, int(round(best_k * n / n_beats)), sr)
    rotated = np.roll(c2, -r, axis=0)
    return (rotated if clip.ndim == 2 else rotated[:, 0]), best_k, r

    # ── the check-and-repeat pass (see `passes` above) ────────────────────────
    if force_slot is None and passes > 1:
        cur, miss = out, _downbeat_miss(out, sr, n_beats)
        for _ in range(passes - 1):
            if miss <= 0.02:                         # slot 0 already is the downbeat
                break
            nxt, slot, _roll = _rotate_to_downbeat(cur, sr, n_beats, passes=1)
            if slot == 0:
                break
            m = _downbeat_miss(nxt, sr, n_beats)
            if m >= miss:                            # no improvement: keep what we have
                break
            cur, miss = nxt, m
        out = cur
    return out, best_k, r


def _pocket_zc(mono: np.ndarray, idx: int, sr: int, *,
               back_s: float = 0.025, fwd_s: float = 0.002) -> int:
    """Zero crossing in the QUIET POCKET just before an attack at `idx` (circular).

    A rotation junction placed exactly on the downbeat can land mid-transient (the plain
    nearest-zc snap is ±symmetric); the natural cut point is the low-energy dip that
    precedes the hit. Search [idx−back_s, idx+fwd_s] for zero crossings and pick the one
    with the smallest local (2 ms) RMS — the junction then sits in the pre-kick pocket
    and a re-trigger catches the full attack. Falls back to `idx` if no crossing.
    """
    n = mono.size
    lo, hi = idx - int(back_s * sr), idx + int(fwd_s * sr)
    seg = mono[np.arange(lo, hi) % n]
    zc = np.where(np.diff(np.signbit(seg)))[0]
    if not zc.size:
        return idx % n
    w = max(4, int(0.002 * sr))
    best, best_rms = None, np.inf
    for z in zc:
        pos = (lo + int(z)) % n
        local = mono[np.arange(pos - w // 2, pos + w // 2) % n]
        r = float(np.sqrt(np.mean(local ** 2)))
        if r < best_rms:
            best_rms, best = r, pos
    return int(best)


def kit_multiples_check(files: list[Path], bpm: float, *, tol_pct: float = 0.5) -> dict:
    """Per-file duration-in-bars at the shared grid `bpm`; flag any off an integer multiple.

    Makes the bar-multiple guarantee EXPLICIT per forge (task 1b): after export, each written
    loop's duration is expressed in bars at the take's grid BPM and compared to the nearest
    integer bar count. A file deviating by more than `tol_pct` (%) is flagged — with a
    length-preserving snap this should never trip, so a flag means a real off-grid slice.
    Returns {"bpm", "bar_len_s", "files": [{name, dur_s, bars, nearest, dev_pct, ok}], "ok"}.
    """
    bar_len_s = (BEATS_PER_BAR * 60.0 / bpm) if bpm > 0 else 0.0
    rows: list[dict] = []
    all_ok = True
    for p in files:
        try:
            info = sf.info(str(p))
            dur = info.frames / info.samplerate
        except Exception:
            continue
        bars = dur / bar_len_s if bar_len_s > 0 else 0.0
        nearest = max(1, round(bars)) if bar_len_s > 0 else 0
        dev_pct = abs(bars - nearest) / nearest * 100.0 if nearest else 0.0
        ok = bar_len_s <= 0 or dev_pct <= tol_pct
        all_ok = all_ok and ok
        rows.append({"name": p.name, "dur_s": round(dur, 4), "bars": round(bars, 4),
                     "nearest": nearest, "dev_pct": round(dev_pct, 3), "ok": ok})
    return {"bpm": round(bpm, 3), "bar_len_s": round(bar_len_s, 6), "files": rows, "ok": all_ok}


def export_take(take: Take, kit: str, workspace: Path, catch: Catch, *,
                keep=("drums", "bass", "other", "vocals"),
                names: Optional[dict[str, str]] = None, peak_norm=False,
                rotate: bool = True,
                samples_root: Optional[Path] = None) -> list[Path]:
    """Write a Take's stem slices to Samples/<kit>/ honoring B1/B6/B8/B10, then link.

    B1 snap is LENGTH-PRESERVING (see `_snap_length_preserving`): the exact bar-multiple
    sample count is held rigid and the whole window slides by one common offset onto
    near-zero edges, so every exported loop's duration is a bar multiple to <1 ms rather
    than drifting inward per-edge. `export_take_report` wraps this with the kit-multiples
    check that surfaces the per-file deviation numbers.

    `rotate` (default True): bar-loops are rotated to the musical downbeat before writing
    (`_rotate_to_downbeat`) so the loop STARTS on its strongest low-band attack instead of
    an arbitrary grid beat (the BCDA/DABC ear-finding). Duration and wrap quality are
    unaffected by construction; pass rotate=False to keep the raw grid phase. Chops
    (bars=0) are never rotated.
    """
    from . import publish
    out_dir = (samples_root or publish.samples_root()) / kit
    out_dir.mkdir(parents=True, exist_ok=True)
    catch_dir = catch.dir(workspace)
    stem_path = {st.name: catch_dir / st.path for st in catch.stems}
    written: list[Path] = []
    n_beats = take.bars * BEATS_PER_BAR

    def _load(name: str, _cache={}):
        if name not in _cache:
            _cache[name] = sf.read(str(stem_path[name]), always_2d=True, dtype="float32")
        return _cache[name]

    def _slice(sl: StemSlice) -> tuple[np.ndarray, int]:
        y, sr = _load(sl.name)
        a = int(sl.start_s * sr)
        mono_ref = y[:, 0]
        # exact bar-multiple length from the take's grid; fall back to the raw span if the
        # slice has no bar count (chops: bars=0) so those keep the per-edge snap behaviour.
        length = bar_len_samples(take.bars, take.bpm, sr) if take.bars > 0 else 0
        if length > 0:
            a, b = _snap_length_preserving(mono_ref, a, sr, length)  # B1, length-rigid
        else:
            b = int(sl.end_s * sr)
            a, b = _snap_zc(mono_ref, a, sr), _snap_zc(mono_ref, b, sr)
        return y[a:b], sr

    slices = [s for s in take.stems if s.name in keep and
              (stem_path.get(s.name) or Path("/nonexistent")).exists()]

    # downbeat slot per shared time window, derived ONCE from the most rhythmic stem of
    # the window (drums > bass > other > vocals) and applied to every stem in it — so
    # co-exported stems keep their relative groove instead of each rotating to its own 1.
    group_slot: dict[float, int] = {}
    if rotate and take.bars > 0:
        prio = {"drums": 0, "bass": 1, "other": 2, "vocals": 3}
        for sl in sorted(slices, key=lambda s: prio.get(s.name, 9)):
            key = round(sl.start_s, 3)
            if key in group_slot:
                continue
            clip, sr = _slice(sl)
            _, slot, _ = _rotate_to_downbeat(clip, sr, n_beats)
            group_slot[key] = slot

    for i, sl in enumerate(slices):
        clip, sr = _slice(sl)
        if rotate and take.bars > 0:                                  # start on the downbeat
            slot = group_slot.get(round(sl.start_s, 3), 0)
            clip, _slot, _roll = _rotate_to_downbeat(clip, sr, n_beats, force_slot=slot)
        if sl.name == "bass" and clip.shape[1] > 1:                   # B6 mono-sum
            clip = clip.mean(axis=1, keepdims=True)
        clip = clip - clip.mean(axis=0, keepdims=True)                # B2: ship DC-free
        if peak_norm:                                                 # B10: off by default
            pk = np.max(np.abs(clip)) + 1e-9
            clip = clip * (10 ** (-1.0 / 20) / pk)
        from . import naming
        nm = (names or {}).get(sl.name) or naming.suggest_name(i, sl.name)  # NN_role (#13)
        dest = out_dir / f"{nm}.wav"
        sf.write(str(dest), clip, EXPORT_SR if sr == EXPORT_SR else sr,
                 subtype="PCM_24")                                    # B8
        written.append(dest)
    if written:
        try:
            publish.link_kit(kit)
        except Exception:
            pass
    return written


def export_take_report(take: Take, kit: str, workspace: Path, catch: Catch,
                       **kwargs) -> dict:
    """`export_take` + the kit-multiples check, returned together for the result/report path.

    Returns {"written": [Path…], "kit", "multiples": <kit_multiples_check dict>}. The
    server surfaces `multiples` so the bar-multiple guarantee is visible per forge; a
    non-ok row is a genuine off-grid warning.
    """
    written = export_take(take, kit, workspace, catch, **kwargs)
    return {"written": written, "kit": kit,
            "multiples": kit_multiples_check(written, take.bpm)}


# ── CLI ───────────────────────────────────────────────────────────────────────
def _main(argv=None) -> int:
    import argparse
    import json
    ap = argparse.ArgumentParser(description="rank loop candidates in a stem")
    ap.add_argument("stem")
    ap.add_argument("--bars", type=int, nargs="+", default=[2, 4, 8])
    ap.add_argument("--top", type=int, default=8)
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args(argv)
    y = G._mono(G.load_audio(a.stem)[0]).astype(np.float32)
    sr = _sr_of(Path(a.stem))
    cands = analyze_stem(y, sr, bars=tuple(a.bars), top_n=a.top)
    if a.json:
        print(json.dumps([c.model_dump() for c in cands], indent=2))
    else:
        for c in cands:
            print(f"  [{c.score:.3f}] {c.bars}-bar {c.start_s:6.1f}–{c.end_s:6.1f}s "
                  f"@{c.bpm:.0f}bpm  struct={c.structural:.2f} nov={c.novelty:.2f} "
                  f"seam={c.seam:.2f} zc={c.zc:.2f}{' ⚠tempo' if c.tempo_unstable else ''}")
    return 0


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