"""kitcheck — the post-export audit: what `grade` cannot see.

`grade.py` answers *is this file mechanically a clean loop* — seam, zero crossings,
DC, level, bar self-consistency. Every one of those can be perfect on a sample that
is musically useless, and this module exists for exactly the failures that survive a
clean grade. It is the substitute for an ear when nobody is listening yet.

Three checks, each closing a hole the rubric has by construction:

**1 · dead bars.** A four-bar loop whose last two bars are the tail of a fade grades
   S: the seam is silent-to-silent, the length is exact. It is still a two-bar loop
   with two bars of nothing after it. Per-bar RMS, flagged when any bar sits far
   below the loudest.

**2 · bar periodicity.** The finder cuts on a beat grid, so its windows are always
   *bar-aligned*; nothing checks they are a *bar-length musical unit*. Onset-envelope
   autocorrelation at the half/whole-bar lag says whether the material actually
   repeats at the claimed rate, or whether four bars of through-composed movement got
   labelled a loop.

**3 · near-duplicates.** A producer pack ships overlapping stems on purpose — ALL
   DRUMS, KIT, MAIN DRUM LOOP and DRUM BREAKS are four views of one groove. Cut the
   best window from each and you ship the same bar four times under four names. A kit
   whose 8 slots hold 3 distinct sounds is worse than a kit of 3.

Plus the two cheap invariants worth asserting rather than assuming: every bar-loop's
duration is an exact bar multiple at its own BPM, and every file is 44.1 kHz.

    python3 -m engine.kitcheck fred_kits.json
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
import soundfile as sf

BAR_TOL_MS = 1.0          # a bar multiple is exact or it is not
DEAD_BAR_DB = 18.0        # a bar this far under the loudest is not carrying material
DUPE_CORR = 0.93          # normalised cross-correlation above this ⇒ same sound

# Bar periodicity is only diagnostic for material that is SUPPOSED to repeat inside
# itself. Measured over 81 bar-loops from the Fred pack, the medians split hard by role:
# drums 0.754, fx 0.549, tonal 0.214, bass 0.173, vox 0.087 — an 8.7x spread between the
# ends. A drum loop that does not repeat at the bar is broken; a sung phrase that does not
# repeat within itself is a sung phrase. Reading one number against one threshold would
# have condemned 12 of 20 vocal loops for the crime of being vocal
# (`feedback_right_lens_per_control`), so the threshold is per family and `None` means
# "measured and reported, never flagged".
PERIODICITY_FLOOR = {
    "drums": 0.30,      # p25 is 0.615 — 0.30 catches only genuine outliers
    "hits": 0.30,
    "fx": 0.05,         # bimodal: rhythmic textures and drones both live here
    "tonal": 0.05,      # a held chord is a legitimate loop with no onset period
    "bass": 0.05,
    "vox": None,        # never flagged
}


@dataclass
class KitCheck:
    path: str
    name: str
    kit: str
    ok: bool = True
    problems: list = field(default_factory=list)
    bar_dev_ms: Optional[float] = None
    bar_rms_db: list = field(default_factory=list)
    periodicity: Optional[float] = None
    dupe_of: Optional[str] = None


def _mono(y):
    return y.mean(axis=1) if y.ndim > 1 else y


def bar_levels(y: np.ndarray, sr: int, bars: int) -> list[float]:
    """RMS in dBFS of each bar of the loop."""
    if bars <= 0:
        return []
    n = len(y) // bars
    return [round(float(20 * np.log10(np.sqrt(np.mean(y[i * n:(i + 1) * n] ** 2)) + 1e-12)), 1)
            for i in range(bars)]


def periodicity(y: np.ndarray, sr: int, bars: int) -> float:
    """How strongly the onset envelope repeats at the loop's own bar period (0..1).

    Autocorrelation of the onset-strength envelope, read at the lag of one bar (and of
    two bars, when the loop is long enough to have them). A rhythmic loop peaks there; a
    through-composed window does not.

    The normalisation matters more than the rest of the function. At lag k only N-k
    frames overlap, so dividing the raw correlation by ac[0] — which sums all N — makes
    every score a function of how much of the loop the lag consumes, and a 1-bar lag
    consumes half of a 2-bar loop but an eighth of an 8-bar one. Measured on a perfectly
    periodic click track, the biased estimator reports 0.418 / 0.623 / 0.754 at 2 / 4 / 8
    bars for identical material; dividing by the overlap count first gives 0.834 / 0.830 /
    0.862, flat as it must be. Without that correction the short loops look arrhythmic
    purely for being short (`feedback_right_lens_per_control`).
    """
    import librosa
    if bars < 2:
        return 1.0                       # a 1-bar loop has no internal bar period
    oenv = librosa.onset.onset_strength(y=y, sr=sr, hop_length=512)
    oenv = oenv - oenv.mean()
    if not np.any(oenv) or len(oenv) < 8:
        return 0.0
    ac = np.correlate(oenv, oenv, mode="full")[len(oenv) - 1:]
    ac = ac / np.arange(len(oenv), 0, -1)      # unbiased: per-lag overlap count
    ac = ac / (ac[0] + 1e-12)
    per_bar = len(oenv) / bars
    best = 0.0
    for k in (1, 2):                      # one bar, and two bars
        lag = int(round(per_bar * k))
        if 0 < lag < len(ac):
            best = max(best, float(ac[max(0, lag - 2):lag + 3].max()))
    return round(max(0.0, min(1.0, best)), 3)


def _fingerprint(y: np.ndarray, sr: int, n: int = 256) -> np.ndarray:
    """A short, tempo-agnostic signature: log-mel energy resampled to n columns.

    Two cuts of the same groove at different lengths must still compare equal, so the
    envelope is normalised in time (resampled to a fixed width) before comparison —
    a raw waveform correlation would call a 2-bar and a 4-bar cut of one break
    unrelated.
    """
    import librosa
    m = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=24, hop_length=512)
    m = librosa.power_to_db(m + 1e-12)
    if m.shape[1] < 2:
        return np.zeros(24 * n)
    idx = np.linspace(0, m.shape[1] - 1, n)
    m = np.vstack([np.interp(idx, np.arange(m.shape[1]), row) for row in m])
    m = m - m.mean()
    return (m / (np.linalg.norm(m) + 1e-12)).ravel()


def _periodicity_floor(family: str) -> Optional[float]:
    """The floor this family's loops are held to, or None if the measure says nothing."""
    return PERIODICITY_FLOOR.get(family, 0.05)


def check_cuts(cuts: list[dict]) -> list[KitCheck]:
    """Audit every exported sample, then cross-check for duplicates within each kit."""
    out: list[KitCheck] = []
    prints: dict[str, list[tuple[str, np.ndarray]]] = {}

    for c in cuts:
        r = KitCheck(path=c["path"], name=c["name"], kit=c["kit"])
        p = Path(c["path"])
        if not p.exists():
            r.ok = False; r.problems.append("missing file"); out.append(r); continue
        y, sr = sf.read(str(p), always_2d=True, dtype="float32")
        m = _mono(y)

        if sr != 44100:
            r.ok = False; r.problems.append(f"sample rate {sr}, expected 44100")

        bars, bpm = c["bars"], c["bpm"]
        if bars > 0 and bpm > 0:
            want = bars * 4 * 60.0 / bpm
            dev = (len(m) / sr - want) * 1000
            r.bar_dev_ms = round(dev, 3)
            if abs(dev) > BAR_TOL_MS:
                r.ok = False
                r.problems.append(f"not a bar multiple: {dev:+.1f} ms off {bars} bars @ {bpm:.1f}")

            r.bar_rms_db = bar_levels(m, sr, bars)
            if r.bar_rms_db:
                lo, hi = min(r.bar_rms_db), max(r.bar_rms_db)
                if hi - lo > DEAD_BAR_DB:
                    r.ok = False
                    dead = [i + 1 for i, v in enumerate(r.bar_rms_db) if hi - v > DEAD_BAR_DB]
                    r.problems.append(f"dead bar(s) {dead}: {hi - lo:.0f} dB spread across bars")

            r.periodicity = periodicity(m, sr, bars)
            floor = _periodicity_floor(c.get("family", ""))
            if floor is not None and r.periodicity < floor:
                r.problems.append(
                    f"weak bar periodicity ({r.periodicity:.2f} < {floor:.2f} for "
                    f"{c.get('family')}) — bar-aligned but not a repeating unit")

        prints.setdefault(c["kit"], []).append((c["name"], _fingerprint(m, sr)))
        out.append(r)

    by_name = {r.name: r for r in out}
    for kit, items in prints.items():
        for i in range(len(items)):
            for j in range(i + 1, len(items)):
                corr = float(items[i][1] @ items[j][1])
                if corr > DUPE_CORR:
                    later = by_name[items[j][0]]
                    if later.dupe_of is None:
                        later.dupe_of = items[i][0]
                        later.problems.append(
                            f"near-duplicate of {items[i][0]} in the same kit (r={corr:.2f})")
    return out


def report(checks: list[KitCheck]) -> str:
    bad = [c for c in checks if c.problems]
    lines = [f"# kit check — {len(checks)} samples, "
             f"{len(checks) - len(bad)} clean, {len(bad)} flagged", ""]
    if not bad:
        lines.append("No problems found.")
    for kit in sorted({c.kit for c in bad}):
        lines.append(f"## `{kit}`")
        for c in [x for x in bad if x.kit == kit]:
            mark = "✗" if not c.ok else "⚠"
            lines.append(f"- {mark} `{c.name}` — " + "; ".join(c.problems))
        lines.append("")
    dupes = [c for c in checks if c.dupe_of]
    lines += ["", f"**Distinct sounds:** {len(checks) - len(dupes)} of {len(checks)} "
              f"({len(dupes)} near-duplicates)"]
    return "\n".join(lines)


def main(argv=None) -> int:
    import argparse
    ap = argparse.ArgumentParser(prog="kitcheck", description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("cuts_json", type=Path, help="fred_kits.json from engine.stempack")
    ap.add_argument("--out", type=Path, default=None)
    a = ap.parse_args(argv)
    cuts = json.loads(a.cuts_json.read_text())
    checks = check_cuts(cuts)
    txt = report(checks)
    print(txt)
    if a.out:
        a.out.write_text(txt)
        (a.out.parent / (a.out.stem + ".json")).write_text(
            json.dumps([c.__dict__ for c in checks], indent=2))
    return 0 if all(c.ok for c in checks) else 1


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