#!/usr/bin/env python3
"""Decision path through vLLM: accuracy vs. the HF/PyTorch reference, throughput, and a mixed AR+decision call.

python3 scripts/vllm_decisions.py --bundle exports/jev-judge-qwen35-9b-v0.8 --hf-ref reports/hf_probs_9b_t30k.npz \
    --out reports/vllm_decisions_9b.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

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

from jev_judge.data import load_split, target_to_slots  # noqa: E402
from jev_judge.metrics import fmt_summary, per_sample_table, summarize  # noqa: E402
from jev_judge.template import SLOT_RANGES, n_options_for  # noqa: E402


def slot_mask_np(df) -> np.ndarray:
    m = np.zeros((len(df), 24), dtype=bool)
    for i, (k, o) in enumerate(zip(df["kind"], df["options"])):
        s, _ = SLOT_RANGES[k]
        m[i, s : s + n_options_for(k, list(o))] = True
    return m


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bundle", required=True)
    ap.add_argument("--data", default="data")
    ap.add_argument("--split", default="test_set_30k")
    ap.add_argument("--hf-ref", default=None, help="npz with `probs` [N,24] from the PyTorch path on the same split")
    ap.add_argument("--out", required=True)
    ap.add_argument("--gpu-mem", type=float, default=0.85)
    args = ap.parse_args()

    from jev_judge.vllm_engine import VllmJudge

    df = load_split(args.data, args.split)
    rows = [{"kind": k, "state": s, "question": q, "options": list(o), "family": f}
            for k, s, q, o, f in zip(df["kind"], df["state"], df["question"], df["options"], df["family"])]
    t0 = time.time()
    judge = VllmJudge(args.bundle, gpu_memory_utilization=args.gpu_mem)
    load_s = time.time() - t0

    judge.decide(rows[:256])  # warm-up (LoRA load, graph capture for these shapes)
    t0 = time.time()
    p = judge.decide(rows)
    dt = time.time() - t0
    n_tok = int(df["n_tokens"].sum())
    q = np.stack([target_to_slots(k, t) for k, t in zip(df["kind"], df["target"])])
    mask = slot_mask_np(df)
    s = summarize(per_sample_table(p, q, mask, df))
    res = {"bundle": args.bundle, "split": args.split, "n": len(df), "load_s": load_s, "decide_s": dt,
           "rows_per_s": len(df) / dt, "prompt_tok_per_s": n_tok / dt, "metrics": s}
    print(f"[vLLM decisions] {fmt_summary(s)}")
    print(f"[vLLM decisions] {len(df)} rows in {dt:.1f}s -> {len(df)/dt:.0f} rows/s, {n_tok/dt:.0f} prompt tok/s (load {load_s:.0f}s)")
    if args.hf_ref and os.path.exists(args.hf_ref):
        ref = np.load(args.hf_ref)["probs"]
        d = np.abs(p - ref)[mask]
        agree = (np.where(mask, p, -1).argmax(1) == np.where(mask, ref, -1).argmax(1)).mean()
        res["vs_hf"] = {"max_abs_dp": float(d.max()), "mean_abs_dp": float(d.mean()), "p99_abs_dp": float(np.percentile(d, 99)),
                        "argmax_agreement": float(agree)}
        print(f"[vs PyTorch path] max|Δp| {d.max():.4f} · mean|Δp| {d.mean():.5f} · p99 {np.percentile(d, 99):.4f} · argmax agreement {agree:.4f}")

    # one engine call, both heads
    from vllm import SamplingParams

    msgs = [[{"role": "user", "content": x}] for x in ["In one sentence, what is safety stock?",
                                                      "Write a Python one-liner that reverses a string."]]
    gen_prompts = [judge.tok.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in msgs]
    t0 = time.time()
    gen_outs, dec = judge.mixed(gen_prompts, SamplingParams(temperature=0.0, max_tokens=48), rows[:64])
    res["mixed_call_s"] = time.time() - t0
    res["mixed_decisions_match"] = float(np.abs(dec - p[:64]).max())
    res["mixed_generation_sample"] = [o.outputs[0].text[:160] for o in gen_outs]
    print(f"[mixed] 2 AR generations + 64 decisions in one call: {res['mixed_call_s']:.2f}s; "
          f"decisions max|Δp| vs decision-only run {res['mixed_decisions_match']:.2e}")
    for t in res["mixed_generation_sample"]:
        print("   gen:", t.replace("\n", " ")[:140])
    with open(args.out, "w") as f:
        json.dump(res, f, indent=1)
    print("->", args.out)


if __name__ == "__main__":
    main()