#!/usr/bin/env python3
"""Sample TF-IDF across the .tidal corpus — derive which samples *identify* a track.

The idea (PLN): treat each `.tidal` score as a document and each Dirt-Samples
folder name as a term. Then:
  - **DF / IDF** tells us how discriminative a sample is. `jungle_breaks` is
    everywhere (low IDF, useless for ID); `cpluck` is rare (high IDF, a *tell*).
  - **TF-IDF per track** ranks each track's signature samples.

This is the principled replacement for hand-picking fingerprint features, and the
weighting layer for locate-matrix L3 (orbit/sample match weighted by IDF, so common
percs don't dominate). Pure text — reads the local git corpus only, never freebox.

Vocabulary = folder names under the *local* Dirt-Samples (custom packs are symlinked
in, e.g. jungle_breaks → MethLab/…). Membership in that set separates samples from
synths (moog, supersaw), control params (gain, room, n), and mininotation noise.

Usage:
    sample_tfidf.py                      # report + write JSON
    sample_tfidf.py --track punkachien   # focus one track's signature
    sample_tfidf.py --sample cpluck      # where is this sample used + how rare
"""
import argparse, math, re, sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "armada" / "tide-table"))
from models import TfidfReport  # noqa: E402

CORPUS = Path("/home/pln/Work/Sound/Tidal")
SC = Path("/home/pln/.local/share/SuperCollider")
DIRT = SC / "downloaded-quarks" / "Dirt-Samples"
# synthdef sources — a synth is a sound source too (PLN: same class for fingerprinting)
SYNTH_DIRS = [SC / "synthdefs", SC / "downloaded-quarks" / "SCLOrkSynths" / "SynthDefs"]
SUPERDIRT_BUILTIN = {"superpiano", "supermandolin", "supergong", "superpwm",
    "superhammond", "supersaw", "supersquare", "super808", "superchip",
    "superhoover", "superzow", "supernoise", "superreese", "superfork",
    "supercomparator", "supervibe", "soskick", "sossnare", "default"}
OUT = CORPUS / "armada" / "tide-table" / "sample_tfidf.json"

# split a quoted mininotation string into candidate tokens. KEEP '_' and digits
# (vec1_acid, jungle_breaks, breaks165); split on whitespace + mininotation ops;
# ':' splits off the sample index (cpluck:1 -> cpluck).
SPLIT = re.compile(r"[\s\[\](){}<>*/.~?!@|,;:%+\-]+")
# Sound-CONTEXT only: a string is a sample pattern iff it's the arg of `s`/`sound`
# or the `# "..."` sound-shorthand. This excludes mask/struct booleans ("t f"),
# note patterns (n/note "..."), MIDI ctrl ("^49"), range args, etc. — which would
# otherwise collide with Dirt folder names ("f", "d", "e") and pollute the signal.
SOUND_CTX = re.compile(r'(?:\bsound\b|\bs\b|#)\s*"([^"]*)"')


def sound_vocab():
    """Authoritative sound-source names = Dirt-Samples folders ∪ synthdef names.
    Returns (vocab_set, kind_map) where kind ∈ {'sample','synth'}."""
    kind = {}
    for p in DIRT.iterdir():
        if not p.name.startswith("."):
            kind[p.name] = "sample"
    for d in SYNTH_DIRS:
        if d.exists():
            for p in d.rglob("*.scd"):
                kind.setdefault(p.stem, "synth")   # don't override a sample name
    for s in SUPERDIRT_BUILTIN:
        kind.setdefault(s, "synth")
    return set(kind), kind


def samples_in(text, vocab):
    """Multiset of sound tokens (samples ∪ synths) in one .tidal, sound-context only."""
    counts = Counter()
    for q in SOUND_CTX.findall(text):
        for tok in SPLIT.split(q):
            if len(tok) > 1 and tok in vocab:   # drop 1-char note/bool collisions
                counts[tok] += 1
    return counts


def build():
    voc, kind = sound_vocab()
    files = sorted(CORPUS.rglob("*.tidal"))
    docs = {}           # rel_path -> Counter(sample -> tf)
    df = Counter()      # sample -> # docs containing it
    for f in files:
        try:
            text = f.read_text(errors="ignore")
        except Exception:
            continue
        c = samples_in(text, voc)
        if not c:
            continue
        rel = str(f.relative_to(CORPUS))
        docs[rel] = c
        for s in c:
            df[s] += 1
    n = len(docs)
    # smoothed idf
    idf = {s: round(math.log((n + 1) / (d + 1)) + 1, 3) for s, d in df.items()}
    tracks = {}
    for rel, c in docs.items():
        tfidf = {s: round(tf * idf[s], 3) for s, tf in c.items()}
        top = sorted(tfidf.items(), key=lambda kv: -kv[1])[:6]
        tracks[rel] = {"n_samples": len(c), "tf": dict(c),
                       "top_tfidf": [{"sample": s, "score": v, "df": df[s],
                                      "kind": kind.get(s, "?")} for s, v in top]}
    used_kinds = {s: kind.get(s, "?") for s in df}
    return {
        "corpus": str(CORPUS), "n_docs": n, "vocab_size": len(voc),
        "n_samples": sum(1 for k in used_kinds.values() if k == "sample"),
        "n_synths": sum(1 for k in used_kinds.values() if k == "synth"),
        "df": dict(df.most_common()),
        "idf": dict(sorted(idf.items(), key=lambda kv: -kv[1])),
        "kinds": used_kinds,
        "tracks": tracks,
    }


def report(data, args):
    n = data["n_docs"]
    df = data["df"]
    idf = data["idf"]
    kinds = data.get("kinds", {})
    K = lambda s: kinds.get(s, "?")
    if args.sample:
        s = args.sample
        if s not in df:
            print(f"'{s}' not used in any .tidal (or not a known sample/synth name).")
            return
        users = [(rel, t["tf"][s]) for rel, t in data["tracks"].items() if s in t["tf"]]
        users.sort(key=lambda x: -x[1])
        print(f"\n■ '{s}' [{K(s)}]  df={df[s]}/{n} docs  idf={idf[s]}  "
              f"({'RARE tell' if df[s] <= 3 else 'common' if df[s] >= 20 else 'mid'})")
        print(f"  used in {len(users)} tracks:")
        for rel, tf in users[:25]:
            print(f"    {tf:>3}×  {rel}")
        return
    if args.track:
        hits = {rel: t for rel, t in data["tracks"].items() if args.track in rel}
        for rel, t in hits.items():
            print(f"\n■ {rel}  ({t['n_samples']} distinct sounds)")
            print("  signature (TF-IDF):")
            for h in t["top_tfidf"]:
                d = df[h["sample"]]
                print(f"    {h['score']:>7}  {h['sample']:<18} {h.get('kind','?'):<6} (df={d}, "
                      f"{'rare' if d <= 3 else 'common' if d >= 20 else 'mid'})")
        if not hits:
            print(f"no track matching '{args.track}'")
        return
    print(f"\nCorpus: {n} .tidal docs · vocab {data['vocab_size']} sound names "
          f"({data.get('n_samples','?')} samples + {data.get('n_synths','?')} synths used)\n")
    rare = [s for s, d in df.items() if d == 1]
    print(f"■ RARE TELLS (df=1, one track only) — {len(rare)} sounds; sample of them:")
    for s in list(df)[::-1][:22]:
        if df[s] <= 2:
            print(f"    df={df[s]}  idf={idf[s]:<6} [{K(s):<6}] {s}")
    print(f"\n■ COMMON / ubiquitous (high df, weak for ID):")
    for s, d in list(df.items())[:16]:
        print(f"    df={d:>4}  idf={idf[s]:<6} [{K(s):<6}] {s}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--track", help="focus one track's signature samples")
    ap.add_argument("--sample", help="where is this sample used + how rare")
    ap.add_argument("--no-write", action="store_true")
    args = ap.parse_args()
    data = build()
    report(data, args)
    if not args.no_write:
        rep = TfidfReport.model_validate(data)   # validate schema on write
        OUT.write_text(rep.model_dump_json(indent=1))
        print(f"\n✓ {OUT}  ({OUT.stat().st_size/1024:.0f} KB)  [validated TfidfReport]")


if __name__ == "__main__":
    main()
