"""Replay check for PREREG_outcome_memory_2026-09-21.md (B).  Method untouched.

At each traced state n with at least one past state of known outcome: pick the past state
j most similar in rho (ratio similarity of rho_n and rho_j; ties -> most recent).  If the
step taken from j bought less than one halving of the bracket, B bisects, otherwise B
takes our step.  One evaluation from the same state for B and for our step; measured:
halvings of the proven bracket.

    python -u diag_outcome_memory.py
"""
from __future__ import annotations

import math

from scipy.optimize._tstutils import _APS_TESTS_DICTS, _CHANDRUPATLA_TESTS_DICTS

from diag_contraction_similarity import CONFIGS, FLAT
from diag_step_extrapolation import halvings_of
from simroot.pure import pure_root

XTOL, MAXFEV, ACC = 1e-12, 200, 1e-6


def replay(f, a, b, r, cfg):
    st = {"trace": []}
    x, _n, conv = pure_root(f, a, b, XTOL, MAXFEV, stats=st, **cfg)
    if not (conv and abs(float(x) - r) <= ACC):
        return []
    tr = st["trace"]
    widths = [abs(s["B"] - s["A"]) for s in tr]
    rhos = [abs(s["fx"] / s["fp"]) if (s["fp"] not in (None, 0.0)) else None for s in tr]
    rows = []
    for n, s in enumerate(tr):
        if rhos[n] is None or rhos[n] == 0.0 or not math.isfinite(rhos[n]):
            continue
        best, best_sim = None, -1.0
        for j in range(n):                         # outcome of j is known once j+1 exists
            if j + 1 > n or rhos[j] is None or rhos[j] == 0.0 or not math.isfinite(rhos[j]):
                continue
            if widths[j] <= 0.0 or widths[j + 1] <= 0.0:
                continue
            sim = min(rhos[n], rhos[j]) / max(rhos[n], rhos[j])
            if sim >= best_sim:                    # >= keeps the most recent on ties
                best, best_sim = j, sim
        if best is None:
            continue
        outcome = math.log2(widths[best] / widths[best + 1])
        A, FA, B, FB = s["A"], s["FA"], s["B"], s["FB"]
        lo, flo, hi, fhi = (A, FA, B, FB) if A < B else (B, FB, A, FA)
        if hi - lo <= 0.0:
            continue
        ours_L = s["x"] + s["raw"]
        if not (lo < ours_L < hi):
            continue
        bisect = outcome < 1.0
        B_L = 0.5 * (lo + hi) if bisect else ours_L
        h_ours = halvings_of(f, lo, flo, hi, fhi, ours_L, r)
        h_B = halvings_of(f, lo, flo, hi, fhi, B_L, r)
        rows.append((bisect, h_ours, h_B))
    return rows


def main():
    out = {}
    for cname, cfg in CONFIGS.items():
        groups = {"APS": [], "CHA fun1": [], "CHA flat fun3-6": [], "CHA other": []}
        for d in _APS_TESTS_DICTS:
            a, b = (float(v) for v in d["bracket"])
            f = (lambda ff, aa: lambda t: ff(t, *aa))(d["f"], d.get("args", ()))
            groups["APS"] += replay(f, a, b, float(d["root"]), cfg)
        for d in _CHANDRUPATLA_TESTS_DICTS:
            a, b = (float(v) for v in d["bracket"])
            fam = str(d["ID"]).split(".")[0]
            g = ("CHA flat fun3-6" if fam in FLAT else "CHA fun1" if fam == "fun1"
                 else "CHA other")
            groups[g] += replay(d["f"], a, b, float(d["root"]), cfg)
        print(f"\n== {cname}")
        print(f"   {'group':<16} {'states':>6} {'B bisects':>10} {'ours':>7} {'B':>7}"
              "   (bracket halvings per eval, all states with a decision)")
        res = {}
        for g, rows in groups.items():
            n = len(rows)
            m = lambda v: sum(v) / len(v) if v else float("nan")   # noqa: E731
            share = 100.0 * sum(r_[0] for r_ in rows) / max(n, 1)
            ho, hb = m([r_[1] for r_ in rows]), m([r_[2] for r_ in rows])
            res[g] = dict(states=n, bisect_share=share, ours=ho, B=hb)
            print(f"   {g:<16} {n:6d} {share:9.1f}% {ho:7.2f} {hb:7.2f}")
        out[cname] = res
    return out


if __name__ == "__main__":
    main()