"""Stage 1 of PREREG_flat_twice_2026-09-21.md: does 'the promise failed twice in a row'
separate the flat families better than a single failure?  Method untouched.

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

from scipy.optimize._tstutils import _APS_TESTS_DICTS, _CHANDRUPATLA_TESTS_DICTS

from simroot.pure import pure_root

XTOL, MAXFEV, ACC = 1e-12, 200, 1e-6
R_STAR = 0.25
CFG = dict(similarity="ratio", scale="derived3", guard="project")
FLAT = ("fun3", "fun4", "fun5", "fun6")


def run_flags(f, a, b, r_true):
    """(model steps judged, single flags, twice-in-a-row flags) for one run."""
    evals = []

    def g(t):
        v = f(t)
        evals.append((t, v))
        return v
    st = {"trace": []}
    x, _n, conv = pure_root(g, a, b, XTOL, MAXFEV, stats=st, **CFG)
    if not (conv and abs(float(x) - r_true) <= ACC):
        return 0, 0, 0
    order = [e[0] for e in evals]
    judged = single = twice = run = 0
    for s in st["trace"]:
        try:
            i = order.index(s["x"])
        except ValueError:
            run = 0
            continue
        if i + 1 >= len(evals):
            continue
        x_next, f_next = evals[i + 1]
        if x_next != s["x"] + s["raw"] or s["fx"] == 0.0:
            run = 0                       # not the model's own step: the run is broken
            continue
        judged += 1
        r = f_next / s["fx"]
        if r >= R_STAR:
            run += 1
            single += 1
            twice += run >= 2
        else:
            run = 0
    return judged, single, twice


def main():
    groups = {g: [0, 0, 0] for g in ("APS", "CHA fun1", "CHA flat fun3-6", "CHA fun7",
                                      "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", ()))
        for k, v in enumerate(run_flags(f, a, b, float(d["root"]))):
            groups["APS"][k] += v
    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 fun7" if fam == "fun7" else "CHA other")
        for k, v in enumerate(run_flags(d["f"], a, b, float(d["root"]))):
            groups[g][k] += v
    print(f"{'group':<18} {'model steps':>12} {'single':>8} {'twice':>8}")
    share = {}
    for g, (j, s1, s2) in groups.items():
        share[g] = (100.0 * s1 / max(j, 1), 100.0 * s2 / max(j, 1))
        print(f"{g:<18} {j:12d} {share[g][0]:7.1f}% {share[g][1]:7.1f}%")
    flat, aps, f1 = share["CHA flat fun3-6"], share["APS"], share["CHA fun1"]
    t1 = flat[1] > aps[1] and flat[1] > f1[1]
    ratio1 = flat[0] / aps[0] if aps[0] else float("inf")
    ratio2 = flat[1] / aps[1] if aps[1] else float("inf")
    print(f"\nT1 twice: flat {flat[1]:.1f}% > APS {aps[1]:.1f}% and > fun1 {f1[1]:.1f}%  ->  "
          f"{'PASS' if t1 else 'FAIL'}")
    print(f"T2 flat/APS ratio: twice {ratio2:.2f} vs single {ratio1:.2f}  ->  "
          f"{'PASS' if ratio2 > ratio1 else 'FAIL'}")
    return share


if __name__ == "__main__":
    main()