#!/usr/bin/env python3
"""Question-level evaluation: accuracy, distribution error, and calibration.

Reported per split, per game, and per question type.  Two controls are printed
alongside every accuracy so the numbers cannot be read as better than they are:
``majority`` is the accuracy of always answering the most common label, and
``uniform`` is the accuracy of guessing uniformly over the offered candidates.
NanoJev's maze atomic accuracy matched its own always-true control exactly,
which is the failure mode these columns are here to expose.
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch

from jevon.data import collate, load_frozen
from jevon.inference import JevonRunner
from jevon.losses import calibration_bins
from jevon.planner import default_iterations


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--checkpoint", default="runs/jevon-final")
    # None means "prefer balanced.pt", the checkpoint every shipped number
    # comes from. See jevon.inference.resolve_weights.
    p.add_argument("--weights", default=None)
    p.add_argument("--splits", default="data/frozen/test.jsonl,data/frozen/ood.jsonl")
    p.add_argument("--batch-states", type=int, default=8)
    p.add_argument("--device", default="mps")
    p.add_argument("--out", default=None)
    args = p.parse_args()

    runner = JevonRunner(args.checkpoint, args.device, args.weights)
    report = {}

    for path in args.splits.split(","):
        path = path.strip()
        if not path or not Path(path).exists():
            continue
        name = Path(path).stem
        states = load_frozen(path)
        buckets: dict[tuple[str, str], list] = defaultdict(list)
        # How many distinct boards a bucket's questions came from. Several
        # questions about one maze would not be several independent tests of
        # whether the model can read a maze -- they share walls, goal and
        # distance field, so they succeed or fail together -- and n alone
        # cannot tell you whether that is happening. Measured here it is not:
        # the generator draws a fresh board per state, so boards tracks n on
        # every split. The column exists to keep that checkable rather than
        # assumed, because it is a property of the generator and generators
        # get edited.
        board_ids: dict[tuple[str, str], set] = defaultdict(set)
        conf, hit = [], []
        with torch.no_grad():
            for start in range(0, len(states), args.batch_states):
                group = states[start:start + args.batch_states]
                batch = collate(group, runner.tokenizer, device=runner.device)
                steps = default_iterations(int(batch["board_size"].max()))
                logits = runner.model(batch, iterations=steps, grad_steps=0)
                probs = logits.float().softmax(-1).cpu()
                target = batch["target_probs"].cpu()
                for row, meta in enumerate(batch["meta"]):
                    k = len(meta["probs"])
                    pr, tg = probs[row, :k], target[row, :k]
                    picked = int(pr.argmax())
                    correct = float(tg[picked] >= tg.max() - 1e-6)
                    tvd = float(0.5 * (pr - tg).abs().sum())
                    brier = float(((pr - tg) ** 2).sum())
                    # Controls computed on the very same questions.
                    majority = float(tg[int(tg.argmax())] >= tg.max() - 1e-6)
                    uniform = float(tg.max() * 0 + (tg >= tg.max() - 1e-6).float().mean())
                    key = (meta["game"], meta["type"])
                    # The uid schema differs per game -- maze is
                    # "maze:<topology>:<size>:<seed>:<row>_<col>" and snake is
                    # "snake:<size>:<seed>:<step>" -- but the last segment is
                    # the within-board position in both, so dropping it names
                    # the board either way. Cross-checked against a hash of
                    # channel 0 and the goal cell: the two agree to within
                    # three coincidental duplicates across all three splits.
                    board = (meta["size"], meta["topology"],
                             meta["uid"].rsplit(":", 1)[0])
                    kind_key = (meta["game"], meta["question"].split("_")[0])
                    buckets[key].append((correct, tvd, brier, uniform))
                    buckets[kind_key].append((correct, tvd, brier, uniform))
                    board_ids[key].add(board)
                    board_ids[kind_key].add(board)
                    del majority
                    conf.append(float(pr.max()))
                    hit.append(correct)
        rows = {}
        for (game, kind), values in sorted(buckets.items()):
            n = len(values)
            rows[f"{game}/{kind}"] = {
                "n": n,
                "boards": len(board_ids[(game, kind)]),
                "accuracy": sum(v[0] for v in values) / n,
                "uniform_control": sum(v[3] for v in values) / n,
                "tvd": sum(v[1] for v in values) / n,
                "brier": sum(v[2] for v in values) / n,
            }
        calib = calibration_bins(torch.tensor(conf), torch.tensor(hit))
        report[name] = {"states": len(states), "questions": len(conf),
                        "ece": calib["ece"], "groups": rows}
        print(f"\n=== {name}: {len(states)} states / {len(conf)} questions "
              f"(ECE {calib['ece']:.4f}) ===")
        print(f"{'group':28s} {'n':>5s} {'brds':>5s} {'acc':>7s} {'unif':>7s} "
              f"{'tvd':>7s} {'brier':>7s}")
        for key, row in rows.items():
            print(f"{key:28s} {row['n']:5d} {row['boards']:5d} "
                  f"{row['accuracy']:7.4f} {row['uniform_control']:7.4f} "
                  f"{row['tvd']:7.4f} {row['brier']:7.4f}")

    if args.out:
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.out).write_text(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()