scripts/analyze_study4.py

"""Prespecified fail-closed analysis for E10 fresh-task retrieval."""

from __future__ import annotations

import argparse
from collections import Counter
import csv
from hashlib import sha256
import itertools
import json
import math
from pathlib import Path
import random
import statistics
import subprocess
from typing import Any, Sequence

import numpy as np
import pandas as pd

from agent_harness.specs import load_repositories, load_task_split, load_tasks


MODELS = ("M002", "M003", "M004")
HARNESSES = ("H000", "H007", "H018")
GATE = {"M002": "P002", "M003": "P003", "M004": "P003"}
BOOTSTRAPS = 20_000
BOOTSTRAP_SEED = 20260720


class Study4AnalysisError(RuntimeError):
    pass


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


def exact_mcnemar(left: Sequence[int], right: Sequence[int]) -> tuple[int, int, float]:
    n10 = sum(a == 1 and b == 0 for a, b in zip(left, right))
    n01 = sum(a == 0 and b == 1 for a, b in zip(left, right))
    n = n10 + n01
    if not n:
        return n10, n01, 1.0
    tail = sum(math.comb(n, k) for k in range(0, min(n10, n01) + 1)) / (2**n)
    return n10, n01, min(1.0, 2 * tail)


def holm_adjust(values: Sequence[float]) -> list[float]:
    order = sorted(range(len(values)), key=lambda index: values[index])
    result = [1.0] * len(values)
    running = 0.0
    for rank, index in enumerate(order):
        adjusted = min(1.0, (len(values) - rank) * values[index])
        running = max(running, adjusted)
        result[index] = running
    return result


def cluster_sign_flip(task_effects: Sequence[float]) -> float:
    observed = abs(statistics.fmean(task_effects))
    extreme = 0
    total = 0
    tolerance = 1e-15
    for signs in itertools.product((-1.0, 1.0), repeat=len(task_effects)):
        value = abs(statistics.fmean(sign * effect for sign, effect in zip(signs, task_effects)))
        extreme += value + tolerance >= observed
        total += 1
    return extreme / total


def cluster_bootstrap(task_effects: Sequence[float]) -> tuple[float, float]:
    generator = random.Random(BOOTSTRAP_SEED)
    n = len(task_effects)
    values = sorted(
        statistics.fmean(task_effects[generator.randrange(n)] for _ in range(n))
        for _ in range(BOOTSTRAPS)
    )
    return values[int(0.025 * BOOTSTRAPS)], values[int(0.975 * BOOTSTRAPS) - 1]


def _mechanism(trajectory: Path, gold_files: set[str]) -> dict[str, Any]:
    searched: list[str] = []
    read: list[str] = []
    accepted_edit_seen = False
    for line in trajectory.read_text(encoding="utf-8").splitlines():
        event = json.loads(line)
        kind = event.get("event_type")
        payload = event.get("payload", {})
        if kind == "edit" and payload.get("accepted"):
            accepted_edit_seen = True
        if accepted_edit_seen:
            continue
        if kind == "retrieval_candidate":
            path = str(payload.get("path", ""))
            if path:
                searched.append(path)
        elif kind == "file_read":
            path = str(payload.get("path", ""))
            if path:
                read.append(path)
    return {
        "gold_retrieved_before_edit": bool(set(searched) & gold_files),
        "gold_read_before_edit": bool(set(read) & gold_files),
        "unique_search_paths_before_edit": len(set(searched)),
        "unique_read_paths_before_edit": len(set(read)),
    }


def discover(root: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    split = load_task_split(root / "tasks" / "splits" / "study4_fresh.txt")
    tasks = load_tasks(root)
    repositories = load_repositories(root)
    url_to_repository = {item.repository_url: item.repository_id for item in repositories.values()}
    expected = {
        (task_id, model, harness, GATE[model])
        for task_id in split for model in MODELS for harness in HARNESSES
    }
    paths = sorted((root / "results" / "raw" / "E10").rglob("final_metrics.json"))
    if len(paths) != 180:
        raise Study4AnalysisError(f"E10 requires 180 final metrics; observed {len(paths)}")
    seen: set[tuple[str, str, str, str]] = set()
    revisions: Counter[str] = Counter()
    raw_hasher = sha256()
    rows: list[dict[str, Any]] = []
    for final_path in paths:
        directory = final_path.parent
        required = [
            directory / "run_manifest.json", directory / "trajectory.jsonl",
            directory / "messages.json", directory / "model.patch", directory / "validation.json",
        ]
        if not all(path.is_file() for path in required):
            raise Study4AnalysisError(f"incomplete run directory: {directory}")
        manifest = json.loads(required[0].read_text(encoding="utf-8"))
        final = json.loads(final_path.read_text(encoding="utf-8"))
        identity = manifest["identity"]
        treatment = str(final["harness_id"])
        if "__" not in treatment:
            raise Study4AnalysisError(f"non-composite E10 treatment: {treatment}")
        harness, interface = treatment.split("__", 1)
        key = (identity["task_id"], identity["model_id"], harness, interface)
        if key not in expected or key in seen:
            raise Study4AnalysisError(f"unexpected or duplicate E10 identity: {key}")
        if final["retrieval_harness_id"] != harness or final["edit_interface_id"] != interface:
            raise Study4AnalysisError(f"manifest/final treatment mismatch: {directory}")
        if identity["context_budget"] != 65536 or identity["seed"] != 0:
            raise Study4AnalysisError(f"E10 inference identity drift: {directory}")
        if any(len(item.get("after_instances", [])) != 1 for item in final["residency_transitions"]):
            raise Study4AnalysisError(f"non-exclusive model residency: {directory}")
        seen.add(key)
        revisions[identity["code_revision"]] += 1
        raw_hasher.update(required[0].read_bytes())
        raw_hasher.update(final_path.read_bytes())
        task = tasks[identity["task_id"]]
        mechanism = _mechanism(required[1], set(task.gold_files))
        rows.append(
            {
                "run_id": manifest["run_id"],
                "task_id": identity["task_id"],
                "repository_id": url_to_repository[task.repository_url],
                "language": task.language,
                "model_id": identity["model_id"],
                "harness_id": harness,
                "interface_id": interface,
                "resolved_at_1": int(bool(final["resolved_at_1"])),
                "accepted_edit_cell": int(bool(final["accepted_edit_cell"])),
                "applicable_final_patch": int(bool(final["applicable_final_patch"])),
                "exact_modified_file_match": int(bool(final["exact_modified_file_match"])),
                "fail_to_pass": int(bool(final["fail_to_pass"])),
                "search_gold_any": int(final["search_localization_metrics"]["file_recall_at_10"] > 0),
                "read_gold_any": int(final["read_localization_metrics"]["file_recall_at_10"] > 0),
                **{key: int(value) if isinstance(value, bool) else value for key, value in mechanism.items()},
                "edit_attempts": int(final["edit_attempts"]),
                "edit_acceptances": int(final["edit_acceptances"]),
                "model_calls": int(final["model_calls"]),
                "tool_calls": int(final["tool_calls"]),
                "test_runs": int(final["test_runs"]),
                "total_tokens": int(final["usage"]["total_tokens"]),
                "elapsed_seconds": float(final["elapsed_seconds"]),
                "model_switch_count": int(final["model_switch_count"]),
                "model_switch_seconds": float(final["model_switch_seconds"]),
                "protocol_violation_count": len(final["protocol_violations"]),
                "failure_stage": final["failure_stage"],
                "finished_reason": final["finished_reason"],
            }
        )
    if seen != expected or len(revisions) != 1:
        raise Study4AnalysisError(f"E10 grid/revision mismatch: cells={len(seen)}, revisions={revisions}")
    preflight_path = root / "results" / "reports" / "study4_preflight.json"
    if not preflight_path.is_file():
        raise Study4AnalysisError("Study 4 preflight report is missing")
    preflight = json.loads(preflight_path.read_text(encoding="utf-8"))
    execution_revision = next(iter(revisions))
    if not preflight.get("passed") or preflight.get("research_code_revision") != execution_revision:
        raise Study4AnalysisError("preflight did not pass on the E10 execution revision")
    return rows, {
        "execution_revision": execution_revision,
        "input_cells": len(rows),
        "raw_manifest_and_metrics_sha256": raw_hasher.hexdigest(),
        "repository_counts": dict(Counter(row["repository_id"] for row in rows)),
    }


def primary(rows: Sequence[dict[str, Any]]) -> dict[str, Any]:
    by = {(row["task_id"], row["model_id"], row["harness_id"]): row for row in rows}
    tasks = sorted({row["task_id"] for row in rows})
    task_effects = [
        statistics.fmean(
            by[(task, model, "H007")]["resolved_at_1"]
            - by[(task, model, "H000")]["resolved_at_1"]
            for model in MODELS
        )
        for task in tasks
    ]
    low, high = cluster_bootstrap(task_effects)
    return {
        "contrast": "H007_vs_H000_resolved_at_1",
        "independent_clusters": len(tasks),
        "model_task_pairs": len(tasks) * len(MODELS),
        "h007_count": sum(by[(task, model, "H007")]["resolved_at_1"] for task in tasks for model in MODELS),
        "h000_count": sum(by[(task, model, "H000")]["resolved_at_1"] for task in tasks for model in MODELS),
        "paired_risk_difference": statistics.fmean(task_effects),
        "cluster_bootstrap_ci_low": low,
        "cluster_bootstrap_ci_high": high,
        "exact_task_cluster_sign_flip_p": cluster_sign_flip(task_effects),
        "task_effects": task_effects,
    }


def _paired_model(rows: Sequence[dict[str, Any]], model: str, endpoint: str) -> dict[str, Any]:
    selected = [row for row in rows if row["model_id"] == model]
    by = {(row["task_id"], row["harness_id"]): row for row in selected}
    tasks = sorted({row["task_id"] for row in selected})
    left = [by[(task, "H007")][endpoint] for task in tasks]
    right = [by[(task, "H000")][endpoint] for task in tasks]
    n10, n01, p = exact_mcnemar(left, right)
    differences = [a - b for a, b in zip(left, right)]
    generator = random.Random(BOOTSTRAP_SEED)
    samples = sorted(statistics.fmean(differences[generator.randrange(len(tasks))] for _ in tasks) for _ in range(BOOTSTRAPS))
    return {
        "contrast": f"{model}_H007_vs_H000_{endpoint}", "model_id": model,
        "endpoint": endpoint, "tasks": len(tasks), "h007_count": sum(left), "h000_count": sum(right),
        "paired_risk_difference": statistics.fmean(differences),
        "ci_low": samples[int(.025 * BOOTSTRAPS)], "ci_high": samples[int(.975 * BOOTSTRAPS)-1],
        "discordant_h007_only": n10, "discordant_h000_only": n01, "mcnemar_p": p,
    }


def secondary(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    endpoints = (
        "resolved_at_1", "gold_retrieved_before_edit", "gold_read_before_edit",
        "accepted_edit_cell", "applicable_final_patch", "exact_modified_file_match", "fail_to_pass",
    )
    values = [_paired_model(rows, model, endpoint) for endpoint in endpoints for model in MODELS]
    adjusted = holm_adjust([row["mcnemar_p"] for row in values])
    for row, p in zip(values, adjusted):
        row["mcnemar_p_holm"] = p
    return values


def summaries(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    result = []
    for model in MODELS:
        for harness in HARNESSES:
            group = [row for row in rows if row["model_id"] == model and row["harness_id"] == harness]
            result.append(
                {
                    "model_id": model, "harness_id": harness, "n": len(group),
                    "resolved_count": sum(row["resolved_at_1"] for row in group),
                    "resolved_rate": statistics.fmean(row["resolved_at_1"] for row in group),
                    "gold_retrieved_before_edit_rate": statistics.fmean(row["gold_retrieved_before_edit"] for row in group),
                    "gold_read_before_edit_rate": statistics.fmean(row["gold_read_before_edit"] for row in group),
                    "accepted_edit_rate": statistics.fmean(row["accepted_edit_cell"] for row in group),
                    "applicable_rate": statistics.fmean(row["applicable_final_patch"] for row in group),
                    "mean_total_tokens": statistics.fmean(row["total_tokens"] for row in group),
                    "mean_elapsed_seconds": statistics.fmean(row["elapsed_seconds"] for row in group),
                    "mean_tool_calls": statistics.fmean(row["tool_calls"] for row in group),
                }
            )
    return result


def hierarchical(rows: Sequence[dict[str, Any]]) -> dict[str, Any]:
    try:
        from statsmodels.genmod.bayes_mixed_glm import BinomialBayesMixedGLM
        frame = pd.DataFrame(rows)
        formula = (
            "resolved_at_1 ~ C(harness_id, Treatment(reference='H000')) * "
            "C(model_id, Treatment(reference='M002')) + C(repository_id, Treatment(reference='R001'))"
        )
        fitted = BinomialBayesMixedGLM.from_formula(formula, {"task": "0 + C(task_id)"}, frame).fit_vb()
        names = list(fitted.model.exog_names)
        means = np.asarray(fitted.params[:len(names)], dtype=float)
        sd = np.asarray(fitted.fe_sd, dtype=float)
        coefficients = [
            {"term": name, "log_odds_mean": float(mean), "log_odds_sd": float(error),
             "odds_ratio": float(math.exp(mean)), "or_low": float(math.exp(mean - 1.96*error)),
             "or_high": float(math.exp(mean + 1.96*error))}
            for name, mean, error in zip(names, means, sd)
        ]
        return {"status": "converged", "formula": formula, "task_random_intercept": True, "coefficients": coefficients}
    except Exception as exc:
        return {"status": "failed", "error": repr(exc), "coefficients": []}


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


def analyze(root: Path) -> dict[str, Any]:
    rows, audit = discover(root)
    primary_result = primary(rows)
    secondary_results = secondary(rows)
    summary_rows = summaries(rows)
    hierarchy = hierarchical(rows)
    failures = {
        "failure_stages": dict(sorted(Counter(row["failure_stage"] for row in rows).items())),
        "finished_reasons": dict(sorted(Counter(row["finished_reason"] for row in rows).items())),
        "protocol_violation_cells": sum(row["protocol_violation_count"] > 0 for row in rows),
    }
    output = root / "results" / "derived" / "study4"; output.mkdir(parents=True, exist_ok=True)
    files = []
    for name, values in (
        ("e10_cells.csv", rows), ("e10_model_harness_summary.csv", summary_rows),
        ("e10_secondary_contrasts.csv", secondary_results),
        ("e10_hierarchical_coefficients.csv", hierarchy["coefficients"]),
    ):
        path = output / name; write_csv(path, values); files.append(path)
    manifest = {
        "schema_version": 1, "experiment_id": "E10", "execution_revision": audit["execution_revision"],
        "analysis_code_revision": subprocess.run(["git", "rev-parse", "HEAD"], cwd=root, check=True, capture_output=True, text=True).stdout.strip(),
        "analysis_script_sha256": sha256_file(root / "scripts" / "analyze_study4.py"),
        "input_cells": 180, "raw_manifest_and_metrics_sha256": audit["raw_manifest_and_metrics_sha256"],
        "bootstrap_samples": BOOTSTRAPS, "bootstrap_seed": BOOTSTRAP_SEED,
        "primary_test": "H007 versus H000 resolved_at_1; exact task-cluster sign flip over 20 tasks",
        "secondary_multiplicity": f"Holm across {len(secondary_results)} prespecified within-model contrasts",
    }
    report = {
        "schema_version": 1, "experiment_id": "E10", "audit": audit,
        "primary": primary_result, "model_harness_summaries": summary_rows,
        "secondary_contrasts": secondary_results, "hierarchical_model": hierarchy,
        "failure_analysis": failures, "analysis_manifest": manifest,
        "claim_boundary": {
            "confirmatory": "pooled H007 versus H000 resolution with task-cluster inference",
            "mediation": "stage decomposition only; no identified natural indirect effect",
            "repository_language": "descriptive; partially confounded and unbalanced",
        },
    }
    manifest_path = output / "analysis_manifest.json"; manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    report_path = output / "e10_analysis.json"; report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    files.extend([manifest_path, report_path])
    checksums = {path.name: sha256_file(path) for path in files}
    (output / "SHA256SUMS.json").write_text(json.dumps(checksums, indent=2, sort_keys=True) + "\n")
    return {**report, "output_directory": str(output), "checksums": checksums}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    args = parser.parse_args()
    try:
        result = analyze(args.root.resolve())
    except Exception as exc:
        print(f"STUDY 4 ANALYSIS FAILED: {exc}")
        return 1
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())