"""Stage 2 - the step guard against the bracket guard, and against no guard.

The bracket guard costs 64 extra evaluations over the eight stage-0 problems
(150 vs 86) because a secant step collapses the ERROR while barely moving the
proven BRACKET.  The step guard tests the thing that is actually converging -
the step length - and should cost close to nothing while keeping the
worst-case bound.
"""
from __future__ import annotations

import math

from simroot.baselines import bisection, brent, itp, ridders, toms748
from simroot.pure import GUARDS, PURE_SIMS, pure_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),
]


def table(methods, title):
    print("\n" + title)
    width = max(len(n) for n, _ in methods) + 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 _, f, a, b, root in PROBLEMS:
            x, n, conv = fn(f, a, b)
            ok = conv and abs(x - root) <= 1e-8
            cells.append(f"{n:>15d}" if ok else f"{'FAIL(' + str(n) + ')':>15s}")
            total += n
            fails += 0 if ok else 1
        print(f"{name:{width}s}" + "".join(cells) + f"{total:>8d}{fails:>6d}")


def main():
    ref = [("bisection", lambda f, a, b: bisection(f, a, b, XTOL, MAXFEV)),
           ("brent", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV)),
           ("ridders", lambda f, a, b: ridders(f, a, b, XTOL, MAXFEV)),
           ("toms748", lambda f, a, b: toms748(f, a, b, XTOL, MAXFEV)),
           ("itp (3 tuned consts)", lambda f, a, b: itp(f, a, b, XTOL, MAXFEV))]

    table(ref + [(f"pure.ratio.{gu}",
                  (lambda g: lambda f, a, b: pure_root(
                      f, a, b, XTOL, MAXFEV, similarity="ratio",
                      scale="derived", guard=g))(gu)) for gu in GUARDS],
          "the three safeguards, similarity = ratio, scale = derived")

    table(ref + [(f"pure.{s}.step",
                  (lambda ss: lambda f, a, b: pure_root(
                      f, a, b, XTOL, MAXFEV, similarity=ss,
                      scale="derived", guard="step"))(s)) for s in PURE_SIMS],
          "every similarity function, scale = derived, guard = step")


if __name__ == "__main__":
    main()