"""Score a model served by `decida serve` on a typed-decisions test file. Standard library only.

    decida serve -m decidabert=helmo/DecidaBERT-large --max-len 2048 &
    python eval_typed_decisions.py --model decidabert --data benchmark_test.jsonl

The file has one case per line: {"state", "questions": "<json>", "gold": "<json>"}. Each gold answer is a label plus a soft distribution
(the average of three samples from a teacher model). A prediction counts as right when its top answer equals the gold label.
Prints one JSON object: accuracy per question type with a 95% Wilson interval, mean confidence, expected calibration error (10 bins),
and the accuracy a random guesser would get.
"""
import argparse
import json
import math
import urllib.request

BATCH = 256


def post(url: str, body: dict) -> dict:
    req = urllib.request.Request(url, json.dumps(body).encode(), {"content-type": "application/json"})
    with urllib.request.urlopen(req, timeout=600) as r:
        return json.load(r)


def wilson(k: int, n: int, z: float = 1.96) -> list[float]:
    if n == 0:
        return [0.0, 0.0]
    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 [round(c - h, 4), round(c + h, 4)]


def top(kind: str, ans: dict) -> tuple[str, float, dict]:
    """The model's top answer as a label string, its probability, and the full distribution over labels."""
    if kind == "noul":
        p = float(ans["noul"])
        dist = {"true": p, "false": 1 - p}
    else:
        dist = {str(k): float(v) for k, v in ans["probabilities"].items()}
    label = max(dist, key=dist.get)
    return label, dist[label], dist


def ece(pairs: list[tuple[float, bool]], bins: int = 10) -> float:
    total = len(pairs)
    err = 0.0
    for b in range(bins):
        lo, hi = b / bins, (b + 1) / bins
        cell = [(c, ok) for c, ok in pairs if lo < c <= hi or (b == 0 and c == 0)]
        if cell:
            err += len(cell) / total * abs(sum(ok for _, ok in cell) / len(cell) - sum(c for c, _ in cell) / len(cell))
    return err


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--url", default="http://127.0.0.1:8000")
    a = ap.parse_args()
    with open(a.data, encoding="utf-8") as f:
        cases = [json.loads(line) for line in f]
    questions = [json.loads(c["questions"]) for c in cases]
    golds = [json.loads(c["gold"]) for c in cases]
    responses: list[dict] = []
    for lo in range(0, len(cases), BATCH):
        chunk = [{"state": c["state"], "questions": q} for c, q in zip(cases[lo:lo + BATCH], questions[lo:lo + BATCH], strict=True)]
        responses += post(f"{a.url}/v1/systemone/batch", {"model": a.model, "requests": chunk})["responses"]
    stats: dict[str, dict] = {}
    for resp, qs, gold in zip(responses, questions, golds, strict=True):
        for qid, q in qs.items():
            kind = q["type"]
            s = stats.setdefault(kind, {"n": 0, "right": 0, "near": 0, "failed": 0, "pairs": [], "chance": 0.0})
            s["n"] += 1
            k = len(q.get("criteria") or {"true": 1, "false": 1})
            s["chance"] += 1 / k
            if "answers" not in resp:
                s["failed"] += 1
                continue
            label, conf, _ = top(kind, resp["answers"][qid])
            ok = label == str(gold[qid]["label"])
            s["right"] += ok
            s["pairs"].append((conf, ok))
            if kind == "score":
                s["near"] += abs(int(label) - int(gold[qid]["label"])) <= 1
    out = {"model": a.model, "cases": len(cases), "types": {}}
    allpairs, right, n = [], 0, 0
    for kind, s in sorted(stats.items()):
        row = {"questions": s["n"], "accuracy": round(s["right"] / s["n"], 4), "ci95": wilson(s["right"], s["n"]),
               "chance": round(s["chance"] / s["n"], 4), "mean_confidence": round(sum(c for c, _ in s["pairs"]) / max(1, len(s["pairs"])), 4),
               "ece": round(ece(s["pairs"]), 4)}
        if kind == "score":
            row["within_one_level"] = round(s["near"] / s["n"], 4)
        if s["failed"]:
            row["failed"] = s["failed"]
        out["types"][kind] = row
        allpairs += s["pairs"]
        right += s["right"]
        n += s["n"]
    out["overall"] = {"questions": n, "accuracy": round(right / n, 4), "ci95": wilson(right, n), "ece": round(ece(allpairs), 4)}
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()