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


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 analyze_stem(y: np.ndarray, sr: int, *, bars=(2, 4, 8), top_n=8,
                 grid=None, weights=None) -> 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.
    """
    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
            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)))
    return _dedup(cands, top_n)


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 export_take(take: Take, kit: str, workspace: Path, catch: Catch, *,
                keep=("drums", "bass", "other", "vocals"),
                names: Optional[dict[str, str]] = None, peak_norm=False,
                samples_root: Optional[Path] = None) -> list[Path]:
    """Write a Take's stem slices to Samples/<kit>/ honoring B1/B6/B8/B10, then link."""
    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] = []
    for i, sl in enumerate(s for s in take.stems if s.name in keep):
        src = stem_path.get(sl.name)
        if not src or not src.exists():
            continue
        y, sr = sf.read(str(src), always_2d=True, dtype="float32")   # (n, ch)
        a, b = int(sl.start_s * sr), int(sl.end_s * sr)
        mono_ref = y[:, 0]
        a, b = _snap_zc(mono_ref, a, sr), _snap_zc(mono_ref, b, sr)   # B1
        clip = y[a:b]
        if sl.name == "bass" and clip.shape[1] > 1:                   # B6 mono-sum
            clip = clip.mean(axis=1, keepdims=True)
        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


# ── 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())
