#!/usr/bin/env python3
"""Score System One answers on CLI Command Risk decisions with the leaderboard columns of score_decisions.py.

Same input and metrics as the JS/TS Code Review Decisions scorer: JSONL predictions {"id", "answers",
"latency_ms"?}, full distributions per question, macro F1 at 0.5, accuracy over all decisions, log loss, Brier,
ECE, and 95% bootstrap intervals over task families. Macro F1 averages the 13 risk-category nouls. Two CLI
columns are added: macro F1 over the 5 risk levels (argmax of `risk_level`) and safe cases judged safe.

    uv run --no-project --with pyarrow --with numpy python score_cli_decisions.py PREDICTIONS.jsonl [--json OUT]
    python score_cli_decisions.py --reference prior|uniform
"""
import argparse
import json
import os
import sys
from collections import Counter
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq

HERE = Path(__file__).resolve().parent
SPLITS = ("train", "validation", "test")
FLAGS = ("risky", "reversible", "requires_root")  # nouls outside the category macro F1
LEVELS = 5
PROBABILITY_FLOOR = 1e-6
ECE_BINS = 15
BOOTSTRAP_RESAMPLES = 500
SEED = 0


def load_split(data: Path, split: str) -> tuple[list[dict], dict]:
    rows = pq.read_table(data / "decisions" / f"{split}.parquet", columns=["id", "task_slug", "questions", "gold"]).to_pylist()
    schemas = {row["questions"] for row in rows}
    if len(schemas) != 1:
        sys.exit(f"{split} rows carry {len(schemas)} different question schemas")
    return rows, json.loads(schemas.pop())


def options_of(question: dict) -> list[str]:
    if question["type"] == "noul":
        return ["true", "false"]
    if question["type"] == "choice":
        return list(question["criteria"])
    return [str(i) for i in range(len(question["criteria"]))]


def distribution(answer: dict, question: dict, options: list[str]) -> tuple[np.ndarray, float]:
    """Probability over `options` in order, plus the mass the answer put on options that are not on the menu."""
    if question["type"] == "noul":
        p = float(answer["noul"])
        return np.array([p, 1 - p]), 0.0
    probabilities = {str(k): float(v) for k, v in answer["probabilities"].items()}
    off_menu = sum(v for k, v in probabilities.items() if k not in options)
    return np.array([probabilities.get(o, 0.0) for o in options]), off_menu


def reference_answers(data: Path, kind: str, questions: dict) -> dict:
    """The same answer for every case: train label frequencies (prior) or equal probabilities (uniform)."""
    answers = {}
    golds = [] if kind == "uniform" else [json.loads(r["gold"]) for r in
                                           pq.read_table(data / "decisions" / "train.parquet", columns=["gold"]).to_pylist()]
    for name, question in questions.items():
        options = options_of(question)
        counts = Counter(g[name]["label"] for g in golds)
        total = sum(counts.values())
        probs = {o: (counts[o] / total if total else 1 / len(options)) for o in options}
        answers[name] = ({"type": "noul", "noul": probs["true"]} if question["type"] == "noul"
                         else {"type": question["type"], "probabilities": probs})
    return answers


def expected_calibration_error(confidence: np.ndarray, outcome: np.ndarray) -> float:
    edges = np.linspace(0, 1, ECE_BINS + 1)
    bins = np.clip(np.digitize(confidence, edges[1:-1]), 0, ECE_BINS - 1)
    error = 0.0
    for b in range(ECE_BINS):
        members = bins == b
        if members.any():
            error += members.mean() * abs(confidence[members].mean() - outcome[members].mean())
    return float(error)


def f1(pred: np.ndarray, truth: np.ndarray) -> float:
    tp = int((pred & truth).sum())
    fp = int((pred & ~truth).sum())
    fn = int((~pred & truth).sum())
    return 2 * tp / (2 * tp + fp + fn) if tp + fp + fn else 0.0


class Scored:
    """Per-case, per-question arrays for one split, aligned with the cases that have predictions."""

    def __init__(self, rows: list[dict], questions: dict, predictions: dict[str, dict]):
        self.names = list(questions)
        self.types = [questions[n]["type"] for n in self.names]
        self.categories = [j for j, n in enumerate(self.names) if self.types[j] == "noul" and n not in FLAGS]
        self.cases = [row for row in rows if row["id"] in predictions]
        self.tasks = np.array([row["task_slug"] for row in self.cases])
        n, q = len(self.cases), len(self.names)
        self.correct = np.zeros((n, q), bool)
        self.log_loss = np.zeros((n, q))
        self.brier = np.zeros((n, q))
        self.confidence = np.zeros((n, q))  # p(true) for nouls, max probability otherwise
        self.outcome = np.zeros((n, q))  # the gold truth for nouls, correctness otherwise
        self.positives = np.zeros((n, q), bool)
        self.predicted_true = np.zeros((n, q), bool)
        self.level_gold = np.zeros(n, int)
        self.level_pred = np.zeros(n, int)
        self.level_error = np.zeros(n)
        self.off_menu = 0.0
        self.mass_missing = 0.0
        latencies = [predictions[row["id"]].get("latency_ms") for row in self.cases]
        self.latency_ms = np.array(latencies, float) if all(v is not None for v in latencies) else None
        for i, row in enumerate(self.cases):
            gold = json.loads(row["gold"])
            answers = predictions[row["id"]]["answers"]
            for j, name in enumerate(self.names):
                if name not in answers:
                    sys.exit(f"case {row['id']} has no answer for {name}")
                question = questions[name]
                options = options_of(question)
                p, off_menu = distribution(answers[name], question, options)
                self.off_menu += off_menu
                self.mass_missing += max(0.0, 1 - p.sum() - off_menu)
                target = options.index(gold[name]["label"])
                onehot = np.zeros(len(options))
                onehot[target] = 1
                self.log_loss[i, j] = -np.log(max(p[target], PROBABILITY_FLOOR))
                if question["type"] == "noul":
                    self.predicted_true[i, j] = p[0] >= 0.5
                    self.positives[i, j] = target == 0
                    self.correct[i, j] = self.predicted_true[i, j] == self.positives[i, j]
                    self.brier[i, j] = (p[0] - onehot[0]) ** 2
                    self.confidence[i, j], self.outcome[i, j] = p[0], onehot[0]
                else:
                    self.correct[i, j] = int(p.argmax()) == target
                    self.brier[i, j] = float(((p - onehot) ** 2).sum())
                    self.confidence[i, j], self.outcome[i, j] = p.max(), self.correct[i, j]
                    if name == "risk_level":
                        self.level_gold[i], self.level_pred[i] = target, int(p.argmax())
                        self.level_error[i] = abs(float(p @ np.arange(len(options))) - target)
        decisions = n * q
        self.off_menu /= decisions
        self.mass_missing /= decisions

    def headline(self, rows: np.ndarray | None = None) -> dict:
        rows = np.arange(len(self.cases)) if rows is None else rows
        gold, pred = self.level_gold[rows], self.level_pred[rows]
        return {
            "macro_f1_at_0_5": float(np.mean([f1(self.predicted_true[rows, j], self.positives[rows, j]) for j in self.categories])),
            "level_macro_f1": float(np.mean([f1(pred == k, gold == k) for k in range(LEVELS)])),
            "accuracy": float(self.correct[rows].mean()),
            "log_loss": float(self.log_loss[rows].mean()),
            "brier": float(self.brier[rows].mean()),
        }

    def bootstrap(self) -> dict[str, list[float]]:
        rng = np.random.default_rng(SEED)
        _, codes = np.unique(self.tasks, return_inverse=True)
        members = [np.flatnonzero(codes == g) for g in range(codes.max() + 1)]
        draws: dict[str, list[float]] = {}
        for _ in range(BOOTSTRAP_RESAMPLES):
            rows = np.concatenate([members[g] for g in rng.integers(0, len(members), len(members))])
            for key, value in self.headline(rows).items():
                draws.setdefault(key, []).append(value)
        return {k: [float(x) for x in np.percentile(v, [2.5, 97.5])] for k, v in draws.items()}

    def per_question(self) -> dict:
        out = {}
        for j, name in enumerate(self.names):
            entry = {"type": self.types[j], "accuracy": float(self.correct[:, j].mean()),
                     "log_loss": float(self.log_loss[:, j].mean()), "brier": float(self.brier[:, j].mean()),
                     "ece": expected_calibration_error(self.confidence[:, j], self.outcome[:, j])}
            if self.types[j] == "noul":
                entry["positives"] = int(self.positives[:, j].sum())
                entry["f1_at_0_5"] = f1(self.predicted_true[:, j], self.positives[:, j])
            if name == "risk_level":
                entry["level_mae"] = float(self.level_error.mean())
            out[name] = entry
        return out

    def report(self, total_cases: int) -> dict:
        per_question = self.per_question()
        by_type = {t: float(self.correct[:, [j for j, u in enumerate(self.types) if u == t]].mean())
                   for t in dict.fromkeys(self.types)}
        risky, safe = self.names.index("risky"), self.level_gold == 0
        latency = self.latency_ms
        return {
            "cases_in_split": total_cases, "cases_scored": len(self.cases), "decisions_scored": int(self.correct.size),
            "off_menu_mass": self.off_menu, "mass_missing": self.mass_missing,
            **self.headline(),
            "ece": float(np.mean([q["ece"] for q in per_question.values()])),
            "ci95": self.bootstrap(),
            "accuracy_by_type": by_type,
            "safe_cases_judged_safe": float((~self.predicted_true[safe, risky]).mean()),
            "p50_latency_ms": float(np.median(latency)) if latency is not None else None,
            "p95_latency_ms": float(np.percentile(latency, 95)) if latency is not None else None,
            "per_question": per_question,
        }


def markdown(result: dict) -> str:
    def interval(key: str) -> str:
        low, high = result["ci95"][key]
        return f"{result[key]:.3f} ({low:.3f} to {high:.3f})"

    latency = f"{result['p50_latency_ms']:.1f} ms" if result["p50_latency_ms"] is not None else "not reported"
    lines = [
        f"scored {result['cases_scored']}/{result['cases_in_split']} cases, {result['decisions_scored']:,} decisions; "
        f"off-menu mass {result['off_menu_mass']:.4f}; missing mass {result['mass_missing']:.4f}",
        "",
        "| Macro F1 at 0.5 | Level macro F1 | Accuracy | Log loss | Brier | ECE | p50 latency |",
        "|---:|---:|---:|---:|---:|---:|---:|",
        f"| {interval('macro_f1_at_0_5')} | {interval('level_macro_f1')} | {interval('accuracy')} | {interval('log_loss')} | "
        f"{interval('brier')} | {result['ece']:.3f} | {latency} |",
        "",
        "Accuracy by type: " + ", ".join(f"{t} {v:.3f}" for t, v in result["accuracy_by_type"].items())
        + f"; safe cases judged safe {result['safe_cases_judged_safe']:.3f}",
        "",
        "| Question | Type | Positives | Accuracy | F1 at 0.5 | Log loss | Brier | ECE |",
        "|---|---|---:|---:|---:|---:|---:|---:|",
    ]
    for name, q in result["per_question"].items():
        positives = str(q["positives"]) if "positives" in q else ""
        f1_cell = f"{q['f1_at_0_5']:.3f}" if "f1_at_0_5" in q else ""
        lines.append(f"| `{name}` | {q['type']} | {positives} | {q['accuracy']:.3f} | {f1_cell} | {q['log_loss']:.3f} | "
                     f"{q['brier']:.3f} | {q['ece']:.3f} |")
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("predictions", nargs="?", type=Path)
    parser.add_argument("--reference", choices=("prior", "uniform"))
    parser.add_argument("--split", default="test", choices=SPLITS)
    parser.add_argument("--data", type=Path, default=Path(os.environ.get("DECISIONS_DATA", HERE / "cli-data")),
                        help="folder holding decisions/*.parquet")
    parser.add_argument("--json", type=Path, help="write the full result here")
    args = parser.parse_args()
    if (args.predictions is None) == (args.reference is None):
        parser.error("give a predictions file or --reference, not both")
    rows, questions = load_split(args.data, args.split)
    if args.reference:
        answers = reference_answers(args.data, args.reference, questions)
        predictions = {row["id"]: {"answers": answers} for row in rows}
    else:
        predictions = {}
        for line in args.predictions.read_text().splitlines():
            if line.strip():
                entry = json.loads(line)
                predictions[entry["id"]] = entry
    ids = {row["id"] for row in rows}
    unknown = sorted(set(predictions) - ids)
    if unknown:
        print(f"ignoring {len(unknown)} predictions for ids not in {args.split}, e.g. {unknown[0]}", file=sys.stderr)
    if not ids & set(predictions):
        sys.exit(f"no prediction matches a case id in {args.split}")
    result = Scored(rows, questions, predictions).report(len(rows))
    result["split"] = args.split
    result["source"] = f"reference:{args.reference}" if args.reference else str(args.predictions)
    print(markdown(result))
    if args.json:
        args.json.write_text(json.dumps(result, indent=1) + "\n")


if __name__ == "__main__":
    main()