#!/usr/bin/env python3
"""Reproduce the addendum's cost-denominator retest.

    python retest_cost_denominator.py

Reads test.parquet and instances.parquet from the working directory and prints
the table in ADDENDUM_2026-07-26_two-surfaces.md §3.

WHAT IT TESTS
-------------
extract_instance_results.py normalises per-instance costs so the mean over
*graded* records (resolved is not None AND cost is not None) equals the published
cost_per_instance. The survey's primary cost policy averaged over a different
population: cost is not None, regardless of resolved.

If the divergent cost cells were an artifact of choosing the wrong denominator,
switching to the script's own population should rescue some of them. It does not.

SCOPE — read this before quoting the numbers
--------------------------------------------
This holds the policy FIXED at one policy per population, in order to isolate the
denominator. It is NOT the survey's census rule. The census counts a cell
reconciled if ANY declared policy reproduces it, which for these same 170 cells
gives 142. The comparison below is 115 vs 114 because both sides run a single
policy. Neither number restates the census; they answer a narrower question.

Tolerance tiers are the arm's, declared in its preregistration, in published
units: T0 = equal at 4 decimal places, T1 = |delta| <= 0.005, T2 = <= 0.05.
"""
import pyarrow.parquet as pq
from collections import defaultdict

CT1, CT2 = 0.005, 0.05
BENCH_TO_CATEGORY = {
    "swe-bench": "issue_resolution",
    "swe-bench-multimodal": "frontend",
    "commit0": "greenfield",
    "swt-bench": "testing",
    "gaia": "information_gathering",
}


def tier(recomputed, published):
    if recomputed is None or published is None:
        return "divergent"
    if abs(round(recomputed, 4) - round(published, 4)) < 1e-7:
        return "T0"
    d = abs(recomputed - published)
    if d <= CT1:
        return "T1"
    if d <= CT2:
        return "T2"
    return "divergent"


def main():
    test = pq.read_table("test.parquet").to_pydict()
    inst = pq.read_table("instances.parquet").to_pydict()

    groups = defaultdict(list)
    for i in range(len(inst["id"])):
        groups[(inst["id"][i], BENCH_TO_CATEGORY[inst["benchmark"][i]])].append(i)

    counts = {"survey": defaultdict(int), "graded": defaultdict(int)}
    moved = []

    for slug in BENCH_TO_CATEGORY.values():
        published = test[f"{slug}_cost"]
        for row, model in enumerate(test["id"]):
            idx = groups.get((model, slug), [])
            if not idx:
                continue
            # the survey's primary cost policy: mean over instances with a cost
            survey_pop = [inst["cost"][i] for i in idx if inst["cost"][i] is not None]
            # the extractor's population: graded records only
            graded_pop = [inst["cost"][i] for i in idx
                          if inst["cost"][i] is not None and inst["resolved"][i] is not None]

            a = tier(sum(survey_pop) / len(survey_pop) if survey_pop else None, published[row])
            b = tier(sum(graded_pop) / len(graded_pop) if graded_pop else None, published[row])
            counts["survey"][a] += 1
            counts["graded"][b] += 1
            if a == "divergent" and b != "divergent":
                moved.append((f"{slug}_cost", model, b))

    print("Per-category cost cells: 170 (5 columns x 34 rows). One policy per population.\n")
    print(f"  {'population':<44} {'recon':>5} {'T0':>4} {'T1':>4} {'T2':>4} {'div':>5}")
    for key, label in (("survey", "survey's primary policy (cost non-null)"),
                       ("graded", "extractor's population (graded records)")):
        c = counts[key]
        rec = c["T0"] + c["T1"] + c["T2"]
        print(f"  {label:<44} {rec:>5} {c['T0']:>4} {c['T1']:>4} {c['T2']:>4} {c['divergent']:>5}")

    print(f"\n  cells moving divergent -> reconciled under the graded population: {len(moved)}")
    for m in moved:
        print("   ", m)
    print("\n  Reminder: the census rule admits ANY declared policy and reconciles 142")
    print("  of these 170. The figures above hold one policy fixed on purpose.")


if __name__ == "__main__":
    main()