"""validate_real_synth_reference.py — Measure the real-synth reference bank on
held-out audio it has never seen.

Reads training_data/synth_ref_split.json (written by build_real_synth_reference)
and queries the live DB with each holdout file, restricted to the same family
whitelist the analyzer's bass/other stems use — so the number this prints is
the number the product path actually experiences.

Reports per-machine and overall:
  - model top-1 accuracy (exact machine named first)
  - family top-1 accuracy (synthesis type right, the reliable claim)
  - confusion pairs (what got mistaken for what)

Usage: python3 tools/validate_real_synth_reference.py [--per-machine-cap N]
"""
import argparse
import json
import os
import random
import sys
from collections import Counter, defaultdict
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

ROOT = Path(__file__).resolve().parent.parent
PACKS = ROOT / "training_data" / "synth_packs_raw" / "extracted"
SPLIT_MANIFEST = ROOT / "training_data" / "synth_ref_split.json"

# Mirrors analyzer.STEM_FAMILIES for the synth-bearing stems.
SYNTH_FAMILIES = {"bass", "subtractive_analog", "fm", "digital_synth",
                  "keyboard", "piano", "electric_guitar", "acoustic_guitar",
                  "strings", "brass", "woodwind"}
DRUM_FAMILIES = {"drum_machine", "percussion"}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--per-machine-cap", type=int, default=40,
                    help="Max holdout queries per machine (keeps runtime sane).")
    args = ap.parse_args()

    import numpy as np
    from backend.mert_embeddings import MERTExtractor, get_reference_db

    with open(SPLIT_MANIFEST) as f:
        manifest = json.load(f)

    db = get_reference_db()
    extractor = MERTExtractor()
    rng = random.Random(42)

    # Holdout embeddings are deterministic per file — cache them so bank
    # iterations (quota changes, family tweaks) revalidate in seconds.
    cache_path = ROOT / "training_data" / "synth_holdout_emb_cache.npz"
    cache = {}
    if cache_path.exists():
        loaded = np.load(cache_path)
        cache = {k: loaded[k] for k in loaded.files}

    def emb_of(rel):
        key = rel.replace("/", "|")
        if key not in cache:
            cache[key] = extractor.extract(str(PACKS / rel))
        return cache[key]

    overall = Counter()
    confusion = Counter()
    print(f"{'machine':20} {'n':>4} {'model@1':>8} {'model@3':>8} {'family@1':>9}")
    for slug, info in manifest.items():
        model, family = info["model"], info["family"]
        holdouts = [p for role in info["roles"].values() for p in role["holdout"]]
        if not holdouts:
            continue
        rng.shuffle(holdouts)
        holdouts = holdouts[: args.per_machine_cap]
        families = DRUM_FAMILIES if family == "drum_machine" else SYNTH_FAMILIES

        n = m1 = m3 = f1 = 0
        for rel in holdouts:
            if not (PACKS / rel).exists():
                continue
            matches = db.find_nearest(emb_of(rel), top_n=3, families=families)
            if not matches:
                continue
            n += 1
            top3_models = [m for m, _, _ in matches]
            if top3_models[0] == model:
                m1 += 1
            else:
                confusion[(model, top3_models[0])] += 1
            if model in top3_models:
                m3 += 1
            if matches[0][1] == family:
                f1 += 1
        if n == 0:
            continue
        overall["n"] += n
        overall["m1"] += m1
        overall["m3"] += m3
        overall["f1"] += f1
        print(f"{slug:20} {n:>4} {m1/n:>7.0%} {m3/n:>7.0%} {f1/n:>8.0%}")

    np.savez_compressed(cache_path, **cache)

    if overall["n"]:
        print("-" * 54)
        print(f"{'OVERALL':20} {overall['n']:>4} "
              f"{overall['m1']/overall['n']:>7.0%} "
              f"{overall['m3']/overall['n']:>7.0%} "
              f"{overall['f1']/overall['n']:>8.0%}")

    print("\nTop confusions (true -> predicted):")
    for (true, pred), c in confusion.most_common(12):
        print(f"  {c:3}  {true}  ->  {pred}")
    return 0


if __name__ == "__main__":
    sys.exit(main())