"""Plausibility check of PREREG_flat_stick before the full grid: which problems enter
the sticky regime, what order the failed promise reads at entry, and where the cost goes.

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

import math
from collections import Counter as Tally

import simroot.pure as P
from sweep_full import ACC, MAXFEV, XTOL, load_problems

FLAT = ("fun3", "fun4", "fun5", "fun6")
_orig = P._order_from_promise
ENTRIES = []


def _spy(r):
    m = _orig(r)
    ENTRIES.append((r, m))
    return m


def run(gate, f, a, b):
    st = {}
    x, n, conv = P.pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3",
                             guard="project", prestep=gate, stats=st)
    return x, n, conv, st["pre_refused"]


def main():
    P._order_from_promise = _spy
    probs = load_problems()
    for s in ("APS", "CHA"):
        groups = {"flat": [0, 0, 0, 0], "other": [0, 0, 0, 0]}   # problems, entered, base, stick
        ms = Tally()
        for st_, pid, f, a, b, root in probs:
            if st_ != s:
                continue
            g = "flat" if s == "CHA" and pid.split(".")[0] in FLAT else "other"
            _x, n0, c0, _ = run("none", f, a, b)
            ENTRIES.clear()
            x1, n1, c1, k = run("flatstick", f, a, b)
            ok0 = c0 and abs(_x - root) <= ACC
            ok1 = c1 and math.isfinite(x1) and abs(x1 - root) <= ACC
            groups[g][0] += 1
            groups[g][1] += 1 if ENTRIES else 0
            groups[g][2] += n0 if ok0 else MAXFEV
            groups[g][3] += n1 if ok1 else MAXFEV
            for r, m in ENTRIES:
                ms["inf" if math.isinf(m) else ("2-3" if m < 3 else "3-10" if m < 10 else ">=10")] += 1
        for g, (np_, ent, c0, c1) in groups.items():
            if np_:
                print(f"{s} {g:<5}: {np_:3d} problems, entered sticky {ent:3d}, "
                      f"cost none {c0:5d} -> flatstick {c1:5d}")
        print(f"{s} order read at entry: {dict(ms)}")
    print("diag done")


if __name__ == "__main__":
    main()