"""Stage 0d - faster similarity functions, and similarity DERIVED FROM DATA.

Two questions, in order.

(1) Ratio-type similarities.  S(x,y) = |x|/|y| and relatives change much
    faster than 1/(1+z) or exp(-z).  Stage 0c derived the contraction rate of
    the pure step rule as

        rate = | 1 - c * |f'(r)| * D / scale | ,   c = slope of (1 - S) at 0

    so a similarity that changes FASTER has a larger c and moves the rate away
    from 1 - but it overshoots just as easily as it undershoots, because
    nothing ties c to the problem.  Measured below, not argued.

(2) Similarity derived from the data.  The same formula says the rate is
    exactly ZERO when

        c = scale / ( |f'(r)| * D ).

    Nothing in that expression has to be chosen: |f'| is estimated by the
    secant slope of the two most recent evaluations, D by the current distance
    to the far end, scale by the value in play.  Substituting it back:

        step = (1 - S)*(far - x) = c*|f|/scale * D = |f| / |f'|

    - the calibrated similarity step IS the Newton/secant step.  So "derive
    the similarity parameter from the data" is not a tweak of the pure idea;
    it is the statement that the value axis has a LOCAL scale, and reading
    that scale off the data is what turns linear convergence into superlinear.

    This is the same move the article makes when it normalises the input space
    from the data instead of choosing membership widths by hand.
"""
from __future__ import annotations

import math

from simroot.baselines import bisection, brent
from simroot.pure import PURE_SIMS, pure_root

XTOL = 1e-12
MAXFEV = 200

# --------------------------------------------------------------------------
# (1) fast, ratio-type [0,1] similarities of a normalised residual z
# --------------------------------------------------------------------------
# The similarity functions and the step rule now live in simroot.pure, which
# is the single implementation behind every number in this project - the
# benchmark, the tables below and the paper all call the same code.  This file
# keeps only the subset of similarities that stage 0d is about, plus two thin
# delegations so the stage-0d and stage-0e tables read unchanged.
FAST_SIMS = {k: PURE_SIMS[k] for k in
             ("ratio", "sqrt", "cbrt", "ratio_soft", "rational", "exp_abs")}


def pure_step(f, a, b, xtol, maxfev, sim):
    """The step rule with the scale CHOSEN: |f| at the starting point."""
    return pure_root(f, a, b, xtol, maxfev, similarity=sim, scale="start",
                     guard="none")


def derived_step(f, a, b, xtol, maxfev, sim="ratio", guard=True):
    """The step rule with the scale DERIVED from the two most recent
    evaluations.  `guard=True` keeps the bracket safeguard that stage 0d
    measured; stage 0e compares it against no safeguard at all."""
    return pure_root(f, a, b, xtol, maxfev, similarity=sim, scale="derived",
                     guard="bracket" if guard else "none")


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:
            xx, n, conv = fn(f, a, b)
            ok = conv and abs(xx - 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 rate_table():
    """The measured contraction rate of each fast similarity, against the
    rate predicted by the stage-0c law."""
    f = lambda x: x * x - 2.0
    root = math.sqrt(2.0)
    fp, D, scale = 2.0 * math.sqrt(2.0), 2.0 - math.sqrt(2.0), 2.0
    print("\ncontraction rate on x^2 - 2, pure rule, scale = |f(x0)|")
    print(f"{'similarity':10s} {'c (slope of 1-S at 0)':>24s} "
          f"{'predicted':>11s} {'measured':>10s}")
    print("-" * 60)
    for sim in FAST_SIMS:
        S = FAST_SIMS[sim]
        h = 1e-7
        c = (1.0 - S(h)) / h          # numerical slope of (1 - S) at 0
        pred = abs(1.0 - c * fp * D / scale) if math.isfinite(c) else float("inf")
        x, errs = 0.0, []
        for _ in range(14):
            x = x + (1.0 - S(abs(f(x)) / scale)) * (2.0 - x)
            errs.append(abs(x - root))
        win = [errs[i + 1] / errs[i] for i in range(len(errs) - 1) if errs[i] > 1e-11]
        meas = sum(win[-3:]) / len(win[-3:]) if win else float("nan")
        cs = f"{c:.3e}" if math.isfinite(c) else "inf"
        ps = f"{pred:.4f}" if math.isfinite(pred) else "inf (overshoot)"
        print(f"{sim:10s} {cs:>24s} {ps:>11s} {meas:>10.4f}")


def main():
    rate_table()
    _table([("bisection", lambda f, a, b: bisection(f, a, b, XTOL, MAXFEV)),
            ("brent [reference]", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV))] +
           [(f"fast.{s}", (lambda ss: lambda f, a, b: pure_step(
               f, a, b, XTOL, MAXFEV, ss))(s)) for s in FAST_SIMS],
           "(1) FAST similarity functions, parameter CHOSEN")
    _table([("bisection", lambda f, a, b: bisection(f, a, b, XTOL, MAXFEV)),
            ("brent [reference]", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV))] +
           [(f"derived.{s}", (lambda ss: lambda f, a, b: derived_step(
               f, a, b, XTOL, MAXFEV, ss))(s)) for s in FAST_SIMS],
           "(2) similarity SCALE DERIVED FROM DATA (secant slope), guarded")


if __name__ == "__main__":
    main()