#!/usr/bin/env python3
"""Verify the A16W8 trunk (prefill) context binary on real hardware.

The trunk is stateless: one forward per prompt, no KV. So this pushes the per-prompt inputs,
runs qnn-net-run once, pulls the hidden state, and compares against the float reference
(trunk_ref.npz, produced on the AWS box by trunk_ref.py) on:
  - cosine of the last-real-position hidden vs float
  - the argmax next token through the tied softcapped lm_head

Usage:
  python verify_trunk.py --adb-serial <serial> --ref /path/to/trunk_ref.npz
"""
import argparse, pathlib, subprocess, sys, numpy as np

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import hostlib

BASE = "/data/local/tmp/gemma"
STEP = f"{BASE}/tstep"
OUTR = f"{BASE}/tout/Result_0"
SEQ = 128
NEG = hostlib.NEG
L_ = pathlib.Path("/tmp/gemma_trunk"); L_.mkdir(exist_ok=True)


def adb(a, serial, **kw):
    return subprocess.run(["adb"] + (["-s", serial] if serial else []) + a,
                          capture_output=True, text=True, **kw)


def cmask(L):
    j = np.arange(SEQ)[None, :]
    i = np.arange(SEQ)[:, None]
    return np.where((j <= i) & (j < L), 0.0, NEG).astype(np.float32).reshape(1, 1, SEQ, SEQ)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--adb-serial", required=True)
    ap.add_argument("--ref", default="/home/azek/.claude/jobs/db9a6094/tmp/trunk_ref.npz")
    args = ap.parse_args()

    ref = np.load(args.ref)
    n = int(ref["n"][0])
    m = hostlib.HostModel()
    adb(["shell", f"mkdir -p {STEP} {BASE}/tout"], args.adb_serial)

    ok = 0
    coss = []
    for k in range(n):
        ids = [int(x) for x in ref[f"ids_{k}"]]
        L = len(ids)
        hf = ref[f"h_{k}"]
        tf = int(ref[f"tok_{k}"][0])

        # host embeddings for the whole padded window
        padded = ids + [0] * (SEQ - L)
        ie = np.concatenate([m.embeds(t)[0].reshape(1, 1, hostlib.H) for t in padded], axis=1)
        ple = np.concatenate([m.embeds(t)[1].reshape(1, 1, hostlib.NL, hostlib.PLD) for t in padded], axis=1)
        mk = cmask(L)
        files = {
            "inputs_embeds": ie.astype(np.float32),
            "per_layer_inputs": ple.astype(np.float32),
            "position_ids": np.arange(SEQ, dtype=np.int32).reshape(1, SEQ),
            "full_mask": mk,
            "sliding_mask": mk.copy(),
        }
        for name, arr in files.items():
            p = L_ / f"{name}.raw"
            arr.tofile(p)
            r = adb(["push", str(p), f"{STEP}/{name}.raw"], args.adb_serial, timeout=300)
            if r.returncode != 0:
                raise RuntimeError(f"push {name}: {r.stderr}")

        r = adb(["shell", f"sh {BASE}/gate_ondevice_trunk.sh"], args.adb_serial, timeout=600)
        if "TRUNK_OK" not in r.stdout:
            raise RuntimeError(f"trunk net-run failed:\n{r.stdout}\n{r.stderr}")
        adb(["pull", f"{OUTR}/hidden.raw", str(L_ / "hidden.raw")], args.adb_serial, timeout=300)

        hd = np.fromfile(L_ / "hidden.raw", np.float32).reshape(SEQ, hostlib.H)[L - 1]
        cos = float(hd @ hf / (np.linalg.norm(hd) * np.linalg.norm(hf) + 1e-9))
        coss.append(cos)
        td = int(m.logits(hd).argmax())
        hit = (td == tf)
        ok += hit
        print(f"  [{k:2d}] cos={cos:.5f} |hw|={np.linalg.norm(hd):7.2f} |float|={np.linalg.norm(hf):7.2f} "
              f"hw={td:6d} {m.decode([td])!r:16s} float={tf:6d} {m.decode([tf])!r:16s} "
              f"{'OK' if hit else 'MISMATCH'}", flush=True)

    print(f"\n  hidden cos: mean={np.mean(coss):.5f} min={np.min(coss):.5f}", flush=True)
    print(f"  NEXT-TOKEN TOP-1 vs float on HARDWARE: {ok}/{n} ({100.0*ok/n:.1f}%)", flush=True)
    print("TRUNK_VERIFY_DONE", flush=True)


if __name__ == "__main__":
    main()