#!/usr/bin/env python3
"""Ground sample families by ANALYSIS, not by name (the project's core principle).

A sample-folder NAME is an unreliable family signal: it may be one instrument, a
heterogeneous kit (`jazz:0`=kick, `:1`=snare…), or a demucs grab named after its
SOURCE song (`wap`, `take5`, `the_revolution`). So instead of guessing from the
name, we LISTEN: CLAP (`laion/clap-htsat-unfused`) zero-shot scores each one-shot
against text prompts for the 12 fleet families, and we aggregate per folder.

We VALIDATE on name-confident folders first (kick/snare/vec1_snare = free ground
truth) before trusting the unresolved ones — build the katana, then cut
([[feedback_build_katana_first]]).

    python3 sample_classify.py validate          # accuracy on name-confident folders
    python3 sample_classify.py run [--all] [--limit N]   # classify → sample_families.json
    python3 sample_classify.py one <folder>      # debug one folder (per-file scores)

Emits `sample_families.json`: per folder → {n, dist:{fam:count}, dominant, conf,
homogeneous, source}. Audio I/O is ffmpeg → 48k mono f32 (no librosa/soundfile dep).
The dominant family + homogeneity feed eda_report (replacing the name-heuristic) and
unblock the viz palette story. Provenance recorded (model + as_of) per
[[feedback_metadata_provenance]]."""
import json
import subprocess
import sys
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np

import models as M
import sample_ontology as ONT

HERE = Path(__file__).resolve().parent
DIRT = Path.home() / ".local/share/SuperCollider/downloaded-quarks/Dirt-Samples"
CV = HERE / "catalog_view.json"
OUT = HERE / "sample_families.json"
MODEL = "laion/clap-htsat-unfused"
SR = 48000
MAX_FILES = 12          # per folder, for speed
MAX_SECONDS = 10        # CLAP htsat window; truncate long loops
AUDIO_EXT = {".wav", ".aif", ".aiff", ".flac", ".ogg", ".mp3"}

# One rich prompt per fleet family (CLAP scores audio↔text similarity). Phrasings
# chosen to separate our 12 families; tune against the validation accuracy report.
PROMPTS = {
    "kick":  "a single kick drum or bass drum, a deep low thump",
    "snare": "a single snare drum hit or a hand clap, a backbeat crack",
    "hat":   "a single hi-hat or cymbal, a crisp bright metallic tick or sizzle",
    "perc":  "a tuned hand percussion one-shot: a conga, bongo, tom-tom, tabla or shaker",
    "break": "a drum break loop or breakbeat, a full looped rhythm with kick snare and hats",
    "bass":  "a sustained low bass note, a sub bass or bass guitar, deep and low-pitched",
    "keys":  "a piano, electric piano, rhodes, organ or clean guitar playing a chord",
    "lead":  "a bright monophonic synth lead melody or arpeggio",
    "synth": "a polyphonic synthesizer chord or electronic synth stab",
    "pad":   "a slow sustained pad, drone or ambient wash, long and evolving",
    "fx":    "a sound effect: a riser, sweep, downlifter, noise or impact, not musical",
    "vox":   "a human voice: singing, speech, a vocal phrase or chant",
}
FAMS = list(PROMPTS)

# MODE: "coarse" = the 12 family prompts above; "fine" = the rich descriptor ontology
# (sample_ontology.py), scored per-descriptor then marginalized to family. Set by CLI.
MODE = "fine"
# METHOD: "clap" (CLAP only) · "panns" (AudioSet only) · "ensemble" (mean of both).
METHOD = "clap"


def active_prompts():
    """(prompt_texts, family_per_prompt) for the current MODE."""
    if MODE == "fine":
        return ONT.PROMPTS, ONT.FAMILY_OF
    return list(PROMPTS.values()), list(PROMPTS)


# ── audio I/O (ffmpeg, no python audio deps) ──────────────────────────────────
def load_audio(path):
    """Decode any file → mono float32 @ 48k, truncated to MAX_SECONDS. None on error."""
    try:
        p = subprocess.run(
            ["ffmpeg", "-v", "quiet", "-i", str(path), "-ac", "1", "-ar", str(SR),
             "-t", str(MAX_SECONDS), "-f", "f32le", "-"],
            capture_output=True, timeout=30)
        a = np.frombuffer(p.stdout, dtype=np.float32)
        return a if a.size >= SR // 100 else None     # need ≥10ms
    except Exception:
        return None


def folder_files(name):
    """Existing Dirt-Samples folder for a corpus sound name → its sample files."""
    d = DIRT / name
    if not d.is_dir():
        return None
    files = sorted(p for p in d.iterdir() if p.suffix.lower() in AUDIO_EXT)
    return files or None


# ── CLAP (lazy: only imported when we actually classify) ──────────────────────
_CLAP = {}


def _clap():
    if not _CLAP:
        import torch
        from transformers import ClapModel, ClapProcessor
        model = ClapModel.from_pretrained(MODEL).eval()
        proc = ClapProcessor.from_pretrained(MODEL)
        prompts, fams = active_prompts()
        # ENCODE THE TEXT TOWER ONCE. The prompts (≈88 fine descriptors) never change,
        # so embedding them per folder re-ran RoBERTa every time — a fixed cost that
        # dominated and made batch size irrelevant. We cache the normalized text embeds
        # + logit scale; per-folder forwards then run only the audio tower (∝ n_files).
        ti = proc(text=prompts, return_tensors="pt", padding=True)
        with torch.no_grad():
            # get_text_features returns a model-output object whose .pooler_output IS
            # the projected 512-d joint embedding (verified identical to a full forward
            # to 1e-7); normalize it the same way forward() does.
            te = model.get_text_features(input_ids=ti["input_ids"],
                                         attention_mask=ti["attention_mask"]).pooler_output
            te = te / te.norm(p=2, dim=-1, keepdim=True)
        # index → family, for marginalizing per-descriptor probs up to families
        fam_keys = list(dict.fromkeys(fams))
        idx = torch.tensor([fam_keys.index(f) for f in fams])
        onehot = torch.zeros(len(fams), len(fam_keys)).scatter_(1, idx[:, None], 1.0)
        _CLAP.update(torch=torch, model=model, proc=proc, text_embeds=te,
                     logit_scale=model.logit_scale_a.exp().detach(),
                     onehot=onehot, fam_keys=fam_keys)
    return _CLAP


def clap_vector(path):
    """CLAP family-probability vector {family: prob} for one sample, or None.
    logits_per_audio → softmax over descriptors → marginalize (sum) to families."""
    return clap_vectors([path])[0]


def clap_vectors(paths):
    """Batched CLAP: ONE forward pass for all decodable files (the text tower is
    computed once, so batching N audios is ~Nx fewer forwards than per-file). Decode
    runs on a small thread pool (ffmpeg is I/O-bound). Returns a list aligned to
    `paths`, with None where decode failed."""
    out = [None] * len(paths)
    with ThreadPoolExecutor(max_workers=4) as ex:
        decoded = list(ex.map(load_audio, paths))
    audios = [(i, a) for i, a in enumerate(decoded) if a is not None]
    if not audios:
        return out
    C = _clap()
    torch = C["torch"]
    with torch.no_grad():
        ai = C["proc"](audio=[a for _, a in audios], sampling_rate=SR,
                       return_tensors="pt", padding=True)
        ae = C["model"].get_audio_features(input_features=ai["input_features"]).pooler_output
        ae = ae / ae.norm(p=2, dim=-1, keepdim=True)
        logits = C["logit_scale"] * ae @ C["text_embeds"].t()  # (n_audio, n_text)
        probs = torch.softmax(logits, dim=-1)
        fam = (probs @ C["onehot"]).tolist()                  # (n_audio, n_fam)
    for (i, _), row in zip(audios, fam):
        out[i] = dict(zip(C["fam_keys"], row))
    return out


def classify_files(files):
    """Batched per-file classification → list of (family, conf)|None aligned to `files`.
    CLAP runs as ONE forward across all files (the expensive tower is shared); PANNs,
    when in play, stays per-file (AudioSet model has no batched path here). This is the
    one place per-file audio prediction lives — classify_file and the resolver both
    route through it so batching is never accidentally lost."""
    clap_vs = clap_vectors(files) if METHOD in ("clap", "ensemble") else [None] * len(files)
    if METHOD in ("panns", "ensemble"):
        import sample_panns
    out = []
    for i, f in enumerate(files):
        vecs = []
        if clap_vs[i]:
            vecs.append(clap_vs[i])
        if METHOD in ("panns", "ensemble"):
            pv = sample_panns.family_vector(f)
            if pv:
                vecs.append(pv)
        if not vecs:
            out.append(None)
            continue
        avg = {k: sum(v.get(k, 0.0) for v in vecs) / len(vecs) for k in ONT.FAMILIES}
        fam = max(avg, key=avg.get)
        out.append((fam, round(avg[fam], 3)))
    return out


def classify_file(path):
    """One sample → (family, confidence) or None, per the active METHOD:
    clap (fine/coarse) · panns (AudioSet) · ensemble (mean of both vectors)."""
    return classify_files([path])[0]


def classify_folder(name):
    """Aggregate per-file predictions for a folder → distribution + dominant.
    CLAP is batched across the folder's files (one forward); PANNs stays per-file."""
    files = folder_files(name)
    if not files:
        return None
    files = files[:MAX_FILES]
    dist, confs = Counter(), []
    for r in classify_files(files):
        if not r:
            continue
        fam, conf = r
        dist[fam] += 1
        confs.append(conf)
    n = sum(dist.values())
    if not n:
        return None
    dominant, dn = dist.most_common(1)[0]
    return {
        "n": n,
        "dist": dict(dist.most_common()),
        "dominant": dominant,
        "conf": round(dn / n, 3),                         # fraction of files agreeing
        "mean_file_conf": round(float(np.mean(confs)), 3),
        "homogeneous": dn / n >= 0.6,                      # else it's a kit / mixed
        "source": f"{METHOD}:{MODE}" + (f":{MODEL}" if METHOD != "panns" else ""),
    }


# ── corpus vocabulary (the names we care about) ───────────────────────────────
def corpus_sounds():
    cv = json.load(open(CV))
    S = set()
    for t in cv["tracks"]:
        for s in t.get("score_sounds", []):
            S.add(s)
    return sorted(S)


# ── commands ──────────────────────────────────────────────────────────────────
def cmd_validate():
    """Accuracy on name-confident folders: does CLAP's dominant == the name's family?
    This is the katana test — if it can't recover the obvious ones, don't trust it."""
    sounds = corpus_sounds()
    truth = [(s, M.classify_sample_family(s)) for s in sounds]
    gt = [(s, fam) for s, fam in truth if fam and folder_files(s)]
    print(f"⛵ validate — {len(gt)} name-confident folders with audio present\n")
    ok, total, per = 0, 0, Counter()
    confmat = Counter()
    for i, (s, fam) in enumerate(gt, 1):
        r = classify_folder(s)
        if not r:
            continue
        total += 1
        hit = r["dominant"] == fam
        ok += hit
        per[fam + (":✓" if hit else ":✗")] += 1
        if not hit:
            confmat[f"{fam}→{r['dominant']}"] += 1
        print(f"  [{i:>2}/{len(gt)}] {'✓' if hit else '✗'} {s:<22} name={fam:<6} "
              f"got={r['dominant']:<6} conf={r['conf']} ({r['n']} files)  "
              f"[running {ok}/{total}={ok/total*100:.0f}%]", flush=True)
    print(f"\n  top-1 accuracy: {ok}/{total} = {ok/total*100:.0f}%" if total else "  no data")
    if confmat:
        print("  confusions:", dict(confmat.most_common()))


def cmd_run(all_folders=False, limit=None):
    """Classify corpus folders (default) or every Dirt-Samples folder → json."""
    if all_folders:
        names = sorted(d.name for d in DIRT.iterdir() if d.is_dir())
    else:
        names = [s for s in corpus_sounds() if folder_files(s)]
    if limit:
        names = names[:limit]
    print(f"⛵ classifying {len(names)} folders with CLAP ({MODEL})…")
    out = {}
    for i, name in enumerate(names, 1):
        r = classify_folder(name)
        if r:
            out[name] = r
            tag = "" if r["homogeneous"] else "  ⚠ mixed/kit"
            print(f"  [{i}/{len(names)}] {name:<24} {r['dominant']:<6} "
                  f"conf={r['conf']}{tag}")
    from datetime import datetime  # noqa — only for stamping the artifact
    payload = {"schema": "sample family analysis (clap zero-shot)",
               "model": MODEL, "n_folders": len(out),
               "prompts": PROMPTS, "families": out}
    OUT.write_text(json.dumps(payload, ensure_ascii=False, indent=1))
    homo = sum(1 for v in out.values() if v["homogeneous"])
    print(f"\n✓ {OUT.name}: {len(out)} folders ({homo} homogeneous, "
          f"{len(out)-homo} mixed/kit)")


def cmd_one(name):
    files = folder_files(name)
    if not files:
        sys.exit(f"no Dirt-Samples folder for {name!r} (synthdef or missing)")
    print(f"⛵ {name} — {len(files)} files (showing ≤{MAX_FILES}):\n")
    for f in files[:MAX_FILES]:
        r = classify_file(f)
        print(f"  {f.name:<28} {r[0]:<6} {r[1]:.2f}" if r else f"  {f.name}: (decode fail)")
    r = classify_folder(name)
    print(f"\n  → dominant={r['dominant']} conf={r['conf']} homogeneous={r['homogeneous']}")
    print(f"    dist={r['dist']}")


def main():
    global MODE, METHOD
    args = sys.argv[1:]
    if "--coarse" in args:
        MODE = "coarse"
    if "--fine" in args:
        MODE = "fine"
    if "--method" in args:
        METHOD = args[args.index("--method") + 1]
    print(f"   (method={METHOD}, clap-mode={MODE})")
    cmd = args[0] if args else "validate"
    if cmd == "validate":
        cmd_validate()
    elif cmd == "run":
        cmd_run(all_folders="--all" in args,
                limit=int(args[args.index("--limit") + 1]) if "--limit" in args else None)
    elif cmd == "one" and len(args) > 1:
        cmd_one(args[1])
    else:
        sys.exit("usage: sample_classify.py [validate | run [--all] [--limit N] | one <folder>]")


if __name__ == "__main__":
    main()
