#!/usr/bin/env python3
"""Compare two paired v1 eval runs, conservatively counting scheduled episodes."""

from __future__ import annotations

import argparse
import json
import math
import tomllib
from collections import Counter
from pathlib import Path


def wilson(successes: int, total: int, z: float = 1.959963984540054) -> list[float] | None:
    if total == 0:
        return None
    p = successes / total
    d = 1 + z * z / total
    c = (p + z * z / (2 * total)) / d
    r = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / d
    return [max(0.0, c - r), min(1.0, c + r)]


def read_traces(path: Path) -> list[dict]:
    path = path / "traces.jsonl" if path.is_dir() else path
    by_id: dict[str, dict] = {}
    for line in path.read_text().splitlines():
        record = json.loads(line)
        for trace in record.get("traces", []):
            by_id[trace["id"]] = trace
    return list(by_id.values())


def task_key(trace: dict) -> str:
    data = (trace.get("task") or {}).get("data") or {}
    return str(data.get("name") or data.get("instance_id") or data.get("idx") or trace["id"])


def clean(trace: dict) -> bool:
    return bool(trace.get("ok") and trace.get("is_completed") and not trace.get("errors"))


def solved(trace: dict) -> bool:
    return float(((trace.get("rewards") or {}).get("solved") or {}).get("score", 0)) > 0


def has_reward(trace: dict) -> bool:
    return isinstance((trace.get("rewards") or {}).get("solved"), dict)


def exact_two_sided_sign_p(a_only: int, b_only: int) -> float | None:
    """Exact paired sign-test p-value over discordant outcomes.

    This is deliberately dependency-free and conditions only on tasks where the
    two harnesses/checkpoints disagree.  It is diagnostic rather than a reason
    to hide either run's conservative Wilson interval.
    """
    discordant = a_only + b_only
    if discordant == 0:
        return None
    tail = sum(math.comb(discordant, k) for k in range(min(a_only, b_only) + 1))
    return min(1.0, 2.0 * tail / (2**discordant))


def expected_episodes(path: Path, observed: int) -> int:
    run_dir = path if path.is_dir() else path.parent
    config_path = run_dir / "config.toml"
    if not config_path.is_file():
        return observed
    config = tomllib.loads(config_path.read_text())
    num_tasks = config.get("num_tasks")
    num_rollouts = config.get("num_rollouts", 1)
    if not isinstance(num_tasks, int) or not isinstance(num_rollouts, int):
        return observed
    return max(observed, num_tasks * num_rollouts)


def summary(traces: list[dict], path: Path) -> dict:
    clean_traces = [trace for trace in traces if clean(trace)]
    scored = [trace for trace in traces if has_reward(trace)]
    successes = sum(solved(trace) for trace in traces)
    calls = [call for trace in traces for call in trace.get("calls", [])]
    prompt_tokens = sum(int((call.get("usage") or {}).get("prompt_tokens", 0)) for call in calls)
    completion_tokens = sum(
        int((call.get("usage") or {}).get("completion_tokens", 0)) for call in calls
    )
    categories = Counter(
        str(((trace.get("task") or {}).get("data") or {}).get("category") or "unknown")
        for trace in traces
        if solved(trace)
    )
    scheduled = expected_episodes(path, len(traces))
    return {
        "scheduled_episodes": scheduled,
        "observed_trace_episodes": len(traces),
        "missing_trace_episodes": scheduled - len(traces),
        "clean_completed": len(clean_traces),
        "reward_bearing": len(scored),
        "errored_or_unscored": scheduled - len(scored),
        "successes": successes,
        "score_scheduled": successes / scheduled if scheduled else None,
        "ci95_wilson_scheduled": wilson(successes, scheduled),
        "diagnostic_score_reward_bearing": successes / len(scored) if scored else None,
        "diagnostic_ci95_reward_bearing": wilson(successes, len(scored)),
        "stop_conditions": dict(Counter(trace.get("stop_condition", "unknown") for trace in traces)),
        "model_calls": len(calls),
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "success_categories": dict(sorted(categories.items())),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("a", type=Path)
    parser.add_argument("b", type=Path)
    parser.add_argument("--label-a", default="a")
    parser.add_argument("--label-b", default="b")
    args = parser.parse_args()

    a_traces = read_traces(args.a)
    b_traces = read_traces(args.b)
    a = {task_key(trace): trace for trace in a_traces}
    b = {task_key(trace): trace for trace in b_traces}
    common = sorted(a.keys() & b.keys())
    paired = Counter()
    for key in common:
        sa, sb = solved(a[key]), solved(b[key])
        paired[
            "both" if sa and sb else "a_only" if sa else "b_only" if sb else "neither"
        ] += 1

    paired_delta = (
        (paired["b_only"] - paired["a_only"]) / len(common) if common else None
    )

    print(
        json.dumps(
            {
                args.label_a: summary(a_traces, args.a),
                args.label_b: summary(b_traces, args.b),
                "paired": {
                    "common_scheduled_tasks": len(common),
                    "both_solve": paired["both"],
                    f"{args.label_a}_only": paired["a_only"],
                    f"{args.label_b}_only": paired["b_only"],
                    "neither": paired["neither"],
                    f"score_delta_{args.label_b}_minus_{args.label_a}": paired_delta,
                    "discordant_tasks": paired["a_only"] + paired["b_only"],
                    "exact_two_sided_sign_test_p": exact_two_sided_sign_p(
                        paired["a_only"], paired["b_only"]
                    ),
                    f"missing_{args.label_a}": len(b.keys() - a.keys()),
                    f"missing_{args.label_b}": len(a.keys() - b.keys()),
                    "both_clean_completed": sum(clean(a[key]) and clean(b[key]) for key in common),
                },
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()