#!/usr/bin/env python3
"""sample_packs — derive the "pack" granularity ABOVE the loose folder convention.

Three granularities of sample grouping, weakest-binding to strongest:

    per-pack    (rample, 808, commodore…)   ← THIS tool: a real-world origin
    per-folder  (rampleA0, 808bd…)           loose convention; sometimes a kit
    per-sample  (BD0000.wav)                 the ONLY ground truth (sample_resolve)

PLN's insight (#85): the FOLDER is a loose convention, not the true grouping.
`rample*` (rampleA0..rampleS57) is ONE sample pack sliced into banks; `808*`
(808bd/cy/hc/mc) is one drum machine split per voice; `bass_commodore`/`fx_commodore`/
`synth_commodore` are one CHIP sliced by role. Folder maps to family "mostly but not
always" — so a pack maps to ONE audio family even less often, and that mismatch is
itself a story: a pack defined by *origin* (a song, a machine, a console) deliberately
spans the family wheel. We DERIVE packs and COMPARE them to the audio clusters; we
never let a pack boundary bias the audio analysis (it's a hint, never a law).

Reads `sample_families.json` (the grounded L1+L3 resolution) — no audio, fast.

    python3 sample_packs.py            # summary table (packs, cohesion, the mismatches)
    python3 sample_packs.py --json     # write sample_packs.json
    python3 sample_packs.py --auto     # also show auto-proposed packs not in the registry

See [[feedback_sample_grouping_loose]], [[reference_fleet_color_language]]."""
import json
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

HERE = Path(__file__).resolve().parent
FAMILIES = HERE / "sample_families.json"
OUT = HERE / "sample_packs.json"

# ── Curated pack registry ────────────────────────────────────────────────────
# A pack is a real-world ORIGIN that got sliced into several Dirt-Samples folders.
# `match` is a predicate over the folder name. Order matters: first hit wins.
# Types:
#   drum_machine — one box, one folder per voice (808 → bd/cy/hc/mc)
#   bank         — one sample library sliced into indexed banks (rample A0..S57)
#   chip         — one console/computer sound chip, sliced by role (commodore)
#   source       — one song / recording / film sliced into stems (come_*, movie_*)
#   genre        — a themed collection spread over tempo/variant folders (nujazz_*)
# Each entry: (pack, type, matcher). Matcher gets the folder name, returns bool.
def _pre(*prefixes):
    return lambda n: any(n.startswith(p) for p in prefixes)

def _exact(*names):
    s = set(names)
    return lambda n: n in s

REGISTRY = [
    # drum machines / sliced-by-voice
    ("808",        "drum_machine", _pre("808")),
    ("h2ogm",      "drum_machine", _pre("h2ogm")),          # General-MIDI drum kit, per voice
    # sliced sample banks
    ("rample",     "bank",         _pre("rample")),
    ("vec1",       "bank",         _pre("vec1_")),
    ("vec2",       "bank",         _pre("vec2_")),
    # chip / console families, sliced by role (the SAME chip across bass/synth/fx/drums)
    ("commodore",  "chip",         lambda n: n.endswith("_commodore")),
    ("gameboy",    "chip",         lambda n: n.endswith("_gameboy")),
    ("nes",        "chip",         lambda n: n.endswith("_nes")),
    ("atari",      "chip",         lambda n: n.endswith("_atari")),
    # one song / film / recording, sliced into stems
    ("come",       "source",       _pre("come_")),
    ("suns",       "source",       _pre("suns_")),
    ("sunny",      "source",       _pre("sunny_")),
    ("movie",      "source",       _pre("movie_")),
    ("rhadamanthe","source",       _pre("rhadamanthe_")),
    ("nujazz",     "source",       _pre("nujazz_")),
    # themed collections over tempo/variant folders
    ("jungle",     "genre",        lambda n: "jungle" in n or n == "jungbass"),
    ("breaks",     "genre",        _pre("breaks")),
    ("fbreak",     "genre",        _pre("fbreak")),
    ("90s",        "genre",        _pre("90s_")),
    ("trance",     "genre",        _pre("trance_")),
    # vocal collections (different singers/langs, grouped as a vocal-samples pack)
    ("voices",     "source",       _pre("voices_")),
    ("vocal",      "source",       _pre("vocal")),            # vocalOoh1, vocalScatAr/J
    ("dr",         "drum_machine", _exact("dr", "dr2", "dr55", "dr_few")),
]


def assign_pack(folder):
    for pack, typ, match in REGISTRY:
        if match(folder):
            return pack, typ, "registry"
    return None


def _auto_proposals(unpacked):
    """Heuristic packs from shared structure, for folders no registry rule claimed.
    Two axes: prefix-before-digit and underscore tokens (head AND tail). A proposal
    needs ≥2 member folders. These are SUGGESTIONS to review, never authoritative."""
    pre, head, tail = defaultdict(list), defaultdict(list), defaultdict(list)
    for f in unpacked:
        m = re.match(r"^([a-z]+?)\d", f)          # alpha run before first digit
        if m and len(m.group(1)) >= 2:
            pre[m.group(1)].append(f)
        if "_" in f:
            h, _, t = f.partition("_")
            head[h].append(f)
            tail[t].append(f)
    out = {}
    for axis, groups in (("prefix", pre), ("head", head), ("tail", tail)):
        for key, members in groups.items():
            if len(members) >= 2:
                out.setdefault(key, {"axis": axis, "folders": sorted(members)})
    return out


def build(families):
    fam = families["families"]
    packs = {}
    unpacked = []
    for folder, info in fam.items():
        hit = assign_pack(folder)
        if not hit:
            unpacked.append(folder)
            continue
        pack, typ, src = hit
        p = packs.setdefault(pack, {
            "type": typ, "source": src, "folders": [],
            "n_samples": 0, "n_resolved": 0,
            "family_dist": Counter(), "folder_dominants": [],
            "agree": 0, "disagree": 0, "na": 0,
        })
        p["folders"].append(folder)
        p["n_samples"] += info.get("n", 0)
        # roll the per-sample (ground-truth) family counts up to the pack
        for r in info.get("per_index", []):
            if r.get("family"):
                p["family_dist"][r["family"]] += 1
                p["n_resolved"] += 1
        p["folder_dominants"].append(info.get("dominant"))
        agr = info.get("folder_agrees")
        p["agree" if agr is True else "disagree" if agr is False else "na"] += 1

    # A pack needs ≥2 member folders to be a pack: within the currently-grounded folder
    # set a 1-folder pack is just a folder (it may gain siblings when #84 grounds all 718).
    for pack in [k for k, p in packs.items() if len(p["folders"]) < 2]:
        unpacked.extend(packs.pop(pack)["folders"])

    # finalize: cohesion verdict (does the pack live in one audio family or span the wheel?)
    for pack, p in packs.items():
        dist = p.pop("family_dist")
        p["folders"].sort()
        p["n_folders"] = len(p["folders"])
        p["families_spanned"] = dict(dist.most_common())
        n = sum(dist.values())
        if n:
            top, tn = dist.most_common(1)[0]
            p["dominant_family"] = top
            p["dominant_frac"] = round(tn / n, 3)
            p["n_families"] = len(dist)
            p["cohesion"] = "single" if tn / n >= 0.6 and len(dist) <= 2 else "multi"
        else:
            p["dominant_family"], p["dominant_frac"] = None, 0.0
            p["n_families"], p["cohesion"] = 0, "unknown"
        # story line: why this pack does/doesn't map to one family
        fams = ", ".join(f"{k} {v}" for k, v in p["families_spanned"].items())
        if p["cohesion"] == "single":
            p["story"] = f"coheres as {p['dominant_family']} ({fams})"
        else:
            p["story"] = (f"{p['type']} pack spans {p['n_families']} families "
                          f"by design — {fams}")
    return packs, sorted(unpacked)


def main():
    args = sys.argv[1:]
    families = json.loads(FAMILIES.read_text())
    packs, unpacked = build(families)
    auto = _auto_proposals(unpacked)

    if "--json" in args:
        OUT.write_text(json.dumps({
            "schema": "sample packs — origin grouping above the folder convention",
            "method": "curated registry over grounded sample_families (per-sample rollup)",
            "n_folders": families["n_folders"],
            "n_packs": len(packs),
            "n_folders_packed": sum(p["n_folders"] for p in packs.values()),
            "provenance": (f"sample_packs.py · REGISTRY + sample_families.json "
                           f"(grounded subset of {families['n_folders']} folders; "
                           f"re-run after #84 grounds all 718)"),
            "packs": packs,
            "unpacked_folders": unpacked,
            "auto_proposals": auto,
        }, indent=1))
        print(f"✓ {OUT.name} — {len(packs)} packs over "
              f"{sum(p['n_folders'] for p in packs.values())} folders")
        return

    print(f"⛵ {len(packs)} packs over {sum(p['n_folders'] for p in packs.values())} "
          f"of {families['n_folders']} folders\n")
    print(f"  {'pack':<13}{'type':<13}{'fld':>4}{'smp':>5}  cohesion  story")
    multi = []
    for pack in sorted(packs, key=lambda k: -packs[k]["n_folders"]):
        p = packs[pack]
        mark = "·" if p["cohesion"] == "single" else "✶"
        print(f"  {pack:<13}{p['type']:<13}{p['n_folders']:>4}{p['n_samples']:>5}"
              f"  {mark} {p['cohesion']:<6}  {p['story']}")
        if p["cohesion"] == "multi":
            multi.append(pack)
    npk = sum(p["n_folders"] for p in packs.values())
    print(f"\n  {len(multi)} packs span multiple audio families (✶) — the origin-vs-sound "
          f"mismatch the viz can tell:\n    {', '.join(multi)}")
    print(f"  {len(unpacked)} folders are singletons (no pack) — folder == origin there.")

    if "--auto" in args and auto:
        print(f"\n  auto-proposals (review, not in registry): {len(auto)}")
        for key, a in sorted(auto.items()):
            print(f"    {key:<14}({a['axis']}) {', '.join(a['folders'])}")


if __name__ == "__main__":
    main()
