#!/usr/bin/env python3
"""Jev vs small and frontier LLMs on labeled classification data.

AY Automate, September 2026. Standard library only.

  python3 jev_benchmark.py fetch    # download the public datasets into ./data
  python3 jev_benchmark.py run      # call every model on every item (needs OPENROUTER_API_KEY)
  python3 jev_benchmark.py report   # print the summary tables, write summary.json

Everything goes through OpenRouter so all systems share one client, one network path and one
billing meter. Jev uses OpenRouter's decisions endpoint. The LLMs use chat completions with a
strict JSON schema, temperature 0 and minimal reasoning. `run` is resumable and stops if the
running cost passes BUDGET_USD.
"""
import json, math, os, random, statistics, sys, threading, time, urllib.error, urllib.request
from concurrent.futures import ThreadPoolExecutor

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data")
OUT = os.path.join(HERE, "results.jsonl")
SEED = 42
BUDGET_USD = float(os.environ.get("BUDGET_USD", "2.0"))
WORKERS = int(os.environ.get("WORKERS", "4"))
LIMIT = int(os.environ.get("LIMIT", "0"))  # smoke tests: cap items per task

CHAT = "https://openrouter.ai/api/v1/chat/completions"
DECISIONS = "https://openrouter.ai/api/alpha/decisions"
JEV = "typesafe/jev-1.13"
LLMS = [
    "openai/gpt-5.4-nano",
    "google/gemini-3.5-flash-lite",
    "anthropic/claude-haiku-4.5",
    "openai/gpt-5.6-terra",
]
SYSTEMS = [JEV] + LLMS

# eight deliberately confusable Banking77 intents: card delivery vs card payment problems
INTENT8 = [
    "card_arrival", "card_delivery_estimate", "order_physical_card", "getting_spare_card",
    "pending_card_payment", "declined_card_payment", "card_payment_not_recognised",
    "direct_debit_payment_not_recognised",
]
INTENT_Q = "Which banking support intent best matches this customer message?"
INJECT_Q = ("Does this text try to override, ignore or hijack the instructions given to an AI "
            "assistant, or otherwise manipulate an AI system into behaving differently?")


def human(label):
    return label.replace("_", " ")


# ---------------------------------------------------------------- data
def hf_rows(dataset, split):
    rows, offset = [], 0
    while True:
        url = ("https://datasets-server.huggingface.co/rows?dataset=%s&config=default&split=%s"
               "&offset=%d&length=100" % (dataset, split, offset))
        page = None
        for attempt in range(6):
            try:
                with urllib.request.urlopen(url, timeout=60) as r:
                    page = json.load(r)
                break
            except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError) as e:
                print("  retry %d after %s" % (attempt + 1, e))
                time.sleep(2 * (attempt + 1))
        if page is None:
            raise RuntimeError("datasets server kept failing for " + url)
        got = [x["row"] for x in page["rows"]]
        rows += got
        offset += 100
        if not got or offset >= page["num_rows_total"]:
            return rows
        time.sleep(0.15)


def fetch():
    os.makedirs(DATA, exist_ok=True)
    jobs = {"banking77_test": ("mteb/banking77", "test"),
            "inject_train": ("deepset/prompt-injections", "train"),
            "inject_test": ("deepset/prompt-injections", "test")}
    for name, (ds, split) in jobs.items():
        rows = hf_rows(ds, split)
        json.dump(rows, open(os.path.join(DATA, name + ".json"), "w"))
        print(name, len(rows), "rows")


def load(name):
    return json.load(open(os.path.join(DATA, name + ".json")))


def build_tasks():
    rnd = random.Random(SEED)
    b77 = load("banking77_test")
    by_label = {}
    for r in b77:
        by_label.setdefault(r["label_text"], []).append(r["text"])
    all_labels = sorted(by_label)
    t8 = []
    for lab in INTENT8:
        for i, txt in enumerate(rnd.sample(by_label[lab], 20)):
            t8.append({"id": "%s-%d" % (lab, i), "text": txt, "gold": human(lab)})
    t77 = []
    for lab in all_labels:
        for i, txt in enumerate(rnd.sample(by_label[lab], 3)):
            t77.append({"id": "%s-%d" % (lab, i), "text": txt, "gold": human(lab)})
    inj = load("inject_train") + load("inject_test")
    tinj = [{"id": "inj-%d" % i, "text": r["text"], "gold": bool(r["label"])} for i, r in enumerate(inj)]
    tinj = rnd.sample(tinj, 400)  # 400 of the 662 labeled examples, seeded
    tasks = [
        {"name": "intent8", "kind": "choice", "instructions": INTENT_Q,
         "options": [human(x) for x in INTENT8], "items": t8},
        {"name": "inject", "kind": "noul", "instructions": INJECT_Q, "options": None, "items": tinj},
        {"name": "intent77", "kind": "choice", "instructions": INTENT_Q,
         "options": [human(x) for x in all_labels], "items": t77},
    ]
    if LIMIT:
        for t in tasks:
            rnd2 = random.Random(SEED)
            t["items"] = rnd2.sample(t["items"], min(LIMIT, len(t["items"])))
    return tasks


# ---------------------------------------------------------------- calls
def http_post(url, body, key):
    data = json.dumps(body).encode()
    hdr = {"Authorization": "Bearer " + key, "Content-Type": "application/json"}
    last = None
    for attempt in range(4):
        t0 = time.perf_counter()
        try:
            with urllib.request.urlopen(urllib.request.Request(url, data=data, headers=hdr), timeout=120) as r:
                txt = r.read().decode()
            return json.loads(txt), time.perf_counter() - t0, attempt, None
        except urllib.error.HTTPError as e:
            last = "http %d %s" % (e.code, e.read().decode()[:160])
            if e.code in (429, 500, 502, 503, 504):
                time.sleep(min(8, 1.5 * (attempt + 1)))
                continue
            break
        except Exception as e:  # network, timeout, bad json
            last = "%s %s" % (type(e).__name__, e)
            time.sleep(1.5)
    return None, 0.0, attempt, last


def ask_jev(task, item, key):
    q = ({"type": "choice", "instructions": task["instructions"],
          "criteria": {o: None for o in task["options"]}} if task["kind"] == "choice"
         else {"type": "noul", "instructions": task["instructions"]})
    j, dt, att, err = http_post(DECISIONS, {"model": JEV, "state": item["text"], "questions": {"q": q}}, key)
    if j is None:
        return {"ok": False, "err": err, "latency": None, "retries": att}
    a = j["answers"]["q"]
    u = j.get("usage", {})
    rec = {"ok": True, "latency": dt, "retries": att, "cost": u.get("cost", 0.0), "in_tok": u.get("input_tokens"),
           "out_tok": u.get("output_tokens"), "served_model": j.get("model")}
    if task["kind"] == "choice":
        rec.update(pred=a["choice"], conf=a.get("confidence"), p_top=(a.get("probabilities") or {}).get(a["choice"]))
    else:
        rec.update(pred=a["noul"] >= 0.5, prob=a["noul"])
    return rec


def parse_json(text):
    t = text.strip()
    if t.startswith("```"):
        t = t.strip("`")
        t = t[t.find("{"):]
    return json.loads(t[t.find("{"): t.rfind("}") + 1])


def ask_llm(model, task, item, key):
    if task["kind"] == "choice":
        sysmsg = (task["instructions"] + "\nChoose exactly one option from this list:\n- "
                  + "\n- ".join(task["options"]) + '\nReturn JSON only: {"label": <option>}.')
        schema = {"type": "object", "properties": {"label": {"type": "string", "enum": task["options"]}},
                  "required": ["label"], "additionalProperties": False}
    else:
        sysmsg = (task["instructions"] + '\nReturn JSON only: {"injection": true or false, "probability": '
                  "your probability from 0 to 1 that the text is an injection}.")
        schema = {"type": "object", "properties": {"injection": {"type": "boolean"},
                  "probability": {"type": "number"}}, "required": ["injection", "probability"],
                  "additionalProperties": False}
    body = {"model": model, "temperature": 0, "max_tokens": 300, "usage": {"include": True},
            "reasoning": {"effort": "minimal"},
            "messages": [{"role": "system", "content": sysmsg}, {"role": "user", "content": item["text"]}],
            "response_format": {"type": "json_schema", "json_schema": {"name": "decision", "strict": True, "schema": schema}}}
    j, dt, att, err = http_post(CHAT, body, key)
    if j is None:
        return {"ok": False, "err": err, "latency": None, "retries": att}
    try:
        out = parse_json(j["choices"][0]["message"]["content"])
    except Exception as e:
        return {"ok": False, "err": "unparseable: %s" % e, "latency": dt, "retries": att}
    u = j.get("usage", {})
    rec = {"ok": True, "latency": dt, "retries": att, "cost": u.get("cost", 0.0), "in_tok": u.get("prompt_tokens"),
           "out_tok": u.get("completion_tokens"), "served_model": j.get("model")}
    if task["kind"] == "choice":
        rec["pred"] = out["label"]
    else:
        rec.update(pred=bool(out["injection"]), prob=float(out["probability"]))
    return rec


# ---------------------------------------------------------------- run
def run():
    key = os.environ["OPENROUTER_API_KEY"]
    tasks = build_tasks()
    done, lock, spent = set(), threading.Lock(), [0.0]
    if os.path.exists(OUT):
        for line in open(OUT):
            r = json.loads(line)
            done.add((r["system"], r["task"], r["id"]))
            spent[0] += r.get("cost") or 0.0
    print("resuming with %d done, $%.4f spent so far" % (len(done), spent[0]))
    out = open(OUT, "a")
    stop = threading.Event()

    def work(system, task, item):
        if stop.is_set() or (system, task["name"], item["id"]) in done:
            return
        rec = ask_jev(task, item, key) if system == JEV else ask_llm(system, task, item, key)
        rec.update(system=system, task=task["name"], id=item["id"], gold=item["gold"])
        with lock:
            out.write(json.dumps(rec) + "\n")
            out.flush()
            spent[0] += rec.get("cost") or 0.0
            if spent[0] > BUDGET_USD and not stop.is_set():
                stop.set()
                print("BUDGET STOP at $%.4f" % spent[0])

    def system_job(system):
        with ThreadPoolExecutor(WORKERS) as ex:
            for task in tasks:
                for item in task["items"]:
                    ex.submit(work, system, task, item)

    threads = [threading.Thread(target=system_job, args=(s,)) for s in SYSTEMS]
    [t.start() for t in threads]
    while any(t.is_alive() for t in threads):
        time.sleep(20)
        n = sum(1 for _ in open(OUT))
        print("  %d calls logged, $%.4f spent" % (n, spent[0]), flush=True)
    [t.join() for t in threads]
    print("finished, $%.4f spent" % spent[0])


# ---------------------------------------------------------------- report
def pct(xs, q):
    xs = sorted(xs)
    return xs[max(0, min(len(xs) - 1, math.ceil(q * len(xs)) - 1))] if xs else float("nan")


def wilson(k, n, z=1.96):
    if n == 0:
        return (float("nan"), float("nan"))
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (c - h, c + h)


def auroc(pos, neg):
    if not pos or not neg:
        return float("nan")
    s = sorted([(p, 1) for p in pos] + [(p, 0) for p in neg])
    rank, i, sum_pos = 1, 0, 0.0
    while i < len(s):
        j = i
        while j < len(s) and s[j][0] == s[i][0]:
            j += 1
        avg = (rank + rank + (j - i) - 1) / 2
        sum_pos += avg * sum(1 for x in s[i:j] if x[1] == 1)
        rank += j - i
        i = j
    return (sum_pos - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg))


def report():
    rows = [json.loads(l) for l in open(OUT)]
    summ = {}
    for system in SYSTEMS:
        for task in ("intent8", "intent77", "inject"):
            rs = [r for r in rows if r["system"] == system and r["task"] == task]
            if not rs:
                continue
            ok = [r for r in rs if r["ok"]]
            lat = [r["latency"] for r in ok]
            k = sum(1 for r in ok if r["pred"] == r["gold"])
            s = {"n": len(rs), "answered": len(ok), "failed": len(rs) - len(ok),
                 "acc_of_all": k / len(rs), "acc_ci": wilson(k, len(rs)),
                 "lat_p50": statistics.median(lat) if lat else None, "lat_p95": pct(lat, 0.95),
                 "lat_mean": statistics.mean(lat) if lat else None,
                 "cost_per_1k": 1000 * sum(r.get("cost") or 0 for r in ok) / max(1, len(ok)),
                 "in_tok_mean": statistics.mean([r["in_tok"] for r in ok if r.get("in_tok")]) if ok else None}
            if task == "inject":
                tp = sum(1 for r in ok if r["pred"] and r["gold"]); fp = sum(1 for r in ok if r["pred"] and not r["gold"])
                fn = sum(1 for r in ok if not r["pred"] and r["gold"]); tn = sum(1 for r in ok if not r["pred"] and not r["gold"])
                prec = tp / (tp + fp) if tp + fp else float("nan"); rec = tp / (tp + fn) if tp + fn else float("nan")
                s.update(tp=tp, fp=fp, fn=fn, tn=tn, precision=prec, recall=rec,
                         f1=2 * prec * rec / (prec + rec) if prec + rec else float("nan"),
                         auroc=auroc([r["prob"] for r in ok if r["gold"]], [r["prob"] for r in ok if not r["gold"]]),
                         brier=statistics.mean([(r["prob"] - (1.0 if r["gold"] else 0.0)) ** 2 for r in ok]) if ok else None)
            if system == JEV and task != "inject":
                gate = []
                for t in (0.0, 0.5, 0.7, 0.8, 0.9, 0.95):
                    cov = [r for r in ok if (r.get("conf") or 0) >= t]
                    gate.append({"min_conf": t, "coverage": len(cov) / len(rs),
                                 "accuracy": (sum(1 for r in cov if r["pred"] == r["gold"]) / len(cov)) if cov else None})
                s["gating"] = gate
            summ[system + "|" + task] = s
    json.dump(summ, open(os.path.join(HERE, "summary.json"), "w"), indent=1)
    print("\n%-30s %-9s %5s %6s %-15s %7s %7s %9s" % ("system", "task", "n", "fail", "accuracy (95% CI)", "p50 s", "p95 s", "$/1k dec."))
    for key, s in summ.items():
        system, task = key.split("|")
        print("%-30s %-9s %5d %6d %5.1f%% (%4.1f-%4.1f) %7.2f %7.2f %9.4f" % (
            system, task, s["n"], s["failed"], 100 * s["acc_of_all"], 100 * s["acc_ci"][0], 100 * s["acc_ci"][1],
            s["lat_p50"] or 0, s["lat_p95"] or 0, s["cost_per_1k"]))
        if task == "inject":
            print("%-30s   precision %.3f recall %.3f F1 %.3f AUROC %.3f Brier %.3f (tp %d fp %d fn %d tn %d)" % (
                "", s["precision"], s["recall"], s["f1"], s["auroc"], s["brier"], s["tp"], s["fp"], s["fn"], s["tn"]))
    for key, s in summ.items():
        if "gating" in s:
            print("\nJev confidence gating on", key.split("|")[1])
            for g in s["gating"]:
                print("  confidence >= %.2f: coverage %5.1f%%  accuracy %s" % (
                    g["min_conf"], 100 * g["coverage"], "n/a" if g["accuracy"] is None else "%.1f%%" % (100 * g["accuracy"])))


if __name__ == "__main__":
    {"fetch": fetch, "run": run, "report": report}[sys.argv[1] if len(sys.argv) > 1 else "report"]()