"""train_synth_classifier.py — Train and export the synth model classifier.

Recipe (chosen by tools/experiment_synth_layers.py --sweep on held-out audio):
LogisticRegression over MERT hidden layer 2, standard-scaled. Held-out
accuracy: 41% top-1 / 69% top-3 across 13 real machines — vs 30-32% / 53-55%
for the nearest-neighbor path the classifier replaces on synth stems.

Trains on ALL labeled embeddings (ref + holdout from the experiment npz):
model selection used the split honestly; the shipped model uses every sample.

Exports plain arrays to training_data/synth_classifier.npz so production
inference is pure numpy (backend/synth_classifier.py) — no sklearn, no pickle.

Usage: python3 tools/train_synth_classifier.py
"""
import os
import sys
from pathlib import Path

import numpy as np

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

ROOT = Path(__file__).resolve().parent.parent
EMB_STORE = ROOT / "training_data" / "synth_layers_experiment.npz"
OUT = ROOT / "training_data" / "synth_classifier.npz"

LAYER = 2
C = 1.0

# model -> family, must mirror build_real_synth_reference.MACHINES
MODEL_FAMILIES = {
    "Roland Juno-106":          "subtractive_analog",
    "Roland Jupiter-8":         "subtractive_analog",
    "Moog Minimoog":            "subtractive_analog",
    "Sequential Prophet-600":   "subtractive_analog",
    "Korg Mono/Poly":           "subtractive_analog",
    "Oberheim Matrix-1000":     "subtractive_analog",
    "Yamaha DX7 (FM family)":   "fm",
    "Yamaha DX200":             "fm",
    "Roland JD-800":            "digital_synth",
    "Roland JV-2080":           "digital_synth",
    "Kawai K4":                 "digital_synth",
    "E-MU Vintage Pro":         "digital_synth",
    "Quasimidi Sirius":         "digital_synth",
    "Yamaha SY35":              "digital_synth",
}


def main():
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler

    d = np.load(EMB_STORE, allow_pickle=True)
    X = np.concatenate([d["X_ref"][:, LAYER], d["X_hold"][:, LAYER]])
    y = np.concatenate([d["y_ref"], d["y_hold"]])
    print(f"training on {len(X)} samples, {len(set(y.tolist()))} classes "
          f"(MERT layer {LAYER})")

    scaler = StandardScaler().fit(X)
    clf = LogisticRegression(max_iter=3000, C=C).fit(scaler.transform(X), y)

    families = [MODEL_FAMILIES.get(c, "digital_synth") for c in clf.classes_]
    np.savez_compressed(
        OUT,
        layer=np.int32(LAYER),
        scaler_mean=scaler.mean_.astype(np.float32),
        scaler_scale=scaler.scale_.astype(np.float32),
        coef=clf.coef_.astype(np.float32),
        intercept=clf.intercept_.astype(np.float32),
        classes=np.array(clf.classes_, dtype=object),
        class_families=np.array(families, dtype=object),
    )
    print(f"exported -> {OUT.name} ({OUT.stat().st_size/1024:.0f} KB)")

    # sanity: round-trip through the pure-numpy inference path
    from backend.synth_classifier import SynthClassifier
    sc = SynthClassifier(OUT)
    preds = sc.predict_proba(X[:5])
    for row, true in zip(preds, y[:5]):
        top = max(row, key=lambda t: t[1])
        print(f"  sanity: true={true:28} pred={top[0]} ({top[1]:.0%})")


if __name__ == "__main__":
    main()