"""Stage 0b - the PURE similarity idea, taken literally as a STEP-SIZE rule.

Valeri's formulation, word for word:

    "in bisection the next point is chosen through half of the interval.
     With similarity the next point is chosen proportionally to the value of
     the similarity function (0 = not similar at all, 1 = identical) of the
     result to the expectation (e.g. zero).  So the step is LARGER when the
     result is not similar to the expectation, and SMALLER when it is."

This is NOT the same object as Stage 0.  Stage 0 asked "where between a and b
does the root sit" (an interpolation rule, which turned out to be regula
falsi identically).  Here the rule is:

    direction  <- given by the sign (towards the side that brackets the root)
    step       <- (1 - S(f(x_n))) * (distance to the far end of the bracket)

with S a genuine [0,1] similarity function of a single value f(x) to the
target 0.  Bisection is the special case S == 1/2 for every input, i.e.
"the similarity function carries no information".

The whole content of the method is therefore the similarity function itself.
The one thing S needs that a two-point rule did not is a SCALE: how large is
|f| = 3?  Any constant here would be a hidden tuning parameter.  Three
scale sources are compared, all taken from the run itself:

    s0      the value at the starting point            (fixed reference)
    srun    the largest |f| seen so far                (growing reference)
    slive   the largest |f| at the current bracket ends (shrinking reference)

and the last one is the one that reproduces Stage 0 / regula falsi for the
linear S, which is exactly the collapse this stage is meant to expose.
"""
from __future__ import annotations

from simroot.baselines import _opposite, _same_sign  # 2026-09-21 fix

import math

from simroot.baselines import Counter, _Budget, bisection, brent

XTOL = 1e-12
MAXFEV = 200

# --------------------------------------------------------------------------
# genuine [0,1] similarity functions of a normalised residual z = f / scale
# S(0) = 1 (identical to the expectation), S(large) -> 0 (not similar at all)
# --------------------------------------------------------------------------
SIMS_01 = {
    # linear ramp, saturating at the reference scale
    "linear":   lambda z: max(0.0, 1.0 - abs(z)),
    # rational, never saturates
    "rational": lambda z: 1.0 / (1.0 + abs(z)),
    # rational in the square
    "rat_sq":   lambda z: 1.0 / (1.0 + z * z),
    # exponential decay of the distance
    "exp_abs":  lambda z: math.exp(-abs(z)),
    # the article's negative squared Euclidean distance, exponentiated to [0,1]
    "exp_sq":   lambda z: math.exp(-z * z),
    # cosine similarity of (1, f) against (1, 0) - article eq. (4) on the
    # value axis lifted to 2-D
    "cosine":   lambda z: 1.0 / math.sqrt(1.0 + z * z),
}

SCALES = ("s0", "srun", "slive")


def pure_similarity_step(f, a, b, xtol, maxfev, sim="rational", scale="s0"):
    """The literal step-size rule.  Bracket kept only for safety and for
    knowing which way to walk; the step LENGTH comes from similarity alone."""
    S = SIMS_01[sim]
    g = Counter(f, maxfev)
    x = 0.5 * (a + b)
    try:
        fa, fb = g(a), g(b)
        if fa == 0.0:
            return a, g.n, True
        if fb == 0.0:
            return b, g.n, True

        # start from the endpoint whose result is already more similar to 0
        if abs(fa) <= abs(fb):
            x, fx, far, ffar = a, fa, b, fb
        else:
            x, fx, far, ffar = b, fb, a, fa

        s0 = abs(fx)
        smax = max(abs(fa), abs(fb))

        while abs(far - x) > xtol:
            if scale == "s0":
                sc = s0
            elif scale == "srun":
                sc = smax
            else:
                sc = max(abs(fx), abs(ffar))
            if sc == 0.0:
                return x, g.n, True

            s = S(abs(fx) / sc)          # 1 = the result already IS zero
            step = (1.0 - s) * (far - x)  # dissimilar -> long step
            x_new = x + step
            lo, hi = (x, far) if x < far else (far, x)
            if not math.isfinite(x_new) or not (lo < x_new < hi):
                x_new = 0.5 * (x + far)

            fnew = g(x_new)
            if fnew == 0.0:
                return x_new, g.n, True
            smax = max(smax, abs(fnew))
            if _opposite(fx, fnew):
                # the root is between the old point and the new one
                far, ffar = x, fx
            x, fx = x_new, fnew
        return x, g.n, True
    except _Budget:
        return x, g.n, False


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 run():
    methods = [("bisection (S=1/2)", lambda f, a, b: bisection(f, a, b, XTOL, MAXFEV)),
               ("brent  [reference]", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV))]
    for sc in SCALES:
        for s in SIMS_01:
            methods.append((
                f"pure.{s}.{sc}",
                (lambda ss, scl: lambda f, a, b: pure_similarity_step(
                    f, a, b, XTOL, MAXFEV, sim=ss, scale=scl))(s, sc)))

    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}")


if __name__ == "__main__":
    run()