"""Run the leakage-controlled tau-bench semantic-trace study.

The outcome is the benchmark's released database/policy reward. Feature
extraction never reads task evaluation criteria, database checks, assertions,
or reward breakdowns. Every prediction is cross-fitted while holding out both
the task group and the agent model.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import platform
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Iterable, Sequence

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import FormatStrFormatter

HERE = Path(__file__).resolve().parent
RESEARCH_ROOT = HERE.parents[1]
sys.path.insert(0, str(RESEARCH_ROOT / "common"))

from plot_style import COLORS, configure, label_panels, save_figure  # noqa: E402

CONFIG_PATH = HERE / "study_config.json"
RESULTS = HERE / "results"
FIGURES = HERE / "figures"

WRITE_VERBS = {
    "add",
    "book",
    "cancel",
    "create",
    "delete",
    "disable",
    "disconnect",
    "enable",
    "exchange",
    "grant",
    "make",
    "modify",
    "reboot",
    "refuel",
    "remove",
    "reseat",
    "reset",
    "resume",
    "return",
    "send",
    "set",
    "toggle",
    "update",
}
READ_VERBS = {"calculate", "can", "check", "find", "get", "list", "run", "search"}
HANDOFF_VERBS = {"transfer"}
GENERIC_RESOURCE_TOKENS = {
    "agent",
    "available",
    "by",
    "details",
    "for",
    "from",
    "human",
    "id",
    "items",
    "pending",
    "service",
    "status",
    "the",
    "to",
    "user",
    "with",
}
CONFIRMATION_PATTERN = re.compile(
    r"\b(yes|confirm|confirmed|proceed|go ahead|please do|do it|sounds good|that is correct)\b",
    flags=re.IGNORECASE,
)
TOOL_TOKEN_PATTERN = re.compile(r"[a-z0-9]+")

TRACE_FEATURES = [
    "domain_retail",
    "domain_telecom",
    "log_messages",
    "log_user_turns",
    "log_assistant_text_turns",
    "log_tool_calls",
    "tool_error_fraction",
    "abnormal_termination",
    "no_tool_call",
]
SEMANTIC_FEATURES = [
    "has_write",
    "log_write_count",
    "write_fraction",
    "log_unique_write_tools",
    "repeated_write_fraction",
    "failed_write_fraction",
    "post_write_read_fraction",
    "confirmed_write_fraction",
    "retry_after_error_fraction",
    "log_write_argument_footprint",
    "log_write_resource_count",
    "last_effect_is_write",
    "log_handoff_count",
    "unclassified_tool_fraction",
]


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def stable_fold(group_id: str, folds: int) -> int:
    digest = hashlib.sha256(group_id.encode("utf-8")).digest()
    return int.from_bytes(digest[:8], "big") % folds


def write_csv(path: Path, rows: Sequence[dict[str, Any]], fields: Sequence[str] | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if fields is None:
        if not rows:
            raise ValueError(f"cannot infer fields for empty table: {path}")
        fields = list(rows[0])
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(fields), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def leaf_count(value: Any) -> int:
    if isinstance(value, dict):
        return sum(leaf_count(child) for child in value.values())
    if isinstance(value, list):
        return sum(leaf_count(child) for child in value)
    return 1


def tool_tokens(name: str) -> list[str]:
    return TOOL_TOKEN_PATTERN.findall(name.lower())


def tool_class(name: str) -> str:
    tokens = tool_tokens(name)
    if not tokens:
        return "other"
    first = tokens[0]
    if first in WRITE_VERBS:
        return "write"
    if first in READ_VERBS:
        return "read"
    if first in HANDOFF_VERBS:
        return "handoff"
    return "other"


def resource_tokens(name: str) -> frozenset[str]:
    excluded = WRITE_VERBS | READ_VERBS | HANDOFF_VERBS | GENERIC_RESOURCE_TOKENS
    return frozenset(token for token in tool_tokens(name) if token not in excluded and len(token) >= 3)


def canonical_signature(name: str, arguments: Any) -> str:
    return name + "|" + json.dumps(arguments or {}, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def extract_events(messages: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    error_by_id: dict[str, bool] = {}
    for message in messages:
        if message.get("role") == "tool" and message.get("id") is not None:
            error_by_id[str(message["id"])] = bool(message.get("error", False))

    events: list[dict[str, Any]] = []
    last_user_text = ""
    for message in messages:
        if message.get("role") == "user" and isinstance(message.get("content"), str):
            last_user_text = message["content"]
        for call in message.get("tool_calls") or []:
            name = str(call.get("name") or "")
            arguments = call.get("arguments") or {}
            call_id = str(call.get("id") or "")
            category = tool_class(name)
            events.append(
                {
                    "name": name,
                    "arguments": arguments,
                    "signature": canonical_signature(name, arguments),
                    "category": category,
                    "resources": resource_tokens(name),
                    "error": int(error_by_id.get(call_id, False)),
                    "confirmed": int(bool(CONFIRMATION_PATTERN.search(last_user_text))),
                }
            )
    return events


def hash_bucket(name: str, buckets: int) -> int:
    digest = hashlib.sha256(name.encode("utf-8")).digest()
    return int.from_bytes(digest[:8], "big") % buckets


def extract_features(simulation: dict[str, Any], domain: str, buckets: int) -> dict[str, float | int]:
    messages = simulation.get("messages") or []
    events = extract_events(messages)
    writes = [event for event in events if event["category"] == "write"]
    reads = [event for event in events if event["category"] == "read"]
    handoffs = [event for event in events if event["category"] == "handoff"]
    other = [event for event in events if event["category"] == "other"]

    user_turns = sum(message.get("role") == "user" for message in messages)
    assistant_text_turns = sum(
        message.get("role") == "assistant" and bool(message.get("content")) for message in messages
    )
    tool_errors = sum(event["error"] for event in events)
    successful_writes = [event for event in writes if not event["error"]]

    signature_counts = Counter(event["signature"] for event in writes)
    repeated_writes = sum(max(0, count - 1) for count in signature_counts.values())

    verified = 0
    retry_after_error = 0
    for index, event in enumerate(events):
        if event["category"] != "write":
            continue
        later = events[index + 1 :]
        if not event["error"] and event["resources"]:
            if any(
                candidate["category"] == "read"
                and bool(event["resources"] & candidate["resources"])
                for candidate in later
            ):
                verified += 1
        earlier_failed = [
            candidate
            for candidate in events[:index]
            if candidate["category"] == "write" and candidate["error"]
        ]
        if event["resources"] and any(event["resources"] & prior["resources"] for prior in earlier_failed):
            retry_after_error += 1

    resource_sets = {tuple(sorted(event["resources"])) for event in writes if event["resources"]}
    argument_footprint = sum(leaf_count(event["arguments"]) for event in writes)
    hashed = [0] * buckets
    for event in events:
        hashed[hash_bucket(event["name"], buckets)] += 1

    termination_reason = str(simulation.get("termination_reason") or "")
    normal_terminations = {"user_stop", "agent_stop", "end_conversation"}
    values: dict[str, float | int] = {
        "domain_retail": int(domain == "retail"),
        "domain_telecom": int(domain == "telecom"),
        "log_messages": math.log1p(len(messages)),
        "log_user_turns": math.log1p(user_turns),
        "log_assistant_text_turns": math.log1p(assistant_text_turns),
        "log_tool_calls": math.log1p(len(events)),
        "tool_error_fraction": tool_errors / len(events) if events else 0.0,
        "abnormal_termination": int(termination_reason not in normal_terminations),
        "no_tool_call": int(not events),
        "has_write": int(bool(writes)),
        "log_write_count": math.log1p(len(writes)),
        "write_fraction": len(writes) / len(events) if events else 0.0,
        "log_unique_write_tools": math.log1p(len({event["name"] for event in writes})),
        "repeated_write_fraction": repeated_writes / len(writes) if writes else 0.0,
        "failed_write_fraction": sum(event["error"] for event in writes) / len(writes) if writes else 0.0,
        "post_write_read_fraction": verified / len(successful_writes) if successful_writes else 0.0,
        "confirmed_write_fraction": sum(event["confirmed"] for event in writes) / len(writes) if writes else 0.0,
        "retry_after_error_fraction": retry_after_error / len(writes) if writes else 0.0,
        "log_write_argument_footprint": math.log1p(argument_footprint),
        "log_write_resource_count": math.log1p(len(resource_sets)),
        "last_effect_is_write": int(bool(events) and events[-1]["category"] == "write"),
        "log_handoff_count": math.log1p(len(handoffs)),
        "unclassified_tool_fraction": len(other) / len(events) if events else 0.0,
        "raw_message_count": len(messages),
        "raw_tool_call_count": len(events),
        "raw_write_count": len(writes),
        "raw_read_count": len(reads),
        "raw_tool_error_count": tool_errors,
        "raw_verified_write_count": verified,
    }
    for index, count in enumerate(hashed):
        values[f"tool_hash_{index:02d}"] = math.log1p(count)
    return values


def source_path(raw_dir: Path, source: dict[str, Any]) -> Path:
    return raw_dir / source["local_filename"]


def load_rows(config: dict[str, Any], raw_dir: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    rows: list[dict[str, Any]] = []
    source_checks: list[dict[str, Any]] = []
    folds = int(config["crossfit"]["task_folds"])
    buckets = int(config["features"]["tool_hash_buckets"])

    for source in config["sources"]:
        path = source_path(raw_dir, source)
        if not path.exists():
            raise FileNotFoundError(f"missing raw source: {path}")
        actual_hash = sha256_file(path)
        if actual_hash != source["sha256"]:
            raise ValueError(f"source hash mismatch for {path.name}: {actual_hash}")
        payload = json.loads(path.read_text(encoding="utf-8"))
        simulations = payload.get("simulations") or []
        if len(simulations) != int(source["expected_trajectories"]):
            raise ValueError(f"trajectory count mismatch for {path.name}")
        task_counts = Counter(str(simulation["task_id"]) for simulation in simulations)
        if len(task_counts) != int(source["expected_tasks"]):
            raise ValueError(f"task count mismatch for {path.name}")
        if set(task_counts.values()) != {int(source["expected_trials_per_task"])}:
            raise ValueError(f"trial multiplicity mismatch for {path.name}: {set(task_counts.values())}")

        successes = 0
        missing_rewards = 0
        for simulation in simulations:
            raw_reward = (simulation.get("reward_info") or {}).get("reward")
            missing_reward = raw_reward is None
            success = int(raw_reward == 1.0)
            successes += success
            missing_rewards += int(missing_reward)
            task_id = str(simulation["task_id"])
            group_id = f"{source['domain']}:{task_id}"
            row_id = f"{source['model_id']}:{source['domain']}:{task_id}:trial-{simulation['trial']}"
            features = extract_features(simulation, source["domain"], buckets)
            rows.append(
                {
                    "row_id": row_id,
                    "model_id": source["model_id"],
                    "model_label": source["model_label"],
                    "domain": source["domain"],
                    "task_id": task_id,
                    "group_id": group_id,
                    "task_fold": stable_fold(group_id, folds),
                    "trial": int(simulation["trial"]),
                    "success": success,
                    "reward_was_missing": int(missing_reward),
                    "termination_reason": str(simulation.get("termination_reason") or ""),
                    **features,
                }
            )
        pass_rate = successes / len(simulations)
        if not math.isclose(pass_rate, float(source["expected_pass_rate"]), abs_tol=5e-12):
            raise ValueError(
                f"pass-rate mismatch for {path.name}: {pass_rate} != {source['expected_pass_rate']}"
            )
        source_checks.append(
            {
                "source_id": source["source_id"],
                "model_id": source["model_id"],
                "domain": source["domain"],
                "sha256": actual_hash,
                "trajectories": len(simulations),
                "tasks": len(task_counts),
                "missing_rewards_counted_as_official_failures": missing_rewards,
                "pass_rate": pass_rate,
                "source_url": source["url"],
            }
        )

    rows.sort(key=lambda row: row["row_id"])
    if len({row["row_id"] for row in rows}) != len(rows):
        raise AssertionError("row identifiers are not unique")
    group_sizes = Counter(row["group_id"] for row in rows)
    expected_group_size = len(config["models"]) * int(config["sources"][0]["expected_trials_per_task"])
    if set(group_sizes.values()) != {expected_group_size}:
        raise AssertionError(f"task clusters are not balanced: {set(group_sizes.values())}")
    return rows, source_checks


def sigmoid(values: np.ndarray) -> np.ndarray:
    output = np.empty_like(values, dtype=float)
    positive = values >= 0
    output[positive] = 1.0 / (1.0 + np.exp(-values[positive]))
    exp_values = np.exp(values[~positive])
    output[~positive] = exp_values / (1.0 + exp_values)
    return output


def fit_ridge_logistic(
    matrix: np.ndarray,
    outcome: np.ndarray,
    penalty: float,
    max_iter: int,
    tolerance: float,
) -> tuple[np.ndarray, int]:
    beta = np.zeros(matrix.shape[1], dtype=float)
    penalty_vector = np.full(matrix.shape[1], penalty, dtype=float)
    penalty_vector[0] = 0.0
    for iteration in range(1, max_iter + 1):
        probabilities = np.clip(sigmoid(matrix @ beta), 1e-8, 1 - 1e-8)
        gradient = matrix.T @ (probabilities - outcome) + penalty_vector * beta
        weights = probabilities * (1.0 - probabilities)
        hessian = matrix.T @ (weights[:, None] * matrix)
        hessian.flat[:: hessian.shape[0] + 1] += penalty_vector + 1e-8
        step = np.linalg.solve(hessian, gradient)
        beta -= step
        if float(np.max(np.abs(step))) < tolerance:
            return beta, iteration
    raise RuntimeError("ridge logistic regression did not converge")


def design_matrices(
    train_rows: Sequence[dict[str, Any]],
    test_rows: Sequence[dict[str, Any]],
    features: Sequence[str],
) -> tuple[np.ndarray, np.ndarray]:
    train = np.asarray([[float(row[name]) for name in features] for row in train_rows], dtype=float)
    test = np.asarray([[float(row[name]) for name in features] for row in test_rows], dtype=float)
    means = train.mean(axis=0)
    scales = train.std(axis=0)
    scales[scales < 1e-12] = 1.0
    train = (train - means) / scales
    test = (test - means) / scales
    return (
        np.column_stack([np.ones(len(train)), train]),
        np.column_stack([np.ones(len(test)), test]),
    )


def crossfit(
    rows: Sequence[dict[str, Any]], config: dict[str, Any], cohort: str
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    folds = int(config["crossfit"]["task_folds"])
    penalty = float(config["model"]["ridge_penalty"])
    max_iter = int(config["model"]["max_iter"])
    tolerance = float(config["model"]["tolerance"])
    buckets = int(config["features"]["tool_hash_buckets"])
    hash_features = [f"tool_hash_{index:02d}" for index in range(buckets)]
    feature_sets = {
        "trace": TRACE_FEATURES,
        "semantic": TRACE_FEATURES + SEMANTIC_FEATURES,
        "toolbag": TRACE_FEATURES + hash_features,
        "full": TRACE_FEATURES + hash_features + SEMANTIC_FEATURES,
    }
    models = [model["id"] for model in config["models"]]
    predictions: list[dict[str, Any]] = []
    split_rows: list[dict[str, Any]] = []

    for held_model in models:
        for held_fold in range(folds):
            test_rows = [
                row for row in rows if row["model_id"] == held_model and row["task_fold"] == held_fold
            ]
            train_rows = [
                row for row in rows if row["model_id"] != held_model and row["task_fold"] != held_fold
            ]
            if not test_rows:
                raise AssertionError(f"empty test block: {held_model}, fold {held_fold}")
            if {row["model_id"] for row in train_rows} & {held_model}:
                raise AssertionError("held-out model leaked into training")
            if {row["group_id"] for row in train_rows} & {row["group_id"] for row in test_rows}:
                raise AssertionError("held-out task leaked into training")

            domain_rates = {}
            for domain in ("airline", "retail", "telecom"):
                domain_outcomes = [row["success"] for row in train_rows if row["domain"] == domain]
                domain_rates[domain] = (sum(domain_outcomes) + 1.0) / (len(domain_outcomes) + 2.0)
            for row in test_rows:
                predictions.append(
                    {
                        "cohort": cohort,
                        "row_id": row["row_id"],
                        "group_id": row["group_id"],
                        "model_id": row["model_id"],
                        "domain": row["domain"],
                        "task_fold": row["task_fold"],
                        "success": row["success"],
                        "reward_was_missing": row["reward_was_missing"],
                        "method": "prevalence",
                        "probability": domain_rates[row["domain"]],
                    }
                )

            iterations = {}
            outcome = np.asarray([row["success"] for row in train_rows], dtype=float)
            for method, features in feature_sets.items():
                train_matrix, test_matrix = design_matrices(train_rows, test_rows, features)
                beta, iteration_count = fit_ridge_logistic(
                    train_matrix, outcome, penalty, max_iter, tolerance
                )
                iterations[method] = iteration_count
                probabilities = np.clip(sigmoid(test_matrix @ beta), 1e-6, 1 - 1e-6)
                for row, probability in zip(test_rows, probabilities, strict=True):
                    predictions.append(
                        {
                            "cohort": cohort,
                            "row_id": row["row_id"],
                            "group_id": row["group_id"],
                            "model_id": row["model_id"],
                            "domain": row["domain"],
                            "task_fold": row["task_fold"],
                            "success": row["success"],
                            "reward_was_missing": row["reward_was_missing"],
                            "method": method,
                            "probability": float(probability),
                        }
                    )
            split_rows.append(
                {
                    "cohort": cohort,
                    "heldout_model": held_model,
                    "heldout_task_fold": held_fold,
                    "training_rows": len(train_rows),
                    "test_rows": len(test_rows),
                    "training_models": ";".join(sorted({row["model_id"] for row in train_rows})),
                    "training_task_groups": len({row["group_id"] for row in train_rows}),
                    "test_task_groups": len({row["group_id"] for row in test_rows}),
                    **{f"{method}_iterations": count for method, count in iterations.items()},
                }
            )

    expected = len(rows) * 5
    if len(predictions) != expected:
        raise AssertionError(f"expected {expected} cross-fit predictions, got {len(predictions)}")
    per_row = Counter(row["row_id"] for row in predictions)
    if set(per_row.values()) != {5}:
        raise AssertionError("each row must receive exactly one prediction from each method")
    predictions.sort(key=lambda row: (row["cohort"], row["row_id"], row["method"]))
    return predictions, split_rows


def auc_score(outcomes: Sequence[int], probabilities: Sequence[float]) -> float:
    y = np.asarray(outcomes, dtype=int)
    p = np.asarray(probabilities, dtype=float)
    positive = int(y.sum())
    negative = len(y) - positive
    if positive == 0 or negative == 0:
        return float("nan")
    order = np.argsort(p, kind="mergesort")
    ranks = np.empty(len(p), dtype=float)
    start = 0
    while start < len(order):
        end = start + 1
        while end < len(order) and p[order[end]] == p[order[start]]:
            end += 1
        average_rank = (start + 1 + end) / 2.0
        ranks[order[start:end]] = average_rank
        start = end
    return float((ranks[y == 1].sum() - positive * (positive + 1) / 2.0) / (positive * negative))


def percentile(values: np.ndarray, probability: float) -> float:
    return float(np.quantile(values, probability, method="linear"))


def cluster_values(
    predictions: Sequence[dict[str, Any]], method: str, value_function
) -> dict[str, list[float]]:
    grouped: dict[str, list[float]] = defaultdict(list)
    for row in predictions:
        if row["method"] == method:
            grouped[row["group_id"]].append(value_function(row))
    return grouped


def bootstrap_cluster_mean(
    grouped_values: dict[str, list[float]], draws: int, seed: int
) -> tuple[float, float, float]:
    group_means = np.asarray([np.mean(grouped_values[key]) for key in sorted(grouped_values)], dtype=float)
    rng = np.random.default_rng(seed)
    replicate_means = np.empty(draws, dtype=float)
    for start in range(0, draws, 500):
        batch = min(500, draws - start)
        indices = rng.integers(0, len(group_means), size=(batch, len(group_means)))
        replicate_means[start : start + batch] = group_means[indices].mean(axis=1)
    return float(group_means.mean()), percentile(replicate_means, 0.025), percentile(replicate_means, 0.975)


def metric_rows(
    predictions: Sequence[dict[str, Any]], config: dict[str, Any]
) -> list[dict[str, Any]]:
    draws = int(config["inference"]["bootstrap_replicates"])
    seed = int(config["inference"]["bootstrap_seed"])
    rows: list[dict[str, Any]] = []
    methods = ("prevalence", "trace", "semantic", "toolbag", "full")
    for cohort in sorted({row["cohort"] for row in predictions}):
        cohort_rows = [row for row in predictions if row["cohort"] == cohort]
        for method_index, method in enumerate(methods):
            selected = [row for row in cohort_rows if row["method"] == method]
            y = np.asarray([row["success"] for row in selected], dtype=float)
            p = np.asarray([row["probability"] for row in selected], dtype=float)
            brier_groups = cluster_values(
                cohort_rows, method, lambda row: (row["probability"] - row["success"]) ** 2
            )
            brier, lower, upper = bootstrap_cluster_mean(
                brier_groups, draws, seed + 100 * method_index + (0 if cohort == "official" else 5000)
            )
            rows.append(
                {
                    "cohort": cohort,
                    "method": method,
                    "trajectories": len(selected),
                    "task_clusters": len({row["group_id"] for row in selected}),
                    "success_rate": float(y.mean()),
                    "brier": brier,
                    "brier_ci_lower": lower,
                    "brier_ci_upper": upper,
                    "log_loss": float(np.mean(-(y * np.log(p) + (1 - y) * np.log(1 - p)))),
                    "auc": auc_score(y.astype(int).tolist(), p.tolist()),
                }
            )
    return rows


def paired_contrasts(
    predictions: Sequence[dict[str, Any]], config: dict[str, Any]
) -> list[dict[str, Any]]:
    draws = int(config["inference"]["bootstrap_replicates"])
    seed = int(config["inference"]["bootstrap_seed"])
    pairs = [
        ("semantic", "trace", "semantic_minus_trace"),
        ("full", "toolbag", "full_minus_toolbag"),
        ("toolbag", "trace", "toolbag_minus_trace"),
        ("full", "semantic", "full_minus_semantic"),
    ]
    output: list[dict[str, Any]] = []
    for cohort_index, cohort in enumerate(sorted({row["cohort"] for row in predictions})):
        cohort_rows = [row for row in predictions if row["cohort"] == cohort]
        lookup = {(row["row_id"], row["method"]): row for row in cohort_rows}
        row_ids = sorted({row["row_id"] for row in cohort_rows})
        for pair_index, (first, second, label) in enumerate(pairs):
            grouped: dict[str, list[float]] = defaultdict(list)
            for row_id in row_ids:
                left = lookup[(row_id, first)]
                right = lookup[(row_id, second)]
                outcome = float(left["success"])
                difference = (float(left["probability"]) - outcome) ** 2 - (
                    float(right["probability"]) - outcome
                ) ** 2
                grouped[left["group_id"]].append(difference)
            estimate, lower, upper = bootstrap_cluster_mean(
                grouped, draws, seed + 9000 + cohort_index * 100 + pair_index
            )
            output.append(
                {
                    "cohort": cohort,
                    "contrast": label,
                    "first_method": first,
                    "second_method": second,
                    "metric": "brier_difference",
                    "difference": estimate,
                    "ci_lower": lower,
                    "ci_upper": upper,
                    "interpretation": "negative favors first method",
                    "task_clusters": len(grouped),
                }
            )
    return output


def stratified_rows(predictions: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    output = []
    official = [row for row in predictions if row["cohort"] == "official"]
    for stratum_type in ("model_id", "domain"):
        for stratum in sorted({row[stratum_type] for row in official}):
            for method in ("prevalence", "trace", "semantic", "toolbag", "full"):
                selected = [
                    row for row in official if row[stratum_type] == stratum and row["method"] == method
                ]
                y = [int(row["success"]) for row in selected]
                p = [float(row["probability"]) for row in selected]
                output.append(
                    {
                        "stratum_type": stratum_type,
                        "stratum": stratum,
                        "method": method,
                        "trajectories": len(selected),
                        "task_clusters": len({row["group_id"] for row in selected}),
                        "success_rate": float(np.mean(y)),
                        "brier": float(np.mean([(left - right) ** 2 for left, right in zip(p, y)])),
                        "auc": auc_score(y, p),
                    }
                )
    return output


def calibration_rows(predictions: Sequence[dict[str, Any]], bins: int = 8) -> list[dict[str, Any]]:
    output = []
    official = [row for row in predictions if row["cohort"] == "official"]
    for method in ("prevalence", "trace", "semantic", "toolbag", "full"):
        selected = sorted(
            [row for row in official if row["method"] == method], key=lambda row: row["probability"]
        )
        for bin_index, indices in enumerate(np.array_split(np.arange(len(selected)), bins)):
            bin_rows = [selected[int(index)] for index in indices]
            output.append(
                {
                    "method": method,
                    "bin": bin_index + 1,
                    "trajectories": len(bin_rows),
                    "mean_predicted": float(np.mean([row["probability"] for row in bin_rows])),
                    "observed_success": float(np.mean([row["success"] for row in bin_rows])),
                }
            )
    return output


def descriptive_bins(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    output: list[dict[str, Any]] = []
    for feature, label in (
        ("raw_write_count", "state-changing calls"),
        ("post_write_read_fraction", "post-write read closure"),
    ):
        ordered = sorted(rows, key=lambda row: (float(row[feature]), row["row_id"]))
        for bin_index, indices in enumerate(np.array_split(np.arange(len(ordered)), 6)):
            selected = [ordered[int(index)] for index in indices]
            output.append(
                {
                    "feature": feature,
                    "feature_label": label,
                    "bin": bin_index + 1,
                    "trajectories": len(selected),
                    "feature_mean": float(np.mean([row[feature] for row in selected])),
                    "success_rate": float(np.mean([row["success"] for row in selected])),
                    "task_clusters": len({row["group_id"] for row in selected}),
                }
            )
    return output


def model_domain_rates(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    output = []
    for model_id in sorted({row["model_id"] for row in rows}):
        for domain in ("airline", "retail", "telecom"):
            selected = [row for row in rows if row["model_id"] == model_id and row["domain"] == domain]
            output.append(
                {
                    "model_id": model_id,
                    "model_label": selected[0]["model_label"],
                    "domain": domain,
                    "trajectories": len(selected),
                    "tasks": len({row["group_id"] for row in selected}),
                    "success_rate": float(np.mean([row["success"] for row in selected])),
                    "mean_tool_calls": float(np.mean([row["raw_tool_call_count"] for row in selected])),
                    "mean_write_calls": float(np.mean([row["raw_write_count"] for row in selected])),
                }
            )
    return output


def plot_observed_cohort(
    rows: Sequence[dict[str, Any]], domain_rates: Sequence[dict[str, Any]], bins: Sequence[dict[str, Any]]
) -> None:
    configure()
    figure, axes = plt.subplots(1, 2, figsize=(7.10, 3.0), gridspec_kw={"width_ratios": [1.05, 1.35]})
    heatmap, burden = axes
    models = sorted({row["model_id"] for row in domain_rates})
    labels = {
        row["model_id"]: row["model_label"] for row in domain_rates
    }
    domains = ["airline", "retail", "telecom"]
    matrix = np.asarray(
        [
            [next(row["success_rate"] for row in domain_rates if row["model_id"] == model and row["domain"] == domain) for domain in domains]
            for model in models
        ]
    )
    heatmap.imshow(matrix, cmap="Blues", vmin=0.55, vmax=1.0, aspect="auto")
    for row_index in range(len(models)):
        for column_index in range(len(domains)):
            value = matrix[row_index, column_index]
            heatmap.text(
                column_index,
                row_index,
                f"{100 * value:.1f}%",
                ha="center",
                va="center",
                color="white" if value > 0.78 else COLORS["ink"],
                fontweight="bold",
                fontsize=7.4,
            )
    heatmap.set_xticks(range(len(domains)), [domain.title() for domain in domains])
    heatmap.set_yticks(range(len(models)), [labels[model] for model in models])
    heatmap.set_title("Released benchmark outcomes")
    heatmap.tick_params(length=0)
    for spine in heatmap.spines.values():
        spine.set_visible(False)
    selected_bins = [row for row in bins if row["feature"] == "raw_write_count"]
    burden.plot(
        [row["feature_mean"] for row in selected_bins],
        [row["success_rate"] for row in selected_bins],
        marker="o",
        color=COLORS["blue"],
    )
    for row in selected_bins:
        burden.annotate(
            f"Q{row['bin']}",
            (row["feature_mean"], row["success_rate"]),
            xytext=(0, 7),
            textcoords="offset points",
            ha="center",
            fontsize=6.8,
            color=COLORS["ink"],
        )
    burden.set(
        xlabel="Mean state-changing calls in burden sextile",
        ylabel="Observed task success",
        title="State-change burden is not monotonic",
        ylim=(0.58, 0.93),
    )
    burden.set_yticks([0.6, 0.7, 0.8, 0.9], ["60%", "70%", "80%", "90%"])
    burden.grid(axis="y", visible=True)
    burden.text(
        0.02,
        0.04,
        f"{len(rows):,} trajectories; {len({row['group_id'] for row in rows})} tasks",
        transform=burden.transAxes,
        fontsize=7.0,
        color=COLORS["gray"],
    )
    label_panels(axes)
    figure.subplots_adjust(left=0.14, right=0.985, bottom=0.18, top=0.84, wspace=0.52)
    save_figure(figure, FIGURES, "observed_tau2_cohort")


def plot_crossfit_evidence(
    metrics: Sequence[dict[str, Any]], contrasts: Sequence[dict[str, Any]], strata: Sequence[dict[str, Any]]
) -> None:
    configure()
    figure, axes = plt.subplots(1, 2, figsize=(7.10, 3.05), gridspec_kw={"width_ratios": [0.92, 1.35]})
    score_axis, delta_axis = axes
    display_methods = ["prevalence", "trace", "semantic", "toolbag", "full"]
    labels = {
        "prevalence": "Domain rate",
        "trace": "Trace",
        "semantic": "Trace + semantic",
        "toolbag": "Trace + tool bag",
        "full": "Full",
    }
    colors = {
        "prevalence": COLORS["mid_gray"],
        "trace": COLORS["gray"],
        "semantic": COLORS["blue"],
        "toolbag": COLORS["orange"],
        "full": COLORS["green"],
    }
    official_metrics = {row["method"]: row for row in metrics if row["cohort"] == "official"}
    positions = np.arange(len(display_methods))[::-1]
    for position, method in zip(positions, display_methods, strict=True):
        row = official_metrics[method]
        score_axis.errorbar(
            row["brier"],
            position,
            xerr=[[row["brier"] - row["brier_ci_lower"]], [row["brier_ci_upper"] - row["brier"]]],
            fmt="o",
            color=colors[method],
            capsize=3,
            markersize=5.5,
        )
    score_axis.set_yticks(positions, [labels[method] for method in display_methods])
    score_axis.set(xlabel="Cross-fitted Brier score (lower is better)", title="Cross-fitted prediction error")
    score_axis.grid(axis="x", visible=True)

    official = [row for row in strata if row["stratum_type"] == "model_id"]
    models = sorted({row["stratum"] for row in official})
    model_labels = {
        "claude-sonnet-4-5": "Claude Sonnet 4.5",
        "glm-5": "GLM-5",
        "gpt-5-2-high": "GPT-5.2 high",
        "qwen3-5-397b": "Qwen3.5-397B",
    }
    y_positions = np.arange(len(models))[::-1]
    for y_position, model in zip(y_positions, models, strict=True):
        semantic = next(row for row in official if row["stratum"] == model and row["method"] == "semantic")
        trace = next(row for row in official if row["stratum"] == model and row["method"] == "trace")
        delta_axis.scatter(semantic["brier"] - trace["brier"], y_position, color=COLORS["blue"], marker="o")
    overall = next(
        row for row in contrasts if row["cohort"] == "official" and row["contrast"] == "semantic_minus_trace"
    )
    delta_axis.errorbar(
        overall["difference"],
        -1.0,
        xerr=[[overall["difference"] - overall["ci_lower"]], [overall["ci_upper"] - overall["difference"]]],
        fmt="D",
        color=COLORS["ink"],
        capsize=3,
        markersize=5.5,
    )
    delta_axis.axvline(0, color=COLORS["ink"], linestyle=(0, (4, 2)), linewidth=1.0)
    delta_axis.set_yticks(
        list(y_positions) + [-1.0],
        [model_labels[model] for model in models] + ["Overall (95% CI)"],
    )
    delta_axis.set(
        xlabel="Brier difference: semantic - trace",
        title="No reliable semantic gain",
        ylim=(-1.7, len(models) - 0.25),
    )
    delta_axis.grid(axis="x", visible=True)
    delta_axis.xaxis.set_major_formatter(FormatStrFormatter("%.3f"))
    delta_axis.text(
        0.02,
        0.02,
        "left favors semantic summary; right favors trace",
        transform=delta_axis.transAxes,
        fontsize=6.8,
        color=COLORS["gray"],
    )
    label_panels(axes)
    figure.subplots_adjust(left=0.15, right=0.985, bottom=0.19, top=0.84, wspace=0.55)
    save_figure(figure, FIGURES, "crossfit_semantic_evidence")


def plot_calibration_and_closure(
    calibration: Sequence[dict[str, Any]], bins: Sequence[dict[str, Any]]
) -> None:
    configure()
    figure, axes = plt.subplots(1, 2, figsize=(7.10, 3.0))
    calibration_axis, closure_axis = axes
    styles = {
        "trace": (COLORS["gray"], "o", "--", "Interaction trace"),
        "semantic": (COLORS["blue"], "s", "-", "+ semantic summary"),
        "toolbag": (COLORS["orange"], "^", ":", "+ hashed tool bag"),
        "full": (COLORS["green"], "D", "-.", "Tool bag + semantic"),
    }
    calibration_axis.plot([0, 1], [0, 1], color=COLORS["mid_gray"], linewidth=1.0, linestyle=(0, (3, 2)))
    for method, (color, marker, linestyle, label) in styles.items():
        selected = [row for row in calibration if row["method"] == method]
        calibration_axis.plot(
            [row["mean_predicted"] for row in selected],
            [row["observed_success"] for row in selected],
            color=color,
            marker=marker,
            linestyle=linestyle,
            label=label,
        )
    calibration_axis.set(
        xlabel="Predicted success",
        ylabel="Observed success",
        title="Cross-fitted calibration",
        xlim=(0.35, 1.01),
        ylim=(0.35, 1.01),
    )
    calibration_axis.grid(True)
    calibration_axis.legend(loc="lower right")

    closure = [row for row in bins if row["feature"] == "post_write_read_fraction"]
    closure_axis.plot(
        [row["feature_mean"] for row in closure],
        [row["success_rate"] for row in closure],
        color=COLORS["blue"],
        marker="o",
    )
    annotation_offsets = {
        1: (-14, 7),
        2: (8, 8),
        3: (-14, 7),
        4: (8, 8),
        5: (0, -17),
        6: (0, 8),
    }
    for row in closure:
        dx, dy = annotation_offsets[int(row["bin"])]
        closure_axis.annotate(
            f"Q{row['bin']}",
            (row["feature_mean"], row["success_rate"]),
            xytext=(dx, dy),
            textcoords="offset points",
            ha="center",
            fontsize=6.8,
        )
    closure_axis.set(
        xlabel="Mean post-write read closure in sextile",
        ylabel="Observed task success",
        title="Read-after-write closure is descriptive",
        xlim=(-0.04, 0.52),
        ylim=(0.58, 0.93),
    )
    closure_axis.set_yticks([0.6, 0.7, 0.8, 0.9], ["60%", "70%", "80%", "90%"])
    closure_axis.grid(axis="y", visible=True)
    label_panels(axes)
    figure.subplots_adjust(left=0.105, right=0.985, bottom=0.18, top=0.84, wspace=0.35)
    save_figure(figure, FIGURES, "calibration_and_evidence_closure")


def build_manifest(config: dict[str, Any], source_checks: Sequence[dict[str, Any]]) -> None:
    payload_paths = [
        RESULTS / "trajectory_features.csv",
        RESULTS / "crossfit_predictions.csv",
        RESULTS / "split_inventory.csv",
        RESULTS / "method_metrics.csv",
        RESULTS / "paired_contrasts.csv",
        RESULTS / "stratified_metrics.csv",
        RESULTS / "calibration_bins.csv",
        RESULTS / "descriptive_bins.csv",
        RESULTS / "model_domain_rates.csv",
        RESULTS / "source_checks.csv",
        RESULTS / "summary.json",
        FIGURES / "observed_tau2_cohort.pdf",
        FIGURES / "observed_tau2_cohort.png",
        FIGURES / "crossfit_semantic_evidence.pdf",
        FIGURES / "crossfit_semantic_evidence.png",
        FIGURES / "calibration_and_evidence_closure.pdf",
        FIGURES / "calibration_and_evidence_closure.png",
    ]
    manifest = {
        "study_id": config["study_id"],
        "execution_date": config["execution_date"],
        "command": "python run_realworld_study.py --raw-dir <verified-source-directory>",
        "environment": {
            "python": platform.python_version(),
            "numpy": np.__version__,
            "matplotlib": plt.matplotlib.__version__,
        },
        "inputs": {
            "study_config.json": sha256_file(CONFIG_PATH),
            "run_realworld_study.py": sha256_file(Path(__file__)),
            "raw_sources": list(source_checks),
        },
        "outputs": {
            str(path.relative_to(HERE)).replace("\\", "/"): sha256_file(path) for path in payload_paths
        },
    }
    (RESULTS / "run_manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw-dir", required=True, type=Path)
    args = parser.parse_args()
    config = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    rows, source_checks = load_rows(config, args.raw_dir)

    official_predictions, official_splits = crossfit(rows, config, "official")
    numeric_rows = [row for row in rows if not row["reward_was_missing"]]
    numeric_predictions, numeric_splits = crossfit(numeric_rows, config, "numeric_reward_only")
    predictions = official_predictions + numeric_predictions
    splits = official_splits + numeric_splits

    metrics = metric_rows(predictions, config)
    contrasts = paired_contrasts(predictions, config)
    strata = stratified_rows(predictions)
    calibration = calibration_rows(predictions)
    bins = descriptive_bins(rows)
    domain_rates = model_domain_rates(rows)

    feature_fields = [
        "row_id",
        "model_id",
        "model_label",
        "domain",
        "task_id",
        "group_id",
        "task_fold",
        "trial",
        "success",
        "reward_was_missing",
        "termination_reason",
        *TRACE_FEATURES,
        *SEMANTIC_FEATURES,
        *[f"tool_hash_{index:02d}" for index in range(int(config["features"]["tool_hash_buckets"]))],
        "raw_message_count",
        "raw_tool_call_count",
        "raw_write_count",
        "raw_read_count",
        "raw_tool_error_count",
        "raw_verified_write_count",
    ]
    write_csv(RESULTS / "trajectory_features.csv", rows, feature_fields)
    write_csv(RESULTS / "crossfit_predictions.csv", predictions)
    write_csv(RESULTS / "split_inventory.csv", splits)
    write_csv(RESULTS / "method_metrics.csv", metrics)
    write_csv(RESULTS / "paired_contrasts.csv", contrasts)
    write_csv(RESULTS / "stratified_metrics.csv", strata)
    write_csv(RESULTS / "calibration_bins.csv", calibration)
    write_csv(RESULTS / "descriptive_bins.csv", bins)
    write_csv(RESULTS / "model_domain_rates.csv", domain_rates)
    write_csv(RESULTS / "source_checks.csv", source_checks)

    plot_observed_cohort(rows, domain_rates, bins)
    plot_crossfit_evidence(metrics, contrasts, strata)
    plot_calibration_and_closure(calibration, bins)

    official_metric_lookup = {row["method"]: row for row in metrics if row["cohort"] == "official"}
    official_contrast_lookup = {
        row["contrast"]: row for row in contrasts if row["cohort"] == "official"
    }
    numeric_contrast_lookup = {
        row["contrast"]: row for row in contrasts if row["cohort"] == "numeric_reward_only"
    }
    summary = {
        "study_id": config["study_id"],
        "study_kind": "retrospective cross-fitted analysis of independently graded public tau-bench trajectories",
        "trajectories": len(rows),
        "task_clusters": len({row["group_id"] for row in rows}),
        "models": len({row["model_id"] for row in rows}),
        "domains": len({row["domain"] for row in rows}),
        "trials_per_model_task": int(config["sources"][0]["expected_trials_per_task"]),
        "official_success_rate": float(np.mean([row["success"] for row in rows])),
        "missing_reward_rows_counted_as_official_failures": sum(row["reward_was_missing"] for row in rows),
        "crossfit": {
            "task_folds": int(config["crossfit"]["task_folds"]),
            "held_out_axes": ["task_cluster", "agent_model"],
            "every_trajectory_scored_once_per_method": True,
            "task_overlap_between_train_and_test": False,
            "model_overlap_between_train_and_test": False,
        },
        "primary_metrics": official_metric_lookup,
        "primary_contrasts": official_contrast_lookup,
        "numeric_reward_only_sensitivity": numeric_contrast_lookup,
        "evidence_boundary": {
            "outcome_source": "released tau-bench reward_info.reward; missing rewards follow the submission's stated failure convention",
            "features_use_evaluation_criteria": False,
            "features_use_database_checks": False,
            "features_use_assertions": False,
            "features_use_reward_breakdown": False,
            "causal_claim": False,
            "benchmark_scope_only": True,
        },
    }
    (RESULTS / "summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    build_manifest(config, source_checks)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
