scripts/analyze_study4_ancillary.py

"""Fail-closed descriptive analysis for frozen E11/E12 sensitivities."""

from __future__ import annotations

from collections import Counter
import csv
from hashlib import sha256
import json
from pathlib import Path
import statistics
from typing import Any


ENDPOINTS = (
    "resolved_at_1", "accepted_edit_cell", "applicable_final_patch",
    "exact_modified_file_match", "fail_to_pass",
)


def file_sha(path: Path) -> str:
    return sha256(path.read_bytes()).hexdigest()


def load_experiment(
    root: Path, experiment_id: str, expected: int
) -> tuple[list[dict[str, Any]], str, str]:
    paths = sorted((root / "results/raw" / experiment_id).rglob("final_metrics.json"))
    if len(paths) != expected:
        raise RuntimeError(f"{experiment_id} requires {expected} cells; observed {len(paths)}")
    rows: list[dict[str, Any]] = []
    revisions: set[str] = set()
    run_ids: set[str] = set()
    raw_hasher = sha256()
    for path in paths:
        manifest = json.loads((path.parent / "run_manifest.json").read_text())
        final = json.loads(path.read_text())
        raw_hasher.update((path.parent / "run_manifest.json").read_bytes())
        raw_hasher.update(path.read_bytes())
        identity = manifest["identity"]
        if manifest["run_id"] in run_ids:
            raise RuntimeError(f"duplicate run id: {manifest['run_id']}")
        run_ids.add(manifest["run_id"])
        revisions.add(identity["code_revision"])
        harness, interface = final["harness_id"].split("__", 1)
        if any(len(item.get("after_instances", [])) != 1 for item in final["residency_transitions"]):
            raise RuntimeError(f"non-exclusive residency: {path.parent}")
        rows.append({
            "run_id": manifest["run_id"],
            "task_id": identity["task_id"],
            "model_id": identity["model_id"],
            "harness_id": harness,
            "interface_id": interface,
            "seed": int(identity["seed"]),
            "context_budget": int(identity["context_budget"]),
            **{endpoint: int(bool(final[endpoint])) for endpoint in ENDPOINTS},
            "patch_sha256": final.get("patch_sha256"),
            "trajectory_sha256": final["trajectory_sha256"],
            "total_tokens": int(final["usage"]["total_tokens"]),
            "elapsed_seconds": float(final["elapsed_seconds"]),
            "failure_stage": final["failure_stage"],
        })
    if len(revisions) != 1:
        raise RuntimeError(f"{experiment_id} spans revisions: {revisions}")
    return rows, next(iter(revisions)), raw_hasher.hexdigest()


def reliability(rows: list[dict[str, Any]]) -> dict[str, Any]:
    groups: dict[tuple[str, str, str], list[dict[str, Any]]] = {}
    for row in rows:
        groups.setdefault((row["task_id"], row["model_id"], row["harness_id"]), []).append(row)
    output = []
    for key, group in sorted(groups.items()):
        if sorted(row["seed"] for row in group) != [0, 1, 2]:
            raise RuntimeError(f"reliability seed drift: {key}")
        output.append({
            "task_id": key[0], "model_id": key[1], "harness_id": key[2],
            "resolved_seeds": sum(row["resolved_at_1"] for row in group),
            "accepted_edit_seeds": sum(row["accepted_edit_cell"] for row in group),
            "unanimous_resolution": len({row["resolved_at_1"] for row in group}) == 1,
            "unanimous_patch": len({row["patch_sha256"] for row in group}) == 1,
            "unanimous_trajectory": len({row["trajectory_sha256"] for row in group}) == 1,
        })
    return {
        "cells": len(rows), "groups": len(output), "group_rows": output,
        "resolution_rate": statistics.fmean(row["resolved_at_1"] for row in rows),
        "accepted_edit_rate": statistics.fmean(row["accepted_edit_cell"] for row in rows),
        "unanimous_resolution_rate": statistics.fmean(row["unanimous_resolution"] for row in output),
        "unanimous_patch_rate": statistics.fmean(row["unanimous_patch"] for row in output),
        "unanimous_trajectory_rate": statistics.fmean(row["unanimous_trajectory"] for row in output),
        "resolved_seed_distribution": dict(sorted(Counter(row["resolved_seeds"] for row in output).items())),
    }


def context(rows: list[dict[str, Any]]) -> dict[str, Any]:
    by = {(row["task_id"], row["harness_id"], row["context_budget"]): row for row in rows}
    pairs = []
    for task_id in sorted({row["task_id"] for row in rows}):
        for harness in ("H000", "H007"):
            low, high = by[(task_id, harness, 16384)], by[(task_id, harness, 65536)]
            pairs.append({
                "task_id": task_id, "harness_id": harness,
                **{f"delta_{endpoint}": high[endpoint] - low[endpoint] for endpoint in ENDPOINTS},
                "delta_total_tokens": high["total_tokens"] - low["total_tokens"],
                "delta_elapsed_seconds": high["elapsed_seconds"] - low["elapsed_seconds"],
            })
    return {
        "cells": len(rows), "pairs": len(pairs), "pair_rows": pairs,
        "mean_paired_differences_65536_minus_16384": {
            endpoint: statistics.fmean(row[f"delta_{endpoint}"] for row in pairs)
            for endpoint in ENDPOINTS
        },
        "resolution_by_context": {
            str(context): sum(row["resolved_at_1"] for row in rows if row["context_budget"] == context)
            for context in (16384, 65536)
        },
    }


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader(); writer.writerows(rows)


def analyze(root: Path) -> dict[str, Any]:
    e11, e11_revision, e11_digest = load_experiment(root, "E11", 18)
    e12, e12_revision, e12_digest = load_experiment(root, "E12", 12)
    preflight = json.loads((root / "results/reports/study4_ancillary_preflight.json").read_text())
    if not preflight.get("passed"):
        raise RuntimeError("ancillary preflight did not pass")
    if {e11_revision, e12_revision} != {preflight.get("research_code_revision")}:
        raise RuntimeError("ancillary execution/preflight revision mismatch")
    declared = {
        "E11": json.loads((root / "configs/reliability/E11_repeat_cells.json").read_text()),
        "E12": json.loads((root / "configs/context/E12_context_cells.json").read_text()),
    }
    expected_e11 = {
        (cell["task_id"], cell["model_id"], cell["harness_id"], 65536, seed)
        for cell in declared["E11"]["cells"] for seed in declared["E11"]["seeds"]
    }
    expected_e12 = {
        (cell["task_id"], cell["model_id"], cell["harness_id"], cell["context_budget"], 0)
        for cell in declared["E12"]["cells"]
    }
    for experiment_id, rows, expected_keys in (
        ("E11", e11, expected_e11), ("E12", e12, expected_e12)
    ):
        observed = {
            (row["task_id"], row["model_id"], row["harness_id"], row["context_budget"], row["seed"])
            for row in rows
        }
        if observed != expected_keys:
            raise RuntimeError(f"{experiment_id} raw identities differ from its frozen manifest")
    reliability_result = reliability(e11)
    context_result = context(e12)
    report = {
        "schema_version": 1, "experiment_ids": ["E11", "E12"],
        "input_cells": 30, "reliability": reliability_result,
        "context_scarcity": context_result,
        "execution_revision": e11_revision,
        "raw_manifest_and_metrics_sha256": {"E11": e11_digest, "E12": e12_digest},
        "frozen_manifest_sha256": {
            "E11": file_sha(root / "configs/reliability/E11_repeat_cells.json"),
            "E12": file_sha(root / "configs/context/E12_context_cells.json"),
        },
        "claim_boundary": "descriptive prespecified sensitivities; never pooled into E10 H1",
        "analysis_script_sha256": file_sha(root / "scripts/analyze_study4_ancillary.py"),
    }
    output = root / "results/derived/study4_ancillary"; output.mkdir(parents=True, exist_ok=True)
    write_csv(output / "e11_cells.csv", e11)
    write_csv(output / "e11_reliability_groups.csv", reliability_result["group_rows"])
    write_csv(output / "e12_cells.csv", e12)
    write_csv(output / "e12_context_pairs.csv", context_result["pair_rows"])
    (output / "e11_e12_analysis.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    files = sorted(path for path in output.iterdir() if path.name != "SHA256SUMS.json")
    checksums = {path.name: file_sha(path) for path in files}
    (output / "SHA256SUMS.json").write_text(json.dumps(checksums, indent=2, sort_keys=True) + "\n")
    return {**report, "checksums": checksums}


if __name__ == "__main__":
    root = Path(__file__).resolve().parents[1]
    print(json.dumps(analyze(root), indent=2, sort_keys=True))