"""Replay check for PREREG_step_extrapolation_2026-09-21.md (A).  Method untouched.

At every traced state with a known previous displacement, three candidate landings from
the SAME state get one evaluation each: our step, step A = d_n/(1 - S) with S the ratio
similarity of the two successive steps (0 if they point opposite ways), and the midpoint.
Measured: halvings of the proven bracket that single evaluation buys.

    python -u diag_step_extrapolation.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_model_order import bracket_step
from simroot.pure import pure_root

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


def halvings_of(f, lo, flo, hi, fhi, landing, r):
    W = hi - lo
    Wn = bracket_step(f, lo, flo, hi, fhi, landing)
    return math.log2(W / max(Wn, math.ulp(r) if r != 0.0 else 5e-324))


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 []
    rows = []
    for s in st["trace"]:
        if s["xp"] is None:
            continue
        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
        d_n, d_p = s["raw"], s["x"] - s["xp"]
        if d_n == 0.0 or d_p == 0.0 or (d_n > 0.0) != (d_p > 0.0):
            S = 0.0
        else:
            S = min(abs(d_n), abs(d_p)) / max(abs(d_n), abs(d_p))
        if S >= 1.0:
            continue                              # a geometric series needs S < 1
        ours_L = s["x"] + d_n
        ext_L = s["x"] + d_n / (1.0 - S)
        mid_L = 0.5 * (lo + hi)
        if not (lo < ours_L < hi):
            continue
        h_ours = halvings_of(f, lo, flo, hi, fhi, ours_L, r)
        h_mid = halvings_of(f, lo, flo, hi, fhi, mid_L, r)
        ext_ok = lo < ext_L < hi
        h_ext = halvings_of(f, lo, flo, hi, fhi, ext_L, r) if ext_ok else None
        rows.append((S, h_ours, h_ext, h_mid, ext_ok))
    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} {'S>0':>6} {'A lands out':>11} "
              f"{'paired n':>8} {'ours':>6} {'step A':>7} {'bisect':>7}   (bracket halvings per eval)")
        res = {}
        for g, rows in groups.items():
            n = len(rows)
            pos = [r_ for r_ in rows if r_[0] > 0.0]
            out_ = sum(1 for r_ in pos if not r_[4])
            paired = [r_ for r_ in pos if r_[4]]
            m = lambda v: sum(v) / len(v) if v else float("nan")   # noqa: E731
            ho = m([p[1] for p in paired])
            he = m([p[2] for p in paired])
            hm = m([p[3] for p in paired])
            res[g] = dict(states=n, s_pos=len(pos), out=out_, paired=len(paired),
                          ours=ho, ext=he, mid=hm)
            print(f"   {g:<16} {n:6d} {100.0 * len(pos) / max(n, 1):5.1f}% "
                  f"{100.0 * out_ / max(len(pos), 1):10.1f}% {len(paired):8d} "
                  f"{ho:6.2f} {he:7.2f} {hm:7.2f}")
        out[cname] = res
    return out


if __name__ == "__main__":
    main()