"""Stage 6b - is the projection invariant actually holding?

The bound test overshoots by exactly one evaluation on the Chandrupatla
problems.  Two possible causes, and they need different fixes:

  (a) the INVARIANT w_{j+1} <= 2^(n_max - j) * eps is violated somewhere -
      then the implementation is wrong and the guarantee is false;
  (b) the invariant holds and the stated BOUND is off by one - then the
      guarantee is true and the statement of it needs correcting.

Guessing between them is not allowed, so the invariant is instrumented
directly on the worst offenders.
"""
from __future__ import annotations

import math

from scipy.optimize._tstutils import _CHANDRUPATLA_TESTS_DICTS

from simroot.pure import PURE_SIMS

XTOL = 1e-12
MAXFEV = 200
EPS = 0.5 * XTOL


def replay(f, a, b, similarity="ratio"):
    """Re-run the projection rule with the invariant checked at every step.
    Mirrors simroot.pure exactly for scale='derived3', guard='project'."""
    S = PURE_SIMS[similarity]
    half = math.ceil(math.log2(abs(b - a) / XTOL))
    n_max = 2 * half
    fa, fb = f(a), f(b)
    A, FA, B, FB = a, fa, b, fb
    if abs(fa) <= abs(fb):
        x, fx, xp, fp = a, fa, b, fb
    else:
        x, fx, xp, fp = b, fb, a, fa
    xp2 = fp2 = None
    j, nfev, worst = 0, 2, 0.0
    violations = []
    while abs(B - A) > XTOL and nfev < MAXFEV:
        w = abs(B - A)
        lo, hi = min(A, B), max(A, B)
        far = B if abs(B - x) >= abs(A - x) else A
        if x == xp:
            slope = 0.0
        else:
            d12 = (fx - fp) / (x - xp)
            if xp2 is None or xp2 in (x, xp):
                slope = d12
            else:
                d01 = (fp - fp2) / (xp - xp2)
                slope = d12 + (d12 - d01) / (x - xp2) * (x - xp)
        sc = abs(slope) * abs(far - x)
        if sc == 0.0 or not math.isfinite(sc):
            x_new = 0.5 * (A + B)
        else:
            x_new = x + (1.0 - S(abs(fx) / sc)) * (far - x)
            # The CLAMPED form, matching simroot.pure exactly.  The earlier
            # radius form r = allowed - w/2 is a difference of two nearly
            # equal numbers, and the cancellation admitted points that broke
            # the invariant by a relative 8.9e-5 on eleven of these problems.
            allowed = 2.0 ** (n_max - j) * EPS
            if math.isfinite(allowed):
                left, right = hi - allowed, lo + allowed
                x_new = (0.5 * (A + B) if left > right
                         else min(max(x_new, left), right))
        if not math.isfinite(x_new) or not (lo < x_new < hi):
            x_new = 0.5 * (A + B)
        if x_new in (x, xp):
            x_new = 0.5 * (A + B)
            if x_new in (x, xp):
                break
        fn = f(x_new)
        nfev += 1
        j += 1
        xp2, fp2, xp, fp, x, fx = xp, fp, x, fx, x_new, fn
        if fn == 0.0:
            break
        if FA * fn < 0.0:
            B, FB = x_new, fn
        else:
            A, FA = x_new, fn
        allowed = 2.0 ** (n_max - (j - 1)) * EPS
        got = abs(B - A)
        worst = max(worst, got / allowed)
        # The clamp bound hi - allowed is not representable when the endpoints
        # are far larger than the allowance: at |hi| ~ 1e10 the spacing of
        # doubles is ~2e-6, so the bound itself is only known to that
        # accuracy.  This is not a property of the method - bisection's own
        # halving is exact only to the same order - so the invariant is
        # checked to the representation error of the endpoints rather than to
        # the last bit.  The operative guarantee, the EVALUATION COUNT, is
        # checked exactly and separately in tests/test_projection.py.
        slack = 4.0 * math.ulp(max(abs(A), abs(B), 1.0))
        if got > allowed + slack:
            violations.append((j, got, allowed))
    return nfev, j, n_max, worst, violations


def main():
    print("projection invariant  w_{j+1} <= 2^(n_max - j) * eps\n")
    print(f"{'problem':10s}{'nfev':>7s}{'steps j':>9s}{'n_max':>8s}"
          f"{'worst w/allowed':>18s}{'violations':>12s}")
    print("-" * 66)
    any_viol = 0
    for d in _CHANDRUPATLA_TESTS_DICTS:
        a, b = float(d["bracket"][0]), float(d["bracket"][1])
        try:
            nfev, j, n_max, worst, viol = replay(d["f"], a, b)
        except Exception as exc:
            print(f"{d['ID']:10s} skipped: {exc}")
            continue
        if viol or j > n_max:
            print(f"{d['ID']:10s}{nfev:>7d}{j:>9d}{n_max:>8d}"
                  f"{worst:>18.6f}{len(viol):>12d}")
            any_viol += len(viol)
    print(f"\ntotal invariant violations across the set: {any_viol}")
    print("steps j never exceeding n_max is the other half of the claim; "
          "rows above\nare the only ones where either could fail.")
    if any_viol == 0:
        print("\nThe invariant holds. The overshoot is therefore in the "
              "STATED bound:\nthe run terminates at step j = n_max, and steps "
              "are counted j = 0..n_max,\nso the count is n_max + 1 steps plus "
              "the two endpoint evaluations,\ni.e. 2*ceil(log2(w/xtol)) + 3, "
              "not + 2.")


if __name__ == "__main__":
    main()