"""Strength gap, census and headroom fit, recomputed with the program's corrections applied (frontier report).
Local re-tabulation only: no model, VM or evaluation.

Inputs (read only):
  forensics/ledger.csv in the paper directory, the 195-row comparison record.
  In the released ledger.csv these free-text cells are empty, and every numeric column and every categorical column
  except team is as produced: note on L009 L036 L039 L042 L043 L052 L073 L079 L083 L086 L091 L099 L100 L101 L102
  L104 L105 L107 L108 L109 L110 L112 L115 L121 L122 L125 L129 L130 L138 L139 L140 L148 L150 L165 L166 L168 L176
  L178 L193; settings_parity on L036 L037 L039 L052 L086 L102 L108 L110 L139 L140 L142 L143 L148 L193; instrument
  on L169 L170 L171 L173; contrast on L169 L170 L171 L172 L173. Agent and team names are replaced by neutral IDs
  or plain study names in the team column and in the instrument, contrast, settings_parity and note cells of L030
  L034 L038 L044 L045 L052 L082 L103 L107-L110 L133 L134 L151 L152 L184-L191 L193. Every source cell says in plain
  words where the comparison was recorded (internal document names, paths and commit ids removed), as do the note
  of L068 and the instrument of L146. No number this script prints reads those cells.
  row rules and fit functions copied verbatim from the record's model code (its sections 2, 3 and 8)
Corrections applied (each from the program's audits of these comparisons):
  census: L112, L158, L164 out (audit findings); L156 -> null (TIE); L165 -> null (corrected sensor - bare +7.1
  [0, +16.7]); L193, L194, L052 out (budget-contract voiding plus the no-thinking defect; model succession UP
  infra flip, withdrawn); L014 -> corrected -0.54 [-1.19, -0.05] (90%); L174, L176 relabelled mixed (Opus 4.8 ->
  5.5 stream)
  strength gap: one inclusion rule for both directions (see RULE below); corrected values for L001, L002, L014, L103,
  L184.
Outputs: prints everything; writes recount_output.json next to this file.

Run (from the paper directory): python3 -I evidence/recount.py > evidence/recount_output.txt
"""
import itertools
import json
import math
import os

import numpy as np
import pandas as pd
from scipy import optimize, stats

SRC = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "forensics")
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "recount_output.json")
L = pd.read_csv(os.path.join(SRC, "ledger.csv"), keep_default_na=False)
for c in ["est", "lo", "hi", "p", "n", "se", "z", "ours", "ctrl", "mx", "headroom_pts"]:
    L[c] = pd.to_numeric(L[c], errors="coerce")
LAYER = {"layer_vs_bare", "layer_vs_prior", "layer_vs_rival", "carried_w2s", "carried_s2w", "carried_self",
         "system_vs_single", "trained_vs_frontier"}
L["structural"] = L["flags"].str.contains("structural")
L["superseded"] = L["flags"].str.contains("harness_bug")
RES = {}


def se_from_ci(lo, hi, level=0.95):
    return (hi - lo) / (2 * stats.norm.ppf(0.5 + level / 2))


# ================================================================== (b) census, as ledgered and corrected
def census(df, label):
    oz = df[df.ctype.isin(LAYER) & ~df.structural & ~df.superseded & (df.model == "opus-5.5") & df.z.notna()]
    sig = np.where(oz.z > 1.96, "win", np.where(oz.z < -1.96, "loss", "null"))
    tab = pd.crosstab(oz.ctype, sig)
    for col in ("win", "null", "loss"):
        if col not in tab:
            tab[col] = 0
    tab = tab[["win", "null", "loss"]]
    tot = {k: int((sig == k).sum()) for k in ("win", "null", "loss")}
    print(f"\n[{label}] Opus 5.5 rows with an SE: {len(oz)}; wins {tot['win']}, nulls {tot['null']}, losses {tot['loss']}")
    print(tab.to_string())
    RES[f"census_{label}"] = dict(N=int(len(oz)), **tot, by_ctype={c: {k: int(v) for k, v in tab.loc[c].items()}
                                                                       for c in tab.index})
    return oz


print("=" * 100)
print("(b) CENSUS")
census(L, "as_ledgered")
C = L.copy()
OUT_ROWS = {"L112": "mislabelled control (audit finding)", "L158": "threshold choice (audit finding)",
            "L164": "impaired Opus-built ACE playbook (audit finding)",
            "L193": "budget-contract voiding + no-thinking swapped turns (audit finding)",
            "L194": "budget-contract voiding (audit finding)",
            "L052": "model succession UP infra flip, withdrawn (program record)"}
C = C[~C.id.isin(OUT_ROWS)].copy()
# L156: preregistered TIE (fails the 1% per-state side condition) -> null
C.loc[C.id == "L156", "z"] = 0.0
# L165: the +9.5 was sensor - layer; corrected sensor - bare +7.1 [0, +16.7] (the program's record of that comparison)
C.loc[C.id == "L165", ["est", "lo", "hi"]] = [7.1, 0.0, 16.7]
C.loc[C.id == "L165", "z"] = 7.1 / se_from_ci(0.0, 16.7)
# L014: corrected -0.54 [-1.19, -0.05], 90% interval (audit correction)
C.loc[C.id == "L014", ["est", "lo", "hi"]] = [-0.54, -1.19, -0.05]
C.loc[C.id == "L014", "z"] = -0.54 / se_from_ci(-1.19, -0.05, 0.90)
# L184: reference-time normalised -0.223 [-0.494, -0.057] (audit correction)
C.loc[C.id == "L184", ["est", "lo", "hi"]] = [-0.223, -0.494, -0.057]
C.loc[C.id == "L184", "z"] = -0.223 / se_from_ci(-0.494, -0.057)
# L174, L176: AppWorld durable stream, Opus 4.8 -> 5.5 between tasks, gain on the 4.8 half
C.loc[C.id.isin(["L174", "L176"]), "model"] = "mixed"
census(C, "corrected")
print("removed:", OUT_ROWS)
print("changed: L156 -> null (TIE); L165 -> corrected +7.1 [0, +16.7] (z %.2f); L014 -> corrected -0.54 (z %.2f); "
      "L184 -> corrected -0.223 (z %.2f); L174, L176 -> mixed"
      % (float(C.loc[C.id == "L165", "z"].iloc[0]), float(C.loc[C.id == "L014", "z"].iloc[0]),
         float(C.loc[C.id == "L184", "z"].iloc[0])))
RIVAL_ADV = ["L130", "L139", "L140", "L142"]
C2 = C[~C.id.isin(RIVAL_ADV)]
census(C2, "corrected_minus_rival_advantage_losses")

# ================================================================== (c) headroom fit with intervals and row sets


def nll(params, d, kind):  # verbatim from the record's model code
    tau = math.exp(params[-1])
    if kind == "zero":
        mu = np.zeros(len(d))
    elif kind == "bias":
        mu = np.full(len(d), -params[0])
    elif kind == "head":
        mu = params[0] * d.h.values
    else:  # headbias
        mu = params[0] * d.h.values - params[1]
    sd = np.sqrt(d.s.values ** 2 + tau ** 2)
    return -stats.norm.logpdf(d.y.values, mu, sd).sum()


def fit(d, kind):  # verbatim from the record's model code
    x0 = {"zero": [0.0], "bias": [0.0, 0.0], "head": [0.3, 0.0], "headbias": [0.3, 0.0, 0.0]}[kind]
    r = optimize.minimize(nll, x0, args=(d, kind), method="Nelder-Mead",
                          options=dict(maxiter=20000, xatol=1e-7, fatol=1e-9))
    return r.x, r.fun, 2 * len(x0) + 2 * r.fun


def profile(d, kind, idx, grid):
    """Profile-likelihood 95% interval for parameter idx (others, incl. tau, re-optimised)."""
    xb, fb, _ = fit(d, kind)
    prof = []
    for g in grid:
        free0 = [v for j, v in enumerate(xb) if j != idx]

        def f(p, g=g):
            full = list(p[:idx]) + [g] + list(p[idx:])
            return nll(full, d, kind)
        prof.append(optimize.minimize(f, free0, method="Nelder-Mead", options=dict(maxiter=8000)).fun)
    prof = np.array(prof)
    ok = grid[prof - fb <= 1.92]
    return float(xb[idx]), float(ok.min()), float(ok.max())


hp = L[L.ctype.isin(LAYER) & ~L.structural & ~L.superseded & L.se.notna() & L.ctrl.notna() & L.mx.notna()
       & L.headroom_pts.notna()].copy()
hp["y"] = hp.est * 100 / hp.mx
hp["s"] = hp.se * 100 / hp.mx
hp["h"] = hp.headroom_pts
A = hp[~hp.ctype.isin(["layer_vs_rival", "system_vs_single"]) & (hp.model == "opus-5.5")].copy()
print("\n" + "=" * 100)
print("(c) HEADROOM FIT, Opus 5.5 rows with a no-layer control (the record's model code, section 3A)")
print(A[["id", "ctype", "contrast", "y", "s", "h", "flags"]].round(2).to_string(index=False))
CARRIED = ["L010", "L011", "L012", "L016"]
SETS = [
    ("all 15 (as published)", []),
    ("drop carried-from-weaker rows L010 L011 L012 L016", CARRIED),
    ("drop L124 (subset of L126)", ["L124"]),
    ("drop L057 (Haiku->Opus mid-task handoff)", ["L057"]),
    ("drop L099 (unaudited) and L166 (post hoc)", ["L099", "L166"]),
    ("drop L174 (mixed-model stream)", ["L174"]),
    ("drop all of the above (R2 combined)", CARRIED + ["L124", "L057", "L099", "L166", "L174"]),
]
GRID_K = np.linspace(-0.5, 1.5, 801)
GRID_B = np.linspace(-20, 20, 801)
rows = []
for name, drop in SETS:
    d = A[~A.id.isin(drop)]
    fits = {k: fit(d, k) for k in ("zero", "bias", "head", "headbias")}
    k, klo, khi = profile(d, "head", 0, GRID_K)
    kb, kblo, kbhi = profile(d, "headbias", 0, GRID_K)
    b, blo, bhi = profile(d, "headbias", 1, GRID_B)
    const = -fits["bias"][0][0]
    row = dict(rows=name, n=int(len(d)), k=round(k, 3), k_ci=[round(klo, 3), round(khi, 3)],
               k_headbias=round(kb, 3), k_headbias_ci=[round(kblo, 3), round(kbhi, 3)],
               b=round(b, 2), b_ci=[round(blo, 2), round(bhi, 2)], const_lift=round(const, 2),
               aic={kk: round(v[2], 1) for kk, v in fits.items()},
               tau_head=round(math.exp(fits["head"][0][-1]), 2))
    rows.append(row)
    print(f"\n-- {name}: n={row['n']}")
    print(f"   head:      k = {row['k']} [{row['k_ci'][0]}, {row['k_ci'][1]}] (profile 95%), tau {row['tau_head']}")
    print(f"   head+bias: k = {row['k_headbias']} [{row['k_headbias_ci'][0]}, {row['k_headbias_ci'][1]}], "
          f"b = {row['b']} [{row['b_ci'][0]}, {row['b_ci'][1]}]")
    print(f"   constant lift (bias model) = {row['const_lift']} pts")
    print(f"   AIC zero {row['aic']['zero']}, constant {row['aic']['bias']}, head {row['aic']['head']}, "
          f"head+bias {row['aic']['headbias']}")
RES["headroom"] = rows

# ================================================================== (a) strength gap under one inclusion rule
print("\n" + "=" * 100)
print("(a) STRENGTH GAP")
RULE = ("A row counts as a carried-knowledge case only if (i) the reader is given knowledge distilled from a "
        "source model's past episodes (a store, playbook, rules, recipes, minted checks or a compiled digest), not the "
        "in-progress work product of the current task; (ii) it is compared with the same reader without that "
        "knowledge; (iii) the reading is not WITHDRAWN or DO-NOT-USE in the program's audits; corrected values replace "
        "ledger values. Applied identically to both directions.")
print("RULE:", RULE)
CAR = [  # id, reader, source: copied from the record's model code, section 8
    ("L001", "haiku", "opus-5.5"), ("L002", "haiku", "opus-5.5"), ("L017", "haiku", "opus-4.8"),
    ("L015", "haiku", "opus-4.8"), ("L026", "haiku", "opus-4.8"), ("L098", "haiku", "opus-5.5"),
    ("L103", "qwen", "opus-5.5"), ("L104", "qwen", "opus-5.5"), ("L105", "qwen", "opus-5.5"),
    ("L111", "opus-4.8", "opus-5.5"), ("L049", "opus-4.8", "opus-5.5"), ("L054", "opus-4.8", "opus-5.5"),
    ("L051", "haiku", "opus-5.5"),
    ("L018", "opus-4.8", "opus-4.8"), ("L058", "haiku", "haiku"), ("L182", "opus-4.8", "opus-4.8"),
    ("L191", "opus-4.8", "opus-4.8"), ("L084", "opus-5.5", "opus-5.5"),
    ("L014", "opus-5.5", "opus-4.8"), ("L016", "opus-5.5", "opus-4.8"), ("L183", "opus-5.5", "opus-4.8"),
    ("L184", "opus-5.5", "opus-4.8"), ("L185", "opus-5.5", "opus-4.8"), ("L186", "opus-5.5", "opus-4.8"),
    ("L187", "opus-5.5", "opus-4.8"), ("L188", "opus-5.5", "opus-4.8"), ("L059", "opus-5.5", "opus-4.8"),
    ("L028", "opus-5.5", "opus-4.8"), ("L057", "opus-5.5", "haiku"), ("L056", "opus-5.5", "haiku"),
]
EXCLUDE = {
    "L049": "(i) mid-task handoff of the current task (DOWN, ours - Opus 4.8 alone)",
    "L054": "(i) mid-task handoff of the current task (a5 - Opus 4.8 alone)",
    "L051": "(i) mid-task handoff of the current task (Opus -> Haiku)",
    "L056": "(i) mid-task handoff of the current task (Haiku -> Opus 5.5, vs transcript)",
    "L057": "(i) mid-task handoff of the current task (Haiku -> Opus 5.5, vs fresh)",
    "L059": "(i) kernel succession chain: the successor continues the predecessor's work product",
    "L104": "(i) guard: hard checks from the ticket's own examples, no distilled knowledge",
    "L105": "(i) commit: commits the deployed tree after every ticket, no distilled knowledge",
    "L098": "(iii) wave-2 lineage layer over tickets 41-50: DO-NOT-USE (audit of the wave-1/2 handover readings)",
}
# corrected values (value, CI lo, CI hi, CI level) replace ledger values
CORR = {
    "L001": (7.0, 2.88, 11.31, 0.95, "all repeats TC 74.74 vs 67.75 (AppWorld inheritance audit)"),
    "L002": (5.16, -0.79, 11.31, 0.95, "all repeats TN 86.11 vs 80.95 (AppWorld inheritance audit)"),
    "L103": (0.476, 0.149, 0.803, None, "12:12Z t-interval, n=6, sign-flip p=0.031 (audit recomputation)"),
    "L014": (-0.54, -1.19, -0.05, 0.90, "offline H12 vs bare k=10 (audit correction)"),
    "L184": (-0.223, -0.494, -0.057, 0.95, "reference-time normalised (audit correction)"),
}
SOURCE = {  # independent source = one store or capital artefact tested by one team on one instrument family
    "L001": "AppWorld: Opus 5.5-compiled AppWorld store read by Haiku",
    "L002": "AppWorld: Opus 5.5-compiled AppWorld store read by Haiku",
    "L017": "Long-running harness: Opus 4.8-built living store read by Haiku",
    "L015": "Long-running harness: Opus 4.8-built living store read by Haiku",
    "L026": "Long-running harness: Opus 4.8-built living store read by Haiku",
    "L103": "Long-horizon: full_v2 capital minted by Opus 5.5",
    "L111": "Long-horizon: full_v2 capital minted by Opus 5.5",
    "L014": "Long-running harness (horizon run): Opus 4.8 kernel store, offline",
    "L183": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L184": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L185": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L186": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L187": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L188": "Training: Opus 4.8 kernel store, served vs not (O-K arms and item-exposure run)",
    "L016": "Long-running harness: Opus 4.8-built living store, AppWorld fresh tasks",
    "L028": "Long-running harness: Opus 4.8-built living store, AppWorld fresh tasks",
}
S = {"haiku": 1, "qwen": 1, "opus-4.8": 2, "opus-5.5": 3}
IDX = L.set_index("id")
inc, exc = [], []
for i, rd, src in CAR:
    if i in EXCLUDE:
        exc.append(dict(id=i, reason=EXCLUDE[i]))
        continue
    r = IDX.loc[i]
    est, z, note = float(r.est), (None if pd.isna(r.z) else float(r.z)), ""
    if i in CORR:
        est, lo, hi, lvl, note = CORR[i]
        if lvl is None:  # t-interval over 6 streams: use the sign-flip test, p = 0.031
            z = est / ((hi - lo) / (2 * stats.t.ppf(0.975, 5)))
        else:
            z = est / se_from_ci(lo, hi, lvl)
    gap = S[rd] - S[src]
    inc.append(dict(id=i, reader=rd, source_model=src, gap=gap, est=round(est, 4), z=None if z is None else round(z, 2),
                    sign=0 if est == 0 else (1 if est > 0 else -1), sig=(z is not None and abs(z) > 1.96),
                    source=SOURCE.get(i, "self (same model)"), contrast=r.contrast, correction=note))
D = pd.DataFrame(inc)
print("\nincluded rows:")
print(D[["id", "source_model", "reader", "gap", "est", "z", "sig", "source"]].to_string(index=False))
print("\nexcluded rows:")
for e in exc:
    print("  ", e["id"], e["reason"])


def side(df, name):
    nz = df[df.sign != 0]
    out = dict(rows=int(len(df)), nonzero=int(len(nz)), positive=int((nz.sign > 0).sum()),
               negative=int((nz.sign < 0).sum()), zero=int((df.sign == 0).sum()),
               significant_positive=int(((df.sign > 0) & df.sig).sum()),
               significant_negative=int(((df.sign < 0) & df.sig).sum()),
               sources=sorted(df.source.unique().tolist()))
    print(f"\n{name}: {out['positive']} of {out['nonzero']} nonzero helped ({out['zero']} tie); significant: "
          f"{out['significant_positive']} helped, {out['significant_negative']} hurt; independent sources "
          f"{len(out['sources'])}:")
    for s_ in out["sources"]:
        sub = df[df.source == s_]
        print(f"     {s_}: signs {sub.sign.tolist()}, significant {sub.sig.tolist()}")
    return out


down = side(D[D.gap < 0], "downhill (source stronger than reader)")
same = side(D[D.gap == 0], "same model")
up = side(D[D.gap > 0], "uphill (source weaker or earlier than reader)")
rho, p = stats.spearmanr(D.gap, D.sign)
print(f"\nrow-level Spearman(gap, sign) under the rule = {rho:.3f} (nominal p {p:.1e}, n {len(D)}; rows not independent)")

# source-level: one vote per independent source on the two sides (same-model rows are not a direction)
src_rows = []
for s_, sub in D[D.gap != 0].groupby("source"):
    net = np.sign(sub.sign.sum())
    src_rows.append(dict(source=s_, side="down" if sub.gap.iloc[0] < 0 else "up", net_sign=int(net),
                         any_sig=bool(sub.sig.any()), members=sub.id.tolist()))
SR = pd.DataFrame(src_rows)
print("\nsource-level table (net sign = sign of the sum of member signs):")
print(SR.to_string(index=False))
dpos = int(((SR.side == "down") & (SR.net_sign > 0)).sum())
dn = int((SR.side == "down").sum())
upos = int(((SR.side == "up") & (SR.net_sign > 0)).sum())
un = int((SR.side == "up").sum())
fe2 = stats.fisher_exact([[dpos, dn - dpos], [upos, un - upos]], alternative="two-sided")[1]
fe1 = stats.fisher_exact([[dpos, dn - dpos], [upos, un - upos]], alternative="greater")[1]
print(f"\nsources helped: downhill {dpos} of {dn}, uphill {upos} of {un}; Fisher exact p two-sided {fe2:.3f}, "
      f"one-sided {fe1:.3f}")
# exact permutation p for source-level Spearman(gap, net sign), gap = mean member gap
SR["gap"] = [float(D[D.source == s_].gap.mean()) for s_ in SR.source]
r0 = stats.spearmanr(SR.gap, SR.net_sign)[0]
perm = [stats.spearmanr(SR.gap, list(pp))[0] for pp in itertools.permutations(SR.net_sign.tolist())]
pperm = float(np.mean([abs(x) >= abs(r0) - 1e-12 for x in perm]))
print(f"source-level Spearman(gap, net sign) = {r0:.3f}, exact permutation p = {pperm:.3f} (n {len(SR)} sources)")
RES["strength_gap"] = dict(rule=RULE, excluded=exc, included=inc, downhill=down, same_model=same, uphill=up,
                           row_spearman=dict(rho=round(float(rho), 3), nominal_p=float(p), n=int(len(D))),
                           source_level=dict(table=src_rows, downhill_helped=f"{dpos} of {dn}",
                                             uphill_helped=f"{upos} of {un}", fisher_two_sided=round(float(fe2), 3),
                                             fisher_one_sided=round(float(fe1), 3),
                                             spearman=round(float(r0), 3), perm_p=round(pperm, 3)))
with open(OUT, "w") as fh:
    json.dump(RES, fh, indent=1, default=float)
print("\nwrote", os.path.basename(OUT))