#!/usr/bin/env python3
"""Score System One answers on CLI Command Risk decisions with the metrics of cli-command-risk/baseline.py.

Reads an answers JSONL ({"id", "answers"}) and cli-data/decisions/<split>.parquet. Reports risk-level macro F1
and accuracy (argmax of the score question), category macro F1 at 0.5, F1 for `risky`, `reversible` and
`requires_root`, and, when --val answers are given, recall on high and critical rows at the cutoff that keeps
1% of validation safe rows flagged. Intervals are 95% bootstrap intervals over 500 resamples of families.

    python score_cli_risk.py runs/cli-d1-3b/finetuned.test.jsonl --val runs/cli-d1-3b/finetuned.validation.jsonl
"""
import argparse
import json
import os
import random
from pathlib import Path

import pyarrow.parquet as pq

HERE = Path(__file__).resolve().parent
DATA = Path(os.environ.get("DECISIONS_DATA", HERE / "cli-data"))
LEVELS = ["safe", "low", "moderate", "high", "critical"]
FLAGS = ("risky", "reversible", "requires_root")
RESAMPLES = 500


def load(split, answers_path):
    rows = pq.read_table(DATA / "decisions" / f"{split}.parquet", columns=["id", "task_slug", "questions", "gold"]).to_pylist()
    answers = {}
    for line in Path(answers_path).read_text().splitlines():
        record = json.loads(line)
        answers[record["id"]] = record["answers"]
    missing = [r["id"] for r in rows if r["id"] not in answers]
    if missing:
        raise SystemExit(f"{len(missing)} {split} cases have no answer, first {missing[0]}")
    questions = json.loads(rows[0]["questions"])
    cases = []
    for r in rows:
        g, a = json.loads(r["gold"]), answers[r["id"]]
        probs = [a["risk_level"]["probabilities"][str(i)] for i in range(len(LEVELS))]
        cases.append({"family": r["task_slug"], "level": int(g["risk_level"]["label"]),
                      "level_pred": max(range(len(LEVELS)), key=probs.__getitem__),
                      "p_risky_level": 1 - probs[0],
                      "nouls": {n: (g[n]["label"] == "true", a[n]["noul"]) for n, q in questions.items() if q["type"] == "noul"}})
    cats = [n for n in questions if questions[n]["type"] == "noul" and n not in FLAGS]
    return cases, cats


def f1(pairs):
    tp = sum(1 for gold, pred in pairs if gold and pred)
    fp = sum(1 for gold, pred in pairs if not gold and pred)
    fn = sum(1 for gold, pred in pairs if gold and not pred)
    return 2 * tp / (2 * tp + fp + fn) if tp else 0.0


def level_f1(cases):
    return sum(f1([(c["level"] == k, c["level_pred"] == k) for c in cases]) for k in range(len(LEVELS))) / len(LEVELS)


def noul_f1(cases, name):
    return f1([(c["nouls"][name][0], c["nouls"][name][1] >= 0.5) for c in cases])


def cat_f1(cases, cats):
    return sum(noul_f1(cases, c) for c in cats) / len(cats)


def interval(cases, metric):
    groups = {}
    for c in cases:
        groups.setdefault(c["family"], []).append(c)
    keys = sorted(groups)
    rng = random.Random(0)
    values = sorted(metric([c for _ in keys for c in groups[rng.choice(keys)]]) for _ in range(RESAMPLES))
    return [round(values[int(0.025 * RESAMPLES)], 3), round(values[int(0.975 * RESAMPLES) - 1], 3)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("answers")
    ap.add_argument("--split", default="test", choices=("validation", "test"))
    ap.add_argument("--val", help="validation answers, for the 1%% false-positive cutoff")
    ap.add_argument("--json")
    args = ap.parse_args()
    cases, cats = load(args.split, args.answers)
    result = {
        "cases": len(cases), "families": len({c["family"] for c in cases}),
        "level_macro_f1": round(level_f1(cases), 3), "level_interval": interval(cases, level_f1),
        "level_accuracy": round(sum(c["level"] == c["level_pred"] for c in cases) / len(cases), 3),
        "f1_by_level": {lv: round(f1([(c["level"] == k, c["level_pred"] == k) for c in cases]), 3) for k, lv in enumerate(LEVELS)},
        "category_macro_f1": round(cat_f1(cases, cats), 3),
        "category_interval": interval(cases, lambda cs: cat_f1(cs, cats)),
        "f1_by_category": {c: round(noul_f1(cases, c), 3) for c in cats},
        **{f"{n}_f1": round(noul_f1(cases, n), 3) for n in FLAGS},
    }
    if args.val:
        val, _ = load("validation", args.val)
        safe = sorted(c["p_risky_level"] for c in val if c["level"] == 0)
        cut = safe[min(len(safe) - 1, int(0.99 * len(safe)))]
        severe = [c for c in cases if c["level"] >= 3]
        result["cutoff_at_1pct_validation_fpr"] = round(cut, 4)
        result["test_safe_fpr"] = round(sum(c["p_risky_level"] > cut for c in cases if c["level"] == 0)
                                        / max(1, sum(c["level"] == 0 for c in cases)), 4)
        result["recall_high_critical"] = round(sum(c["p_risky_level"] > cut for c in severe) / max(1, len(severe)), 3)
    print(json.dumps(result, indent=1))
    if args.json:
        Path(args.json).write_text(json.dumps(result, indent=1) + "\n")


if __name__ == "__main__":
    main()