"""
Jev vs LLM benchmark via OpenRouter: accuracy, calibration, latency, cost.

uv run jev_bench.py --task spam --backend jev
uv run jev_bench.py --task spam --backend llm --llm-model openai/gpt-6-luna
Requires OPENROUTER_API_KEY (or OPENROUTER_KEY) in the environment.
"""
import argparse, json, os, random, statistics, time
from pathlib import Path

import numpy as np
import requests
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score

API_KEY = (os.environ.get("OPEN_ROUTER_KEY") or os.environ.get("OPENROUTER_API_KEY")
           or os.environ.get("OPENROUTER_KEY"))
BASE = "https://openrouter.ai/api"
HEADERS = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
SESSION = requests.Session()


# ---------------- data ----------------
def load_task(task, n, seed):
    if task == "spam":
        ds = load_dataset("ucirvine/sms_spam", split="train")
        texts, labels = ds["sms"], ["spam" if l == 1 else "ham" for l in ds["label"]]
        label_set = ["ham", "spam"]
    elif task == "banking77":
        import csv
        # original PolyAI CSV (HF copy is a legacy loading script)
        with open(Path(__file__).parent / "data" / "banking77_test.csv") as f:
            rows = list(csv.DictReader(f))
        texts, labels = [r["text"] for r in rows], [r["category"] for r in rows]
        label_set = sorted(set(labels))
    else:
        raise ValueError(task)
    idx = list(range(len(texts)))
    random.Random(seed).shuffle(idx)
    idx = idx[:n]
    return [texts[i] for i in idx], [labels[i] for i in idx], label_set


# ---------------- backends ----------------
class JevRunner:
    def __init__(self, task, label_set, model):
        self.task, self.model = task, model
        if task == "spam":
            self.q = {"spam": {"type": "noul", "instructions":
                      "Is this SMS message unsolicited spam (ads, scams, prize offers, phishing)?"}}
        else:
            self.q = {"intent": {"type": "choice", "instructions":
                      "Which banking customer-support intent does this message express?",
                      "criteria": {l: l.replace("_", " ") for l in label_set}}}

    def predict(self, text):
        body = {"model": self.model, "state": text, "questions": self.q}
        t0 = time.perf_counter()
        r = SESSION.post(f"{BASE}/alpha/decisions", headers=HEADERS, json=body, timeout=60)
        latency = (time.perf_counter() - t0) * 1000
        r.raise_for_status()
        d = r.json()
        if self.task == "spam":
            p = float(d["answers"]["spam"]["noul"])
            probs = {"spam": p, "ham": 1 - p}
        else:
            a = d["answers"]["intent"]
            probs = a.get("probabilities") or {a["choice"]: a.get("confidence", 1.0)}
            probs = {k: float(v) for k, v in probs.items()}
        u = d.get("usage", {}) or {}
        cost = u.get("cost", u.get("total_cost"))
        return probs, latency, cost, d.get("model")


class LLMRunner:
    SYSTEM = ("You are a text classifier. Reply ONLY with JSON: "
              '{"label": <one of the allowed labels>, "confidence": <0..1 probability '
              "that your label is correct>}.")

    def __init__(self, task, label_set, model, reasoning):
        self.model, self.label_set, self.reasoning = model, label_set, reasoning
        self.instr = ("Classify the SMS as 'spam' or 'ham'." if task == "spam"
                      else "Classify the banking support message into one intent.")

    def predict(self, text):
        prompt = f"{self.instr}\nAllowed labels: {', '.join(self.label_set)}\n\nText:\n{text}"
        body = {"model": self.model, "temperature": 0,
                "response_format": {"type": "json_object"},
                "messages": [{"role": "system", "content": self.SYSTEM},
                             {"role": "user", "content": prompt}],
                "usage": {"include": True}}
        if self.reasoning:
            body["reasoning"] = {"effort": self.reasoning}
        t0 = time.perf_counter()
        r = SESSION.post(f"{BASE}/v1/chat/completions", headers=HEADERS, json=body, timeout=120)
        latency = (time.perf_counter() - t0) * 1000
        r.raise_for_status()
        d = r.json()
        try:
            content = d["choices"][0]["message"]["content"].strip().strip("`")
            content = content[content.find("{"): content.rfind("}") + 1]
            out = json.loads(content)
            label, conf = out.get("label"), float(out.get("confidence", 0.5))
        except Exception:
            label, conf = None, 0.0
        conf = min(max(conf, 0.0), 1.0)
        if label not in self.label_set:
            label = "__invalid__"
        probs = {label: conf}
        if len(self.label_set) == 2 and label in self.label_set:
            other = [l for l in self.label_set if l != label][0]
            probs[other] = 1 - conf
        cost = (d.get("usage") or {}).get("cost")
        return probs, latency, cost, d.get("model")


# ---------------- metrics ----------------
def ece(conf, correct, bins=15):
    conf, correct = np.array(conf), np.array(correct, dtype=float)
    edges, total = np.linspace(0, 1, bins + 1), 0.0
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.any():
            total += m.mean() * abs(conf[m].mean() - correct[m].mean())
    return float(total)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", choices=["spam", "banking77"], required=True)
    ap.add_argument("--backend", choices=["jev", "llm"], required=True)
    ap.add_argument("--n", type=int, default=500)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--jev-model", default="typesafe/jev-1.13")
    ap.add_argument("--llm-model", default="openai/gpt-6-luna")
    ap.add_argument("--reasoning", default="minimal", help="'' to disable")
    ap.add_argument("--out", default="results")
    ap.add_argument("--max-cost", type=float, default=None, help="abort when spend exceeds this (USD)")
    a = ap.parse_args()
    if not API_KEY:
        raise SystemExit("OPENROUTER_API_KEY is not set")

    texts, gold, label_set = load_task(a.task, a.n, a.seed)
    runner = (JevRunner(a.task, label_set, a.jev_model) if a.backend == "jev"
              else LLMRunner(a.task, label_set, a.llm_model, a.reasoning))

    rows, errors = [], 0
    for i, (t, g) in enumerate(zip(texts, gold)):
        try:
            probs, lat, cost, model_id = runner.predict(t)
        except Exception as e:
            errors += 1
            print(f"[{i}] error: {str(e)[:200]}")
            if errors >= 10 and not rows:
                raise SystemExit("too many errors, aborting")
            continue
        pred = max(probs, key=probs.get)
        rows.append(dict(text=t, gold=g, pred=pred, conf=probs[pred],
                         p_spam=probs.get("spam"), latency_ms=lat, cost=cost, model=model_id))
        spent = sum(r["cost"] or 0 for r in rows)
        if a.max_cost is not None and spent > a.max_cost:
            print(f"cost cap ${a.max_cost} reached at {i+1} examples (${spent:.3f}), stopping")
            break
        if i % 25 == 0:
            print(f"{i}/{len(texts)}  last={lat:.0f}ms", flush=True)

    y, p = [r["gold"] for r in rows], [r["pred"] for r in rows]
    lats = sorted(r["latency_ms"] for r in rows)
    costs = [r["cost"] for r in rows if r["cost"] is not None]
    summary = dict(
        task=a.task, backend=a.backend, model=rows[0]["model"] if rows else None,
        n=len(rows), errors=errors,
        invalid_outputs=sum(r["pred"] == "__invalid__" for r in rows),
        accuracy=accuracy_score(y, p),
        macro_f1=f1_score(y, p, average="macro"),
        ece=ece([r["conf"] for r in rows], [r["gold"] == r["pred"] for r in rows]),
        latency_p50_ms=statistics.median(lats),
        latency_p95_ms=lats[int(0.95 * (len(lats) - 1))],
        cost_total_usd=sum(costs) if costs else None,
        cost_per_1k_usd=(sum(costs) / len(costs) * 1000) if costs else None,
    )
    if a.task == "spam":
        summary["brier"] = float(np.mean([(r["p_spam"] - (r["gold"] == "spam")) ** 2
                                          for r in rows if r["p_spam"] is not None]))

    out = Path(a.out); out.mkdir(exist_ok=True)
    model_tag = (a.jev_model if a.backend == "jev" else a.llm_model).replace("/", "_")
    stem = f"{a.task}_{a.backend}_{model_tag}"
    (out / f"{stem}.json").write_text(json.dumps(summary, indent=2))
    (out / f"{stem}_preds.jsonl").write_text("\n".join(json.dumps(r) for r in rows))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()