scripts/freeze_study5_e16.py

#!/usr/bin/env python3
"""Select Study 5 archetypes and freeze the held-out E16 cell manifest.

The script is intentionally outcome-aware for E13/E15 screening and outcome-blind
to every E16 task.  It refuses to run after any E16 raw artifact exists.
"""

from __future__ import annotations

from collections import defaultdict
from hashlib import sha256
import json
from pathlib import Path
from statistics import mean
from typing import Any, Iterable

from agent_harness.specs import (
    load_edit_interfaces,
    load_experiments,
    load_harnesses,
    load_models,
    load_task_split,
    load_tasks,
)


ROOT = Path(__file__).resolve().parents[1]
SCREENING_REVISION = "7f4de67853deab34aca5a9ceaaf7e9f85b901088"
SCREENING_EXPERIMENTS = {"E13": 1440, "E15": 540}
ROBUSTNESS_CANDIDATES = ("H008", "H010")
EXPECTED_SELECTION = 6
EXPECTED_TASKS = 17


def canonical_hash(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"))
    return sha256(payload.encode("utf-8")).hexdigest()


def _raw_rows(root: Path, experiment_id: str) -> tuple[list[dict[str, Any]], str]:
    paths = sorted((root / "results" / "raw" / experiment_id).glob("*/*/*/final_metrics.json"))
    expected = SCREENING_EXPERIMENTS[experiment_id]
    if len(paths) != expected:
        raise RuntimeError(f"{experiment_id} requires {expected} finalized cells, found {len(paths)}")
    rows = [json.loads(path.read_text(encoding="utf-8")) for path in paths]
    identities = {
        (row["task_id"], row["retrieval_harness_id"], row["edit_interface_id"], row["model_id"])
        for row in rows
    }
    if len(identities) != expected:
        raise RuntimeError(f"{experiment_id} contains duplicate cell identities")
    digest = sha256()
    for path in paths:
        digest.update(str(path.relative_to(root)).encode("utf-8"))
        digest.update(b"\0")
        digest.update(path.read_bytes())
        digest.update(b"\0")
    return rows, digest.hexdigest()


def _screening_report(root: Path, experiment_id: str, rows: list[dict[str, Any]]) -> dict[str, Any]:
    candidates = sorted(
        (root / "results" / "reports").glob(
            f"{experiment_id}_{SCREENING_REVISION[:12]}_*.json"
        )
    )
    expected = SCREENING_EXPERIMENTS[experiment_id]
    paths = [
        path
        for path in candidates
        if json.loads(path.read_text(encoding="utf-8")).get("run_count") == expected
    ]
    if len(paths) != 1:
        raise RuntimeError(
            f"{experiment_id} requires exactly one complete frozen final report, found {len(paths)}"
        )
    path = paths[0]
    report = json.loads(path.read_text(encoding="utf-8"))
    observed = {
        "run_count": len(rows),
        "accepted_edit_count": sum(bool(row["accepted_edit_cell"]) for row in rows),
        "applicable_patch_count": sum(bool(row["applicable_final_patch"]) for row in rows),
        "resolved_count": sum(bool(row["resolved_at_1"]) for row in rows),
    }
    if report.get("code_revision") != SCREENING_REVISION or report.get("run_count") != expected:
        raise RuntimeError(f"{experiment_id} final report revision/count mismatch")
    if any(report.get(key) != value for key, value in observed.items()):
        raise RuntimeError(f"{experiment_id} raw ledger disagrees with final report")
    return {
        "path": str(path.relative_to(root)),
        "sha256": sha256(path.read_bytes()).hexdigest(),
        **observed,
    }


def _metrics(rows: Iterable[dict[str, Any]]) -> dict[str, float | int]:
    values = list(rows)
    count = len(values)
    return {
        "cells": count,
        "resolved_rate": sum(bool(row["resolved_at_1"]) for row in values) / count,
        "accepted_edit_rate": sum(bool(row["accepted_edit_cell"]) for row in values) / count,
        "applicable_patch_rate": sum(bool(row["applicable_final_patch"]) for row in values) / count,
        "mean_total_tokens": mean(float(row["usage"]["total_tokens"]) for row in values),
        "mean_wall_seconds": mean(float(row["elapsed_seconds"]) for row in values),
    }


def _rank_key(item: tuple[str, dict[str, float | int]]) -> tuple[float, float, float, float, float, str]:
    harness_id, metric = item
    return (
        -float(metric["resolved_rate"]),
        -float(metric["applicable_patch_rate"]),
        -float(metric["accepted_edit_rate"]),
        float(metric["mean_total_tokens"]),
        float(metric["mean_wall_seconds"]),
        harness_id,
    )


def _accepted_key(item: tuple[str, dict[str, float | int]]) -> tuple[float, float, float, float, str]:
    harness_id, metric = item
    return (
        -float(metric["accepted_edit_rate"]),
        -float(metric["applicable_patch_rate"]),
        float(metric["mean_total_tokens"]),
        float(metric["mean_wall_seconds"]),
        harness_id,
    )


def _dominates(left: dict[str, float | int], right: dict[str, float | int]) -> bool:
    left_values = (
        float(left["resolved_rate"]),
        float(left["applicable_patch_rate"]),
        float(left["accepted_edit_rate"]),
        -float(left["mean_total_tokens"]),
        -float(left["mean_wall_seconds"]),
    )
    right_values = (
        float(right["resolved_rate"]),
        float(right["applicable_patch_rate"]),
        float(right["accepted_edit_rate"]),
        -float(right["mean_total_tokens"]),
        -float(right["mean_wall_seconds"]),
    )
    return all(a >= b for a, b in zip(left_values, right_values)) and any(
        a > b for a, b in zip(left_values, right_values)
    )


def select_archetypes(metrics: dict[str, dict[str, float | int]]) -> tuple[list[str], list[dict[str, str]], list[str]]:
    selected: list[str] = []
    decisions: list[dict[str, str]] = []

    def retain(role: str, harness_id: str) -> None:
        if harness_id in selected:
            raise RuntimeError(f"selection role {role} repeated {harness_id}")
        selected.append(harness_id)
        decisions.append({"role": role, "harness_id": harness_id})

    retain("mandatory_exact_baseline", "H000")
    unselected = lambda: [(key, value) for key, value in metrics.items() if key not in selected]
    retain("highest_resolution", min(unselected(), key=_rank_key)[0])

    best_resolution = max(float(value["resolved_rate"]) for value in metrics.values())
    near_best = [
        item for item in unselected() if float(item[1]["resolved_rate"]) >= best_resolution - 0.05
    ]
    retain(
        "lowest_tokens_within_5pp_resolution",
        min(
            near_best,
            key=lambda item: (
                float(item[1]["mean_total_tokens"]),
                *_rank_key(item),
            ),
        )[0],
    )
    retain("highest_accepted_edit", min(unselected(), key=_accepted_key)[0])

    robustness = [(key, metrics[key]) for key in ROBUSTNESS_CANDIDATES if key not in selected]
    retain("best_prespecified_robustness_candidate", min(robustness, key=_rank_key)[0])

    frontier = sorted(
        harness_id
        for harness_id, metric in metrics.items()
        if not any(
            other_id != harness_id and _dominates(other, metric)
            for other_id, other in metrics.items()
        )
    )
    for harness_id in frontier:
        if len(selected) == EXPECTED_SELECTION:
            break
        if harness_id not in selected:
            retain("pareto_frontier_fill", harness_id)
    if len(selected) != EXPECTED_SELECTION:
        raise RuntimeError(f"selection produced {len(selected)} harnesses, expected {EXPECTED_SELECTION}")
    return sorted(selected), decisions, frontier


def build_manifest(root: Path, selected: list[str]) -> dict[str, Any]:
    experiment = load_experiments(root)["E16"]
    if sorted(experiment.harness_ids) != selected:
        raise RuntimeError(
            f"E16 specification harnesses {sorted(experiment.harness_ids)} do not match selection {selected}"
        )
    tasks = load_tasks(root)
    harnesses = load_harnesses(root)
    interfaces = load_edit_interfaces(root)
    models = load_models(root)
    split = load_task_split(root / "tasks" / "splits" / "study5_fresh.txt")
    if len(split) != EXPECTED_TASKS:
        raise RuntimeError(f"E16 requires {EXPECTED_TASKS} fresh tasks, found {len(split)}")
    gate = json.loads((root / "configs" / "gates" / "E09_model_interface_gate.json").read_text())["selected"]
    cells: list[dict[str, Any]] = []
    for task_index, task_id in enumerate(split):
        task = tasks[task_id]
        if task.validation_status != "end_to_end_ready":
            raise RuntimeError(f"{task_id} is not end-to-end ready")
        for model_index, model_id in enumerate(experiment.model_ids):
            treatments = list(selected)
            offset = (task_index + model_index) % len(treatments)
            treatments = treatments[offset:] + treatments[:offset]
            interface_id = str(gate[model_id])
            for treatment_index, harness_id in enumerate(treatments):
                cells.append(
                    {
                        "order": len(cells),
                        "task_id": task_id,
                        "repository_sha": task.base_commit,
                        "harness_id": harness_id,
                        "harness_hash": harnesses[harness_id].config_hash,
                        "interface_id": interface_id,
                        "interface_hash": interfaces[interface_id].config_hash,
                        "model_id": model_id,
                        "model_hash": models[model_id].config_hash,
                        "context_budget": experiment.context_budgets[0],
                        "seed": experiment.seeds[0],
                        "within_model_order": treatment_index,
                    }
                )
    manifest = {
        "schema_version": 1,
        "study": "Study 5 end-to-end harness behavior",
        "experiment_id": "E16",
        "outcome_blind": True,
        "planned_cells": len(cells),
        "cells": cells,
    }
    manifest["design_sha256"] = canonical_hash(manifest)
    return manifest


def main() -> None:
    root = ROOT.resolve()
    raw_e16 = root / "results" / "raw" / "E16"
    if raw_e16.exists() and any(raw_e16.rglob("*")):
        raise RuntimeError("E16 outcomes already exist; selection and manifest cannot be re-frozen")
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    source_digests: dict[str, str] = {}
    source_reports: dict[str, dict[str, Any]] = {}
    # E14 is excluded because each harness is repeated under three action interfaces;
    # E13 and E15 provide one gate-selected action per model-task-harness cell.
    primary_rows: dict[str, list[dict[str, Any]]] = {}
    for experiment_id in SCREENING_EXPERIMENTS:
        rows, source_digests[experiment_id] = _raw_rows(root, experiment_id)
        primary_rows[experiment_id] = rows
        source_reports[experiment_id] = _screening_report(root, experiment_id, rows)
    # H007 is the zero-hop anchor in E15 but its primary screening estimate is
    # the larger E13 component panel.  Each harness therefore contributes one
    # and only one prespecified panel to cross-harness archetype selection.
    for row in primary_rows["E13"]:
        grouped[str(row["retrieval_harness_id"])].append(row)
    e13_harnesses = set(grouped)
    for row in primary_rows["E15"]:
        if str(row["retrieval_harness_id"]) not in e13_harnesses:
            grouped[str(row["retrieval_harness_id"])].append(row)
    metrics = {harness_id: _metrics(rows) for harness_id, rows in sorted(grouped.items())}
    selected, decisions, frontier = select_archetypes(metrics)
    manifest = build_manifest(root, selected)
    if manifest["planned_cells"] != EXPECTED_TASKS * EXPECTED_SELECTION * 3:
        raise RuntimeError("E16 manifest cell count mismatch")
    selection = {
        "schema_version": 1,
        "study": "Study 5 end-to-end harness behavior",
        "selection_is_outcome_aware_to_screening": True,
        "selection_is_outcome_blind_to_e16": True,
        "screening_revision": SCREENING_REVISION,
        "screening_experiments": sorted(SCREENING_EXPERIMENTS),
        "excluded_screening_experiment": {
            "experiment_id": "E14",
            "reason": "retrieval harnesses are repeated across action interfaces and are reserved for retrieval-by-action inference",
        },
        "source_raw_sha256": source_digests,
        "source_final_reports": source_reports,
        "primary_screening_panel": {
            **{f"H{index:03d}": "E13" for index in range(8)},
            **{f"H{index:03d}": "E15" for index in range(8, 16)},
        },
        "metrics": metrics,
        "robustness_candidates": list(ROBUSTNESS_CANDIDATES),
        "pareto_frontier": frontier,
        "decisions": decisions,
        "selected_harnesses": selected,
        "e16_manifest_sha256": manifest["design_sha256"],
        "e16_planned_cells": manifest["planned_cells"],
    }
    selection["selection_sha256"] = canonical_hash(selection)
    output_dir = root / "configs" / "study5"
    (output_dir / "E16_selection.json").write_text(
        json.dumps(selection, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    (output_dir / "E16_cells.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(selection, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()