"""MNLI validation-matched accuracy: CPU HF reference vs the served TT model (/classify).

  python scripts/eval_mnli.py --n 500 [--url http://localhost:8000] [--skip-cpu] [--cpu-only]

Writes results/mnli_<n>.json with per-example labels/probs for both sides.
"""
import argparse
import json
import os
import time
import urllib.request
from pathlib import Path

import torch

TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
# MNLI native labels: 0 entailment, 1 neutral, 2 contradiction -> openjev order (0 con, 1 ent, 2 neu)
MNLI_TO_OPENJEV = {0: 1, 1: 2, 2: 0}
CKPT = os.environ.get("HF_MODEL", "/home/ttuser/experiments/openjev/weights/openjev-qwen3.5-4b-nli")


def load_mnli(n, seed=0):
    from datasets import load_dataset

    ds = load_dataset("nyu-mll/multi_nli", split="validation_matched")
    ds = ds.shuffle(seed=seed).select(range(n))
    texts = [TEMPLATE.format(premise=p.strip(), hypothesis=h.strip()) for p, h in zip(ds["premise"], ds["hypothesis"])]
    gold = [MNLI_TO_OPENJEV[int(l)] for l in ds["label"]]
    return texts, gold


def post(url, payload, timeout=1800):
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def run_cpu(texts, bs=16):
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(CKPT)
    tok.padding_side = "right"
    model = AutoModelForSequenceClassification.from_pretrained(CKPT, dtype=torch.bfloat16).eval()
    model.config.get_text_config().pad_token_id = tok.pad_token_id
    out = []
    t0 = time.time()
    with torch.no_grad():
        for i in range(0, len(texts), bs):
            enc = tok(texts[i : i + bs], truncation=True, max_length=4096, padding=True, return_tensors="pt")
            out.append(torch.softmax(model(**enc).logits.float(), -1))
            print(f"  cpu {min(i + bs, len(texts))}/{len(texts)} ({time.time() - t0:.0f}s)", flush=True)
    return torch.cat(out, 0)


def run_server(texts, url, model_name, bs=32):
    out = []
    t0 = time.time()
    for i in range(0, len(texts), bs):
        res = post(f"{url}/classify", {"model": model_name, "input": texts[i : i + bs]})
        out.extend([d["probs"] for d in res["data"]])
        print(f"  server {min(i + bs, len(texts))}/{len(texts)} ({time.time() - t0:.1f}s)", flush=True)
    return torch.tensor(out, dtype=torch.float32), time.time() - t0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=500)
    ap.add_argument("--url", default="http://localhost:8000")
    ap.add_argument("--model", default="AlexWortega/openjev")
    ap.add_argument("--skip-cpu", action="store_true")
    ap.add_argument("--cpu-only", action="store_true")
    ap.add_argument("--batch", type=int, default=32)
    args = ap.parse_args()

    texts, gold = load_mnli(args.n)
    gold_t = torch.tensor(gold)
    res_dir = Path(__file__).resolve().parent.parent / "results"
    res_dir.mkdir(exist_ok=True)
    out_path = res_dir / f"mnli_{args.n}.json"
    result = json.load(open(out_path)) if out_path.exists() else {"n": args.n, "texts": texts, "gold": gold}

    if not args.skip_cpu:
        cpu = run_cpu(texts)
        result["cpu_probs"] = cpu.tolist()
    cpu = torch.tensor(result["cpu_probs"]) if "cpu_probs" in result else None
    if cpu is not None:
        print(f"CPU bf16 accuracy: {(cpu.argmax(-1) == gold_t).float().mean():.4f}")

    if not args.cpu_only:
        tt, dt = run_server(texts, args.url, args.model, bs=args.batch)
        result["tt_probs"] = tt.tolist()
        result["tt_seconds"] = dt
        print(f"TT accuracy: {(tt.argmax(-1) == gold_t).float().mean():.4f}  ({dt:.1f}s, {dt / args.n * 1e3:.0f} ms/example)")
        if cpu is not None:
            print(f"TT vs CPU label agreement: {(tt.argmax(-1) == cpu.argmax(-1)).float().mean():.4f}; "
                  f"max|dprob| {(tt - cpu).abs().max():.4f} mean|dprob| {(tt - cpu).abs().mean():.4f}")
    json.dump(result, open(out_path, "w"))
    print("saved", out_path)


if __name__ == "__main__":
    main()