"""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) -> 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).
    """
    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 _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,
               progress: Optional[Callable[[str], None]] = None) -> list[Take]:
    """Per-stem candidates → loosely grouped Takes (a shared time window).

    One beat grid for the whole catch: computed from the most rhythmic stem (drums if
    present) 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 fix).
    """
    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)
    grid = None
    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_stem(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())
