"""grade — the loop grader (the katana, task #10).

A pure `grade(wav) -> LoopGrade` that applies the empirical A-profile + the B
mechanical rulebook (`research/PHASE2_FINDINGS.md` §2–§3, §5a) to ANY sample. It
answers one question per file: *how good a loop is this, mechanically?* — seam
click, zero-crossing cleanliness, DC offset, bar-length self-consistency, level
sanity, bass mono-compatibility, and the one-shot/loop class.

Two consumers share this code (DRY, `feedback_build_katana_first`):
  • the **tierlist** (task #11) grades existing `Samples/**/*.wav`;
  • the **finder** (task #5) grades hypothetical windows — it reuses the same
    sub-scorers (`seam_score`, `zc_score`, …) on candidate slices, then adds the
    structural/novelty terms it gets from its SSM context.

Runs under system python3 (numpy/scipy/soundfile/librosa already present). No
network, no GPU. The composite weights + thresholds are **v1, provisional** —
validation is task #11 (grade↔.tidal-usage correlation: core-palette kits must
tier high) and task #7. Tune `WEIGHTS`/`THRESH` there, not by guessing here.

    python3 -m engine.grade  path/to/loop.wav [more.wav …]   # → grades as JSON
"""
from __future__ import annotations

import math
from pathlib import Path
from typing import Optional

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

# ── tunables (v1; calibrate in #11/#7, see module docstring) ─────────────────
THRESH = {
    "one_shot_s": 0.75,        # B9: below this it's a one-shot, not a loop
    "dc_ok": 1e-4,             # B2: |mean| above this is flagged
    "dc_zero": 1e-2,           # |mean| at which the DC score bottoms out
    "seam_db_clean": -20,      # wrap step this far below natural motion ⇒ clean seam
    "seam_db_click": 20,       # wrap step this far above natural motion ⇒ audible click
    "zc_tol_ms": 10.0,         # B1: ±10 ms snap tolerance
    "rms_lo_dbfs": -28.0,      # A-profile moderate-level band …
    "rms_hi_dbfs": -12.0,      # … RMS in here is ideal
    "clip_dbfs": -0.1,         # peak at/above this ⇒ clipping
    "bars": (1, 2, 4, 8),      # B-rule bar interpretations to test
    "beats_per_bar": 4,
    "tempo_lo": 60.0,
    "tempo_hi": 200.0,
}

# Composite weights. Loops and one-shots are graded on different criteria — a
# one-shot is never looped, so seam/bar self-consistency don't apply to it.
WEIGHTS = {
    "loop": {"seam": 0.35, "bar": 0.25, "zc": 0.15, "level": 0.10, "dc": 0.10, "bass_mono": 0.05},
    "one_shot": {"level": 0.5, "dc": 0.3, "zc": 0.2},
}

TIERS = [("S", 0.90), ("A", 0.80), ("B", 0.70), ("C", 0.55), ("D", 0.0)]

_EPS = 1e-12


def _clamp01(x: float) -> float:
    return 0.0 if x < 0 else 1.0 if x > 1 else float(x)


def _db(x: float) -> float:
    return 20.0 * math.log10(max(abs(x), _EPS))


# ── data model ───────────────────────────────────────────────────────────────
class LoopGrade(BaseModel):
    """A mechanical loop-quality grade for one audio file."""

    path: str
    kind: str                      # "loop" | "one_shot"
    grade: float                   # composite 0..1
    tier: str                      # S|A|B|C|D
    sub: dict[str, float]          # per-rule sub-scores, each 0..1
    metrics: dict[str, float]      # raw measured values (dur, dbfs, bpm, …)
    flags: list[str] = Field(default_factory=list)  # human-readable defects

    def tier_of(self) -> str:
        for name, lo in TIERS:
            if self.grade >= lo:
                return name
        return "D"


# ── audio loading ─────────────────────────────────────────────────────────────
def load_audio(path: str | Path) -> tuple[np.ndarray, int]:
    """Read any soundfile-supported file → (channels, n) float64, native sr."""
    y, sr = sf.read(str(path), always_2d=True, dtype="float64")
    return y.T, int(sr)            # (n, ch) → (ch, n)


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


# ── pure sub-scorers (reused by the finder on candidate slices) ───────────────
def dc_score(mono: np.ndarray) -> tuple[float, float]:
    """B2 — DC offset. Returns (score 0..1, raw mean)."""
    dc = float(np.mean(mono))
    a = abs(dc)
    if a <= THRESH["dc_ok"]:
        return 1.0, dc
    span = THRESH["dc_zero"] - THRESH["dc_ok"]
    return _clamp01(1.0 - (a - THRESH["dc_ok"]) / span), dc


def seam_score(mono: np.ndarray) -> tuple[float, float]:
    """Loop the file head-to-tail; measure the wrap click as a continuity defect.

    §3 wants *waveform continuity* at the seam, not a single-sample diff. The
    click is a discontinuity in the signal's TRAJECTORY at the wrap: does x[0]
    continue smoothly from x[-1], x[-2]? The second difference there —
    |x[0] − 2·x[-1] + x[-2]| — is the curvature/jerk of the wrap, and it captures
    both a value jump and a slope break. Crucially it does NOT punish a loop for
    wrapping at a zero crossing (the steepest point) the way a step-vs-mean ratio
    would. Normalized by the signal's typical curvature. Returns (score 0..1,
    seam-jerk dB above typical motion): very negative ⇒ seamless.
    """
    if mono.size < 3:
        return 1.0, -60.0
    # curvature at the wrap: the sequence is … x[-2], x[-1], x[0], x[1] …
    jerk = abs(float(mono[0]) - 2.0 * float(mono[-1]) + float(mono[-2]))
    # Reference = the loop's TYPICAL curvature. The mean |2nd diff| collapses to ≈0
    # on sparse material (dub drums: mostly silence, a few hits) — any real wrap
    # curvature then reads as a huge click (a genuinely clean sparse loop measured
    # 129× the mean, a false positive). The 90th percentile of |2nd diff| tracks the
    # signal's active-region curvature and stays finite on sparse input (2–3× on the
    # same loop). Floor it to a fraction of peak amplitude so a near-silent slice
    # can't make the denominator vanish either.
    d2 = np.abs(np.diff(mono, n=2))
    peak = float(np.max(np.abs(mono))) + _EPS
    typical = max(float(np.percentile(d2, 90)) if d2.size else 0.0,
                  1e-3 * peak) + _EPS       # 90th-pct curvature, amplitude-floored
    click_db = _db(jerk / typical)
    clean, click = THRESH["seam_db_clean"], THRESH["seam_db_click"]
    # ≤ clean dB (jerk ≈ natural curvature) ⇒ 1.0 ; ≥ click dB (audible) ⇒ 0
    score = _clamp01((click - click_db) / (click - clean))
    return score, click_db


def _zero_crossings(mono: np.ndarray) -> np.ndarray:
    # sign flips, robust to exact zeros (signbit treats 0 as positive)
    return np.where(np.diff(np.signbit(mono)))[0]


def zc_score(mono: np.ndarray, sr: int) -> tuple[float, float, float]:
    """B1 — endpoint zero-crossing cleanliness.

    Combines two facts: how near the endpoints sit to zero (a clean loop ends at
    a crossing) and how close a real ZC is to snap to. Returns
    (score 0..1, start dist ms, end dist ms).
    """
    peak = float(np.max(np.abs(mono))) + _EPS
    end_amp = (abs(float(mono[0])) + abs(float(mono[-1]))) / (2 * peak)
    zcs = _zero_crossings(mono)
    if zcs.size:
        d0 = float(np.min(np.abs(zcs - 0))) / sr * 1000.0
        d1 = float(np.min(np.abs(zcs - (mono.size - 1)))) / sr * 1000.0
    else:
        d0 = d1 = THRESH["zc_tol_ms"] * 4
    # within tol ⇒ snappable; amplitude-at-ends near zero ⇒ already clean
    snap = 1.0 - min(1.0, ((d0 + d1) / 2) / (THRESH["zc_tol_ms"] * 2))
    score = _clamp01(0.5 * (1.0 - end_amp) + 0.5 * snap)
    return score, d0, d1


def level_score(mono: np.ndarray) -> tuple[float, float, float, bool]:
    """A-profile level sanity + clip guard. Returns (score, rms dBFS, peak dBFS, clipped)."""
    rms = math.sqrt(float(np.mean(mono ** 2)) + _EPS)
    peak = float(np.max(np.abs(mono))) + _EPS
    rms_db, peak_db = _db(rms), _db(peak)
    lo, hi = THRESH["rms_lo_dbfs"], THRESH["rms_hi_dbfs"]
    if lo <= rms_db <= hi:
        score = 1.0
    elif rms_db < lo:                          # too quiet → 0 by -60 dBFS
        score = _clamp01(1.0 - (lo - rms_db) / (60.0 + lo))
    else:                                      # too hot → 0.4 by 0 dBFS
        score = _clamp01(1.0 - 0.6 * (rms_db - hi) / (0.0 - hi))
    clipped = peak_db >= THRESH["clip_dbfs"]
    if clipped:
        score *= 0.5
    return score, rms_db, peak_db, clipped


def bass_mono_score(y: np.ndarray) -> tuple[float, float]:
    """B6 — inter-channel correlation (mono-sum safety). Mono files → 1.0 (N/A)."""
    if y.ndim == 1 or y.shape[0] < 2:
        return 1.0, 1.0
    a, b = y[0], y[1]
    if np.std(a) < _EPS or np.std(b) < _EPS:
        return 1.0, 1.0
    corr = float(np.corrcoef(a, b)[0, 1])
    return _clamp01((corr + 1.0) / 2.0), corr


def bar_consistency_score(mono: np.ndarray, sr: int) -> tuple[float, dict]:
    """§3 tempo self-check — does the duration map to a clean {1,2,4,8}-bar count
    at a musical tempo? Returns (score 0..1, {detected_bpm, best_bars, implied_bpm, tempo_err}).
    """
    dur = mono.size / sr
    info = {"detected_bpm": 0.0, "best_bars": 0, "implied_bpm": 0.0, "tempo_err": 1.0}
    if dur <= 0:
        return 0.0, info

    detected = 0.0
    try:                                       # librosa optional/slow — guard it
        import librosa
        t = librosa.feature.tempo(y=mono.astype(np.float32), sr=sr, aggregate=np.median)
        detected = float(np.atleast_1d(t)[0])
    except Exception:
        detected = 0.0
    info["detected_bpm"] = round(detected, 2)

    bpb = THRESH["beats_per_bar"]
    best = None
    for bars in THRESH["bars"]:
        implied = bars * bpb * 60.0 / dur
        if not (THRESH["tempo_lo"] <= implied <= THRESH["tempo_hi"]):
            continue
        if detected > 0:                       # compare across ±octave
            err = min(abs(implied - detected * f) / (detected * f) for f in (0.5, 1, 2))
        else:                                  # no tempo read → reward musical range only
            err = 0.05
        if best is None or err < best[2]:
            best = (bars, implied, err)
    if best is None:                           # no interpretation lands in-range
        return 0.0, info
    bars, implied, err = best
    info.update(best_bars=bars, implied_bpm=round(implied, 2), tempo_err=round(err, 4))
    score = _clamp01(1.0 - err / 0.10)         # ±2% ≈ 0.8 ; ≥10% ⇒ 0
    return score, info


# ── the grader ────────────────────────────────────────────────────────────────
def grade_array(y: np.ndarray, sr: int, *, path: str = "<array>",
                role: Optional[str] = None) -> LoopGrade:
    """Grade an in-memory signal (ch, n) or (n,). `role` ('bass'…) is a hint only."""
    mono = _mono(y)
    dur = mono.size / sr
    one_shot = dur < THRESH["one_shot_s"]
    kind = "one_shot" if one_shot else "loop"

    s_dc, dc = dc_score(mono)
    s_seam, seam_db = seam_score(mono)
    s_zc, d0, d1 = zc_score(mono, sr)
    s_lvl, rms_db, peak_db, clipped = level_score(mono)
    s_bass, corr = bass_mono_score(y)
    s_bar, bar = bar_consistency_score(mono, sr) if not one_shot else (1.0, {})

    sub = {"seam": s_seam, "bar": s_bar, "zc": s_zc,
           "level": s_lvl, "dc": s_dc, "bass_mono": s_bass}
    w = WEIGHTS[kind]
    grade = sum(w[k] * sub[k] for k in w)

    flags: list[str] = []
    if one_shot:
        flags.append("one-shot")
    if abs(dc) > THRESH["dc_ok"]:
        flags.append("dc-offset")
    if s_seam < 0.4 and not one_shot:
        flags.append("seam-click")
    if not one_shot and bar.get("best_bars", 0) == 0:
        flags.append("off-grid")
    if clipped:
        flags.append("clipping")
    if rms_db < -45:
        flags.append("near-silent")
    if (role == "bass" or y.ndim > 1) and corr < 0.2:
        flags.append("mono-incompatible")

    metrics = {
        "dur_s": round(dur, 3), "sr": sr,
        "channels": 1 if y.ndim == 1 else y.shape[0],
        "rms_dbfs": round(rms_db, 2), "peak_dbfs": round(peak_db, 2),
        "dc": round(dc, 6), "seam_click_db": round(seam_db, 2),
        "zc_start_ms": round(d0, 2), "zc_end_ms": round(d1, 2),
        "interchannel_corr": round(corr, 4), **bar,
    }
    g = LoopGrade(path=path, kind=kind, grade=round(grade, 4), tier="?",
                  sub={k: round(v, 4) for k, v in sub.items()},
                  metrics=metrics, flags=flags)
    g.tier = g.tier_of()
    return g


def grade(path: str | Path, *, role: Optional[str] = None) -> LoopGrade:
    """Grade a file on disk. `role` hint ('bass', 'drums', …) tunes a couple checks."""
    y, sr = load_audio(path)
    return grade_array(y, sr, path=str(path), role=role)


# ── CLI ───────────────────────────────────────────────────────────────────────
def main(argv: Optional[list[str]] = None) -> int:
    import argparse
    import json

    ap = argparse.ArgumentParser(description="grade audio files as loops (the katana)")
    ap.add_argument("wavs", nargs="+", help="audio file(s) to grade")
    ap.add_argument("--role", default=None, help="role hint: bass|drums|other|vocals")
    ap.add_argument("--json", action="store_true", help="emit JSON array")
    a = ap.parse_args(argv)

    grades = []
    for w in a.wavs:
        try:
            grades.append(grade(w, role=a.role))
        except Exception as e:               # one bad file shouldn't sink a batch
            print(f"✗ {w}: {e}")
    if a.json:
        print(json.dumps([g.model_dump() for g in grades], indent=2))
    else:
        for g in grades:
            fl = f"  ⚑ {', '.join(g.flags)}" if g.flags else ""
            print(f"[{g.tier}] {g.grade:.3f} {g.kind:8} {Path(g.path).name}{fl}")
            print(f"      {'  '.join(f'{k}={v:.2f}' for k, v in g.sub.items())}")
    return 0


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