"""Evaluation of a Groot checkpoint on decision rows and public benchmarks.

    groot-eval --model runs/groot-1-s/final --rows data/test.jsonl [--max-rows 2000] [--out eval.json]
    groot-eval --model ... --bench typed_decisions   # LocalLLaMA/typed-decisions test split

Reports accuracy per type, NLL, Brier, ECE (with and without the stored
temperatures), coverage/accuracy at confidence thresholds, and per-source
breakdowns. Rows are the JSONL format of data.py.
"""
from __future__ import annotations

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

import numpy as np
import torch

from .calibration import Record, evaluate_records, softmax
from .config import QTYPES
from .data import CONVERTERS, read_jsonl
from .model import GrootDecisionModel
from .schema import confidence_choice, confidence_noul, confidence_score
from .serialize import Serializer
from .train import evaluate


def confidence_of(probabilities: list[float], qtype: str) -> float:
    if qtype == "choice":
        return confidence_choice(probabilities)
    if qtype == "score":
        return confidence_score(probabilities)
    return confidence_noul(probabilities[1])


def coverage_table(records: list[Record], temps: list[float], thresholds=(0.3, 0.5, 0.7, 0.9)) -> dict:
    """Accuracy of answers whose System One confidence passes each threshold."""
    rows = []
    for r in records:
        p = softmax(r.logits, temps[r.qtype]).tolist()
        c = confidence_of(p, QTYPES[r.qtype])
        rows.append((c, int(np.argmax(p) == int(r.target.argmax()))))
    out = {}
    for th in thresholds:
        kept = [a for c, a in rows if c >= th]
        out[str(th)] = {"coverage": len(kept) / max(1, len(rows)), "accuracy": float(np.mean(kept)) if kept else None}
    return out


def per_source(rows: list[dict], records: list[Record], temps: list[float]) -> dict:
    groups: dict[str, list[Record]] = defaultdict(list)
    for row, rec in zip(rows, records):
        groups[row.get("source", "?")].append(rec)
    return {src: evaluate_records(recs, temps).to_dict() for src, recs in groups.items()}


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(prog="groot-eval")
    ap.add_argument("--model", required=True)
    ap.add_argument("--rows", default=None)
    ap.add_argument("--bench", default=None, help="name of a converter in data.py (uses its test split)")
    ap.add_argument("--max-rows", type=int, default=None)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--device", default=None)
    ap.add_argument("--out", default=None)
    ap.add_argument("--set", action="append", default=[], help="config override key=value (e.g. state_retrieval=true)")
    ap.add_argument("--calibrate-on", default=None, help="fit per-type temperatures on this JSONL before evaluating")
    ap.add_argument("--save-calibration", action="store_true", help="write the fitted temperatures into the checkpoint")
    args = ap.parse_args(argv)

    from transformers import AutoTokenizer

    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"))
    model = GrootDecisionModel.from_pretrained(args.model, device=device)
    tok = AutoTokenizer.from_pretrained(Path(args.model) / "tokenizer")
    if args.set:
        from .config import GrootConfig, parse_overrides

        model.cfg = GrootConfig.from_dict({**model.cfg.to_dict(), **parse_overrides(args.set)})
    ser = Serializer(tok, model.cfg)
    if args.calibrate_on:
        from .calibration import fit_temperatures

        cal_rows = list(read_jsonl(args.calibrate_on))
        cal_records, _ = evaluate(model, cal_rows, ser, device, args.batch_size, None)
        temps = fit_temperatures(cal_records, n_types=len(model.cfg.qtypes))
        model.temperature.copy_(torch.tensor(temps, dtype=model.temperature.dtype, device=model.temperature.device))
        print(json.dumps({"fitted_temperatures": temps, "on": args.calibrate_on, "rows": len(cal_rows)}), flush=True)
        if args.save_calibration:
            (Path(args.model) / "calibration.json").write_text(json.dumps({"temperature": temps}, indent=2) + "\n")
    if args.rows:
        rows = list(read_jsonl(args.rows, args.max_rows))
    elif args.bench:
        rows = [r for r in CONVERTERS[args.bench](limit=args.max_rows) if r["split"] == "test"]
    else:
        raise SystemExit("pass --rows or --bench")
    records, raw = evaluate(model, rows, ser, device, args.batch_size, None)
    temps = model.temperature.tolist()
    report = {
        "model": args.model,
        "rows": len(rows),
        "raw": raw,
        "calibrated": evaluate_records(records, temps).to_dict(),
        "coverage": coverage_table(records, temps),
        "per_source": per_source(rows, records, temps),
    }
    text = json.dumps(report, indent=2)
    print(text)
    if args.out:
        Path(args.out).write_text(text + "\n")


if __name__ == "__main__":
    main()