"""Stage 1 - does the SIM formulation repair the failures found in Stage 0?

Stage 0 showed the raw similarity-proportional step (= regula falsi) beats
bisection on smooth problems by 2-3x but fails catastrophically when one end
of the bracket goes stale.  Stage 1 asks whether the article's actual
machinery - prototype memory + softmax aggregation - removes that failure
WITHOUT introducing a tuning parameter.

Same 8 problems as Stage 0, plus bisection and Brent for orientation.
"""
from __future__ import annotations

import math

from simroot.baselines import bisection, brent, regula_falsi
from simroot.sim import SIMILARITIES, sim_root

XTOL = 1e-12
MAXFEV = 200

PROBLEMS = [
    ("x^2 - 2", lambda x: x * x - 2.0, 0.0, 2.0, math.sqrt(2.0)),
    ("cos(x) - x", lambda x: math.cos(x) - x, 0.0, 1.0, 0.7390851332151607),
    ("exp(x) - 2", lambda x: math.exp(x) - 2.0, 0.0, 2.0, math.log(2.0)),
    ("x^3 - x - 2", lambda x: x ** 3 - x - 2.0, 1.0, 2.0, 1.5213797068045676),
    ("sin(x) - 0.5", lambda x: math.sin(x) - 0.5, 0.0, 1.5, math.asin(0.5)),
    ("x^10 - 1  STALL", lambda x: x ** 10 - 1.0, 0.0, 1.3, 1.0),
    ("exp(-x)(x-1) FLAT", lambda x: math.exp(-x) * (x - 1.0), -1.0, 40.0, 1.0),
    ("atan(x)", lambda x: math.atan(x), -1.0, 3.0, 0.0),
]

METHODS = [
    ("bisection", lambda f, a, b: bisection(f, a, b, XTOL, MAXFEV)),
    ("regula_falsi", lambda f, a, b: regula_falsi(f, a, b, XTOL, MAXFEV)),
    ("brent", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV)),
]
for _s in SIMILARITIES:
    for _p in ("point", "pair", "pair2d"):
        for _g in (False, True):
            METHODS.append((
                f"SIM.{_p}.{_s}" + (".guard" if _g else ""),
                (lambda s, p, gu: lambda f, a, b: sim_root(
                    f, a, b, XTOL, MAXFEV, similarity=s, prototypes=p,
                    guard=gu))(_s, _p, _g),
            ))


def run():
    names = [n for n, _ in METHODS]
    width = max(len(n) for n in names) + 1
    header = f"{'method':{width}s}" + "".join(f"{p[0][:13]:>15s}" for p in PROBLEMS) + f"{'TOTAL':>8s}{'FAIL':>6s}"
    print(header)
    print("-" * len(header))
    for name, fn in METHODS:
        cells, total, fails = [], 0, 0
        for pname, f, a, b, root in PROBLEMS:
            x, n, conv = fn(f, a, b)
            err = abs(x - root)
            ok = conv and err <= 1e-8
            if not ok:
                fails += 1
                cells.append(f"{'FAIL(' + str(n) + ')':>15s}")
            else:
                cells.append(f"{n:>15d}")
            total += n
        print(f"{name:{width}s}" + "".join(cells) + f"{total:>8d}{fails:>6d}")


if __name__ == "__main__":
    run()