"""npu-forge-ear: train a tiny voice-verifier HEAD on EmbeddingGemma-300m
embeddings. At runtime the embeddings come from FLM's NPU (embed-gemma:300m,
unquantized) and the head is ~20KB of JSON applied in plain JS — a fully
NPU-accelerated 'trained ear' with zero extra runtimes.

  modal run train_ear_head.py
"""
import modal

app = modal.App("npu-forge-ear")
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch==2.5.1", "transformers==4.56.2", "sentence-transformers==5.1.0", "numpy", "huggingface_hub")
    .add_local_dir("C:/Users/Forgemind/Desktop/voice-harness/scratch/verifier-data", remote_path="/data")
)
vol = modal.Volume.from_name("npu-forge-out", create_if_missing=True)


@app.function(image=image, gpu="T4", timeout=2400, volumes={"/out": vol}, secrets=[modal.Secret.from_name("huggingface-token")])
def train():
    import json, numpy as np, torch
    from sentence_transformers import SentenceTransformer

    def load(p):
        return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]

    labels = json.load(open("/data/labels.json"))
    lab2id = {l: i for i, l in enumerate(labels)}
    train_rows, eval_rows, probe_rows = load("/data/train.jsonl"), load("/data/eval.jsonl"), load("/data/probe.jsonl")

    candidates = ["google/embeddinggemma-300m", "unsloth/embeddinggemma-300m", "onnx-community/embeddinggemma-300m-ONNX"]
    model = None
    for cid in candidates:
        try:
            model = SentenceTransformer(cid, device="cuda")
            print("[embed model] " + cid)
            break
        except Exception as e:
            print(f"[skip] {cid}: {str(e)[:120]}")
    if model is None:
        raise RuntimeError("no embeddinggemma variant loadable")

    def embed(rows):
        return np.asarray(model.encode([r["text"] for r in rows], batch_size=64, show_progress_bar=False, normalize_embeddings=True))

    Xtr, Xev, Xpr = embed(train_rows), embed(eval_rows), embed(probe_rows)
    ytr = np.array([lab2id[r["label"]] for r in train_rows])

    # multinomial logistic head, full-batch
    d, k = Xtr.shape[1], len(labels)
    W = torch.zeros(d, k, requires_grad=True)
    b = torch.zeros(k, requires_grad=True)
    Xt = torch.tensor(Xtr, dtype=torch.float32)
    yt = torch.tensor(ytr)
    opt = torch.optim.Adam([W, b], lr=0.05)
    for i in range(400):
        loss = torch.nn.functional.cross_entropy(Xt @ W + b, yt) + 1e-4 * W.pow(2).sum()
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"[head] final loss {loss.item():.4f}, dim={d}")

    def acc(X, rows):
        pred = (torch.tensor(X, dtype=torch.float32) @ W + b).argmax(-1).numpy()
        per = {}
        for r, p in zip(rows, pred):
            s = per.setdefault(r["label"], [0, 0]); s[1] += 1
            if labels[p] == r["label"]: s[0] += 1
        overall = sum(v[0] for v in per.values()) / max(1, len(rows))
        return overall, {l: f"{v[0]}/{v[1]}" for l, v in per.items()}

    ev_overall, ev_per = acc(Xev, eval_rows)
    pr_overall, pr_per = acc(Xpr, probe_rows)

    head = {"labels": labels, "dim": d, "normalize": True,
            "W": W.detach().numpy().round(5).tolist(), "b": b.detach().numpy().round(5).tolist()}
    import os
    os.makedirs("/out/ear-head", exist_ok=True)
    with open("/out/ear-head/head.json", "w") as f:
        json.dump(head, f)
    vol.commit()
    return {"eval_overall": round(ev_overall, 4), "eval_per_class": ev_per,
            "probe_recognized": pr_per, "probe_overall": round(pr_overall, 4),
            "head_kb": round(len(json.dumps(head)) / 1024)}


@app.local_entrypoint()
def main():
    import json
    print(json.dumps(train.remote(), indent=2))