tools/validate_real_drum_reference.py

"""validate_real_drum_reference.py — Held-out validation of drum-machine gear
ID against the MERT reference bank.

Held-out axes vs the reference recipes (build_real_drum_reference.py):
  - one-shot variants: "held" half (odd-indexed files) where pools allow
  - seeds: 7/8/9 (reference used 1/2)
  - same styles (four_floor/electro/breaks) at rng-drawn BPMs

Reports per-machine top-1 model accuracy, top-1 family accuracy, and a
confusion matrix. Also runs a harder single-hit test (isolated kick/snare/hat).

Usage: python3 tools/validate_real_drum_reference.py
"""
import logging
import os
import sys
import tempfile

import numpy as np
import soundfile as sf

logging.basicConfig(level=logging.ERROR)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from backend.mert_embeddings import MERTExtractor, get_reference_db
from tools.real_kits import (SR, discover_kits, render_real_loop,
                             render_solo_kit, pool_overlap, _pool, _load)
from tools.fetch_real_samples import DISPLAY_NAMES

LOOP_CASES = [  # (style, seed)
    ("four_floor", 7), ("four_floor", 8),
    ("electro", 7), ("breaks", 7),
]
SINGLE_HIT_VOICES = ("kick", "snare", "hat_closed")


def _embed(extractor, y):
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf:
        path = tf.name
    try:
        sf.write(path, y, SR)
        return extractor.extract(path)
    finally:
        os.unlink(path)


def main():
    kits = discover_kits()
    db = get_reference_db()
    machines = {e["model"] for e in db.entries if e["family"] == "drum_machine"}
    print(f"Reference DB: {len(db.entries)} entries, {len(machines)} drum machines\n")

    extractor = MERTExtractor()
    confusion: dict = {}
    loop_hits = loop_total = fam_hits = 0

    print("── Loop test (held-out variants, unseen seeds/BPMs) ──")
    for slug, kit in kits.items():
        truth = DISPLAY_NAMES.get(slug, slug)
        overlap = " (single-variant pools: ref/held overlap)" if pool_overlap(kit) else ""
        row_hits = 0
        for style, seed in LOOP_CASES:
            y = render_real_loop(kit, style=style, seed=seed, variant_half="held")
            matches = db.find_nearest(_embed(extractor, y), top_n=3)
            top_model, top_family, top_sim = matches[0]
            loop_total += 1
            fam_hits += int(top_family == "drum_machine")
            ok = top_model == truth
            row_hits += int(ok)
            loop_hits += int(ok)
            confusion.setdefault(truth, {}).setdefault(top_model, 0)
            confusion[truth][top_model] += 1
            if not ok:
                print(f"  ✗ {truth} [{style} s{seed}] -> {top_model} ({top_sim:.0%})")
        print(f"  {truth:<28} {row_hits}/{len(LOOP_CASES)}{overlap}")

    print(f"\nLoop top-1 model:  {loop_hits}/{loop_total} = {loop_hits/loop_total:.0%}")
    print(f"Loop top-1 family: {fam_hits}/{loop_total} = {fam_hits/loop_total:.0%} (drum_machine)")

    print("\n── Confusion (rows=truth, only misses shown above) ──")
    for truth, row in confusion.items():
        parts = ", ".join(f"{m}×{c}" for m, c in
                          sorted(row.items(), key=lambda kv: -kv[1]))
        print(f"  {truth}: {parts}")

    print("\n── Single-hit test (harder: one isolated held-out hit) ──")
    hit_ok = hit_fam = hit_total = 0
    for slug, kit in kits.items():
        truth = DISPLAY_NAMES.get(slug, slug)
        for voice in SINGLE_HIT_VOICES:
            if voice not in kit:
                continue
            pool = _pool(kit[voice], "held")
            y = _load(pool[0])
            matches = db.find_nearest(_embed(extractor, y), top_n=3)
            top_model, top_family, _ = matches[0]
            hit_total += 1
            hit_ok += int(top_model == truth)
            hit_fam += int(top_family == "drum_machine")
    print(f"Single-hit top-1 model:  {hit_ok}/{hit_total} = {hit_ok/hit_total:.0%}")
    print(f"Single-hit top-1 family: {hit_fam}/{hit_total} = {hit_fam/hit_total:.0%}")


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