"""Stage 1 of PREREG_flat_regime_2026-09-21.md: is the 'unkept promise' flag selective?

Flag after a step when: the step was the model's own (the evaluation that followed a traced
state is exactly x + raw, so no gate or safeguard replaced it), f did not change sign, and
r = f(x_new)/f(x) >= 1/4 - the break-even of a tangent step against a halving at m = 2.

    python -u diag_flat_flag.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 flags(f, a, b, r_true):
    """(model steps judged, flags raised) for one run of the proposed rule."""
    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
    order = [e[0] for e in evals]
    judged = raised = 0
    for s in st["trace"]:
        try:
            i = order.index(s["x"])
        except ValueError:
            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:
            continue                     # not the model's own step
        judged += 1
        r = f_next / s["fx"]
        if r > 0.0 and r >= R_STAR:
            raised += 1
    return judged, raised


def main():
    groups = {"APS": [0, 0], "CHA fun1": [0, 0], "CHA flat fun3-6": [0, 0],
              "CHA fun7": [0, 0], "CHA other": [0, 0]}
    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", ()))
        j, r = flags(f, a, b, float(d["root"]))
        groups["APS"][0] += j
        groups["APS"][1] += r
    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")
        j, r = flags(d["f"], a, b, float(d["root"]))
        groups[g][0] += j
        groups[g][1] += r
    print(f"{'group':<18} {'model steps':>12} {'flagged':>8} {'share':>7}")
    out = {}
    for g, (j, r) in groups.items():
        out[g] = (j, r, 100.0 * r / max(j, 1))
        print(f"{g:<18} {j:12d} {r:8d} {out[g][2]:6.1f}%")
    flat, aps, f1 = (out["CHA flat fun3-6"][2], out["APS"][2], out["CHA fun1"][2])
    print(f"\nF1 flat > APS and flat > fun1: {flat:.1f}% vs {aps:.1f}% / {f1:.1f}%  ->  "
          f"{'PASS' if flat > aps and flat > f1 else 'FAIL'}")
    print("bar: m>=2 fires 0.2% canonical / 5.8% flat")
    return out


if __name__ == "__main__":
    main()