"""End-to-end check of the served openjev classifier: POST the reference NLI pairs to /classify and
compare the returned probabilities with the CPU reference (ref/reference.pt).

  python scripts/check_classify_api.py --url http://localhost:8000 [--model AlexWortega/openjev]
"""
import argparse
import json
import time
import urllib.request

import torch

TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
LABELS = ["contradiction", "entailment", "neutral"]


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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://localhost:8000")
    ap.add_argument("--model", default="AlexWortega/openjev")
    ap.add_argument("--ref", default="/home/ttuser/experiments/openjev/ref/reference.pt")
    ap.add_argument("--batch", type=int, default=20, help="pairs per /classify request")
    args = ap.parse_args()

    ref = torch.load(args.ref)
    texts = ref["texts"]
    rprobs = ref["fp32"]["probs"]
    got = []
    t0 = time.time()
    for i in range(0, len(texts), args.batch):
        chunk = texts[i : i + args.batch]
        t1 = time.time()
        out = post(f"{args.url}/classify", {"model": args.model, "input": chunk})
        dt = time.time() - t1
        print(f"batch {i // args.batch}: {len(chunk)} inputs in {dt:.2f}s ({dt / len(chunk) * 1e3:.0f} ms/input)")
        got.extend(out["data"])
    total = time.time() - t0
    probs = torch.tensor([d["probs"] for d in got], dtype=torch.float32)
    labels = [d["label"] for d in got]
    agree = (probs.argmax(-1) == rprobs.argmax(-1)).float().mean().item()
    dprob = (probs - rprobs).abs()
    print(f"\n{len(texts)} pairs in {total:.2f}s; label agreement vs CPU fp32 reference: {agree:.3f}; "
          f"max|dprob| {dprob.max():.4f} mean|dprob| {dprob.mean():.4f}")
    for i, (p, h) in enumerate(ref["pairs"]):
        print(f"  {i:2d} server={labels[i][:3].upper()} {[round(x, 4) for x in probs[i].tolist()]} "
              f"ref={LABELS[int(rprobs[i].argmax())][:3].upper()} {[round(x, 4) for x in rprobs[i].tolist()]}")
    assert agree == 1.0, "label mismatch vs reference"
    print("OK")


if __name__ == "__main__":
    main()