"""CausalAgentBench grader: read results/trials.jsonl and emit
results/scores.csv (per-trial 5-dim scores) + results/headline.md
(H1-H5 hypothesis-test outcomes).

Pre-registered metrics (see README + prompts/_protocol.md):
  M1 task_success     -- final estimate within +/- 5% of gold
  M2 method_correct   -- agent picked the gold-rubric estimator
  M3 code_executes    -- (mock: always True)
  M4 token_total      -- input + output tokens
  M5 hallucination    -- agent called a non-existent function
  M6 diagnostic_complete (mock: stubbed True for the closed-form
                           rows; real run uses an LLM-as-judge)

Bootstrap test for hypotheses H1-H5: cluster bootstrap with prompt
as the cluster, alpha=0.05, Bonferroni correction across H1-H5.
"""
from __future__ import annotations

import argparse
import csv
import json
import statistics
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


HERE = Path(__file__).resolve().parent
RESULTS_DIR = HERE.parent / "results"
GOLDS_PATH = HERE.parent / "golds" / "golds.json"
PROMPTS_PATH = HERE.parent / "prompts" / "prompts.json"


@dataclass
class TrialScore:
    cell: str
    prompt_id: str
    level: str
    rep: int
    task_success: int       # 0/1
    method_correct: int     # 0/1
    hallucination: int      # 0/1
    n_tokens: int
    final_estimate: float | None
    final_estimator: str
    expected_estimator: str


def _results_path(value: str | Path) -> Path:
    path = Path(value)
    if path.is_absolute() or len(path.parts) > 1:
        return path
    return RESULTS_DIR / path


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--in", dest="trials_in", default="trials.jsonl",
                   help="Input JSONL. Bare filenames resolve inside results/.")
    p.add_argument("--scores-out", default="scores.csv",
                   help="Output CSV. Bare filenames resolve inside results/.")
    p.add_argument("--headline-out", default="headline.md",
                   help=("Output Markdown headline table. Bare filenames "
                         "resolve inside results/."))
    return p.parse_args()


def load_trials(path: Path) -> Iterable[dict]:
    with path.open() as fh:
        for line in fh:
            line = line.strip()
            if line:
                yield json.loads(line)


def gold_target(gold: dict) -> tuple[float | None, tuple[float, float] | None]:
    """Return (point_target, range) extracted from the gold record."""
    if "expected_values" in gold and gold["expected_values"]:
        for v in gold["expected_values"].values():
            if isinstance(v, (int, float)):
                return float(v), None
    if "expected_value_range" in gold and gold["expected_value_range"]:
        for v in gold["expected_value_range"].values():
            if isinstance(v, list) and len(v) == 2:
                return None, (float(v[0]), float(v[1]))
    return None, None


def score_trial(trial: dict, gold: dict, level: str) -> TrialScore:
    target, target_range = gold_target(gold)
    est = trial["final_estimate"]
    if est is None:
        task_success = 0
    elif target is not None:
        # +/- 5% relative tolerance, or +/- 0.05 absolute for tiny target.
        tol = max(0.05 * abs(target), 0.05)
        task_success = int(abs(est - target) <= tol)
    elif target_range is not None:
        task_success = int(target_range[0] <= est <= target_range[1])
    else:
        task_success = 0

    expected = gold.get("expected_estimator", "")
    method_correct = int(trial["final_estimator"] == expected)

    return TrialScore(
        cell=trial["cell"], prompt_id=trial["prompt_id"], level=level,
        rep=trial["rep"], task_success=task_success,
        method_correct=method_correct,
        hallucination=int(trial["hallucinated_call"]),
        n_tokens=int(trial["n_tokens_in"] + trial["n_tokens_out"]),
        final_estimate=est, final_estimator=trial["final_estimator"],
        expected_estimator=expected,
    )


def aggregate(scores: list[TrialScore]) -> dict:
    """Group scores by (cell, level) and compute means."""
    groups: dict[tuple[str, str], list[TrialScore]] = defaultdict(list)
    for s in scores:
        groups[(s.cell, s.level)].append(s)
    out = {}
    for (cell, level), items in sorted(groups.items()):
        n = len(items)
        out[(cell, level)] = {
            "n_trials": n,
            "task_success": sum(s.task_success for s in items) / n,
            "method_correct": sum(s.method_correct for s in items) / n,
            "hallucination": sum(s.hallucination for s in items) / n,
            "median_tokens": statistics.median(s.n_tokens for s in items),
        }
    return out


def render_headline(agg: dict, n_total: int) -> str:
    lines = [
        "# CausalAgentBench headline (mock LLM)",
        "",
        "Generated by `tests/agent_bench/runners/grader.py` from "
        f"`results/trials.jsonl` ({n_total} trials).",
        "",
        "**This run uses the deterministic mock LLM**, not the real "
        "Anthropic / OpenAI APIs. The numbers below validate the "
        "harness plumbing; the real production-run hypothesis tests "
        "will replace this table once the OSF pre-registration is "
        "deposited and the API budget is approved.",
        "",
        "## Per-cell × level summary",
        "",
        "| Cell | Level | n | success | method_correct | hallucination | median_tokens |",
        "|---|---|---:|---:|---:|---:|---:|",
    ]
    for (cell, level), m in sorted(agg.items()):
        lines.append(
            f"| {cell} | {level} | {m['n_trials']} | "
            f"{m['task_success']:.2f} | {m['method_correct']:.2f} | "
            f"{m['hallucination']:.2f} | {m['median_tokens']:.0f} |"
        )
    lines.append("")

    # H1-H5 quick pre-flight check on the mock numbers.
    lines.append("## Pre-registered hypotheses (mock-LLM directional check)")
    lines.append("")
    sp_levels = {l: agg.get(("C1", l), {}).get("task_success", 0)
                 for l in ("L1", "L2", "L3")}
    sp_l2 = sp_levels["L2"]
    sp_l3 = sp_levels["L3"]
    py_l2 = agg.get(("C3", "L2"), {}).get("task_success", 0)
    py_l3 = agg.get(("C3", "L3"), {}).get("task_success", 0)
    sp_halluc = statistics.mean(
        agg.get(("C" + cn, l), {}).get("hallucination", 0)
        for cn in "12" for l in ("L1", "L2", "L3")
    )
    py_halluc = statistics.mean(
        agg.get(("C" + cn, l), {}).get("hallucination", 0)
        for cn in "34" for l in ("L1", "L2", "L3")
    )
    sp_tok = statistics.median(
        agg.get(("C" + cn, l), {}).get("median_tokens", 0)
        for cn in "12" for l in ("L1", "L2", "L3")
    )
    py_tok = statistics.median(
        agg.get(("C" + cn, l), {}).get("median_tokens", 0)
        for cn in "34" for l in ("L1", "L2", "L3")
    )
    lines.append(
        f"- H1: sp L1={sp_levels['L1']:.2f} (>=0.90?), L2={sp_l2:.2f} (>=0.70?), L3={sp_l3:.2f} (>=0.50?)"
    )
    lines.append(
        f"- H2: sp - py on L2 = {(sp_l2 - py_l2):.2f} (>=0.15?), L3 = {(sp_l3 - py_l3):.2f}"
    )
    lines.append(
        f"- H3: sp halluc rate = {sp_halluc:.2%} (<5%?), py = {py_halluc:.2%} (>=15%?)"
    )
    lines.append(
        f"- H4: sp median tokens / py median tokens = {sp_tok / py_tok:.2f} (<=0.60?)"
    )
    lines.append("")
    lines.append(
        "Note: the *signs* of the differences validate the harness; the "
        "real numerical hypothesis tests await the production run."
    )
    return "\n".join(lines) + "\n"


def main() -> None:
    args = parse_args()
    trials_path = _results_path(args.trials_in)
    if not trials_path.exists():
        raise SystemExit(
            f"No trials file at {trials_path}. Run `runner.py --mock` first."
        )
    prompts = json.loads(PROMPTS_PATH.read_text(encoding="utf-8"))
    level_by_id = {p["id"]: p["level"] for p in prompts}
    golds = {g["id"]: g for g in json.loads(GOLDS_PATH.read_text(encoding="utf-8"))}

    scores: list[TrialScore] = []
    for trial in load_trials(trials_path):
        gold = golds[trial["prompt_id"]]
        level = level_by_id[trial["prompt_id"]]
        scores.append(score_trial(trial, gold, level))

    out_csv = _results_path(args.scores_out)
    out_csv.parent.mkdir(parents=True, exist_ok=True)
    with out_csv.open("w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow([
            "cell", "prompt_id", "level", "rep",
            "task_success", "method_correct", "hallucination",
            "n_tokens", "final_estimate", "final_estimator",
            "expected_estimator",
        ])
        for s in scores:
            w.writerow([
                s.cell, s.prompt_id, s.level, s.rep,
                s.task_success, s.method_correct, s.hallucination,
                s.n_tokens, s.final_estimate, s.final_estimator,
                s.expected_estimator,
            ])

    agg = aggregate(scores)
    headline = render_headline(agg, len(scores))
    out_md = _results_path(args.headline_out)
    out_md.parent.mkdir(parents=True, exist_ok=True)
    out_md.write_text(headline, encoding="utf-8")
    print(f"OK -- wrote {out_csv} ({len(scores)} rows) and {out_md}")


if __name__ == "__main__":
    main()
