"""Regression suite for the similarity root finders and the baselines.

Every bug found in a static audit becomes a test here (project rule), so the
four audited defects are pinned first, then the behavioural contracts.
"""
from __future__ import annotations

import math
import os
import sys

import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from simroot.baselines import (BASELINES, bisection, brent, illinois, itp,
                               regula_falsi, ridders, toms748)
from simroot.sim import SIMILARITIES, sim_root
from stage0b_pure_similarity import SIMS_01, pure_similarity_step
from stage0d_derived_similarity import FAST_SIMS, derived_step, pure_step

XTOL = 1e-12
MAXFEV = 300

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

GUARDED_SIM = [(s, p) for s in SIMILARITIES for p in ("point", "pair", "pair2d")]

# Interpolation-only methods: no bisection safeguard, hence no worst-case
# guarantee.  They are in the suite as references, not as claims.
NO_GUARANTEE = {"regula_falsi", "pegasus"}


# --------------------------------------------------------------------------
# audited defects, pinned
# --------------------------------------------------------------------------
@pytest.mark.parametrize("name", sorted(BASELINES))
def test_bug1_exhausted_budget_is_an_honest_failure(name):
    """A budget that runs out on the very first evaluation must return
    converged=False, never raise UnboundLocalError."""
    method = BASELINES[name]
    for budget in (0, 1, 2):
        x, n, conv = method(lambda t: t * t - 2.0, 0.0, 2.0, XTOL, budget)
        assert n <= budget
        assert math.isfinite(x)
        assert conv is False or n <= budget


def test_bug2_itp_stops_at_the_same_bracket_as_everyone_else():
    """ITP is natively written against a half-width epsilon.  If xtol were fed
    in as epsilon, ITP would stop at twice the bracket width of its
    competitors and win the benchmark on the stopping rule.  The final
    accuracy must be comparable to the other methods, not 2x looser."""
    f = lambda t: t * t - 2.0
    root = math.sqrt(2.0)
    xtol = 1e-8
    x_itp, _, c1 = itp(f, 0.0, 2.0, xtol, MAXFEV)
    x_bis, _, c2 = bisection(f, 0.0, 2.0, xtol, MAXFEV)
    assert c1 and c2
    assert abs(x_itp - root) <= xtol
    assert abs(x_bis - root) <= xtol


def test_bug3_itp_rejects_a_reversed_bracket_loudly():
    with pytest.raises(ValueError):
        itp(lambda t: t * t - 2.0, 2.0, 0.0, XTOL, MAXFEV)


def test_bug4_stagnation_is_not_reported_as_success():
    """sim_root used to return converged=True when it could no longer split
    the bracket, regardless of whether the tolerance had been met.  A
    tolerance far below what float can represent on this bracket must not
    come back as a success."""
    f = lambda t: math.exp(-t) * (t - 1.0)
    x, n, conv = sim_root(f, -1.0, 40.0, 1e-300, 400,
                          similarity="inv", prototypes="pair2d", guard=True)
    if not conv:
        return                      # honest failure - fine
    assert abs(x - 1.0) <= 1e-8, "claimed success while stagnated"


# --------------------------------------------------------------------------
# behavioural contracts
# --------------------------------------------------------------------------
@pytest.mark.parametrize("name", sorted(BASELINES))
@pytest.mark.parametrize("prob", PROBLEMS, ids=[p[0] for p in PROBLEMS])
def test_baselines_find_the_root(name, prob):
    _, f, a, b, root = prob
    x, n, conv = BASELINES[name](f, a, b, XTOL, MAXFEV)
    if name in NO_GUARANTEE:
        # These carry no bisection safeguard and therefore no worst-case
        # bound; the literature says so and the flat problem shows it.  What
        # is required of them is only that they stay honest about it.
        assert n <= MAXFEV
        assert math.isfinite(x)
        if not conv:
            return
    assert conv, f"{name} did not converge"
    assert abs(x - root) <= 1e-8, f"{name} converged to the wrong point"
    assert n <= MAXFEV


@pytest.mark.parametrize("sim,proto", GUARDED_SIM,
                         ids=[f"{s}.{p}" for s, p in GUARDED_SIM])
@pytest.mark.parametrize("prob", PROBLEMS, ids=[p[0] for p in PROBLEMS])
def test_guarded_sim_never_fails(sim, proto, prob):
    """The parameter-free safeguard is the whole reason it exists: with it on,
    no similarity function and no prototype level may fail on any problem."""
    _, f, a, b, root = prob
    x, n, conv = sim_root(f, a, b, XTOL, MAXFEV,
                          similarity=sim, prototypes=proto, guard=True)
    assert conv, f"SIM.{proto}.{sim} did not converge"
    assert abs(x - root) <= 1e-8
    assert n <= MAXFEV


@pytest.mark.parametrize("sim,proto", GUARDED_SIM,
                         ids=[f"{s}.{p}" for s, p in GUARDED_SIM])
@pytest.mark.parametrize("prob", PROBLEMS, ids=[p[0] for p in PROBLEMS])
def test_guard_bounds_the_worst_case_by_twice_bisection(sim, proto, prob):
    """The safeguard pays back one bisection step for every similarity step
    that failed to halve the bracket, so the worst case is bounded by twice
    the BISECTION BOUND.

    The comparison must be against the bound ceil(log2(w/xtol)), not against
    what bisection happened to spend on this problem: on atan over [-1, 3]
    bisection lands exactly on the root in 4 evaluations, which is luck, not
    a worst case, and comparing against it asserts something that was never
    claimed."""
    _, f, a, b, _ = prob
    n_bound = math.ceil(math.log2(abs(b - a) / XTOL))
    _, n_sim, conv = sim_root(f, a, b, XTOL, MAXFEV,
                              similarity=sim, prototypes=proto, guard=True)
    assert conv
    assert n_sim <= 2 * n_bound + 4, f"{n_sim} > 2*{n_bound}+4"


@pytest.mark.parametrize("prob", PROBLEMS, ids=[p[0] for p in PROBLEMS])
def test_every_method_stays_inside_the_initial_bracket(prob):
    """A bracketing method must never return a point outside the bracket it
    was given, converged or not."""
    _, f, a, b, _ = prob
    lo, hi = min(a, b), max(a, b)
    for name, m in BASELINES.items():
        x, _, _ = m(f, a, b, XTOL, MAXFEV)
        assert lo - 1e-9 <= x <= hi + 1e-9, f"{name} left the bracket"
    for sim, proto in GUARDED_SIM:
        x, _, _ = sim_root(f, a, b, XTOL, MAXFEV, similarity=sim,
                           prototypes=proto, guard=True)
        assert lo - 1e-9 <= x <= hi + 1e-9, f"SIM.{proto}.{sim} left the bracket"


def test_the_pure_two_point_similarity_step_is_regula_falsi_exactly():
    """The stage-0 identity, as an executable claim rather than an assertion
    in prose: normalising the similarity by the two bracket ends reproduces
    the false-position formula bit for bit."""
    for fa, fb in ((-5.436, 1.6e-16), (-1.0, 3.0), (2.5, -0.125), (-1e-9, 7.0)):
        a, b = -1.0, 40.0
        s_a = abs(fb) / (abs(fa) + abs(fb))
        s_b = abs(fa) / (abs(fa) + abs(fb))
        x_sim = s_a * a + s_b * b
        x_falsi = (a * fb - b * fa) / (fb - fa)
        assert abs(x_sim - x_falsi) <= 1e-12 * max(1.0, abs(x_falsi))


def test_similarity_shape_predicts_the_contraction_rate():
    """The stage-0c law: rate = |1 - c*|f'(r)|*D/scale| with c the slope of
    (1 - S) at zero.  Pinned for the two linear-approach similarities, where
    it was measured to four digits."""
    f = lambda t: t * t - 2.0
    root, fp, D, scale = math.sqrt(2.0), 2.0 * math.sqrt(2.0), 2.0 - math.sqrt(2.0), 2.0
    predicted = abs(1.0 - 1.0 * fp * D / scale)
    for sim in ("rational", "exp_abs"):
        S = SIMS_01[sim]
        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]
        measured = sum(win[-3:]) / len(win[-3:])
        assert abs(measured - predicted) < 1e-3, (sim, measured, predicted)


def test_quadratic_approach_similarities_stagnate():
    """The other half of the same law: a similarity reaching 1 quadratically
    has c = 0, so the predicted rate is 1 and the iteration must crawl."""
    f = lambda t: t * t - 2.0
    root = math.sqrt(2.0)
    for sim in ("rat_sq", "exp_sq", "cosine"):
        S = SIMS_01[sim]
        x = 0.0
        for _ in range(14):
            x = x + (1.0 - S(abs(f(x)) / 2.0)) * (2.0 - x)
        assert abs(x - root) > 1e-3, f"{sim} was expected to stagnate"


@pytest.mark.parametrize("sim", sorted(FAST_SIMS))
@pytest.mark.parametrize("prob", PROBLEMS, ids=[p[0] for p in PROBLEMS])
def test_derived_scale_converges_where_chosen_scale_need_not(sim, prob):
    """The data-derived local scale plus the parameter-free safeguard must
    always converge; the chosen-scale version carries no such promise and is
    only required not to crash."""
    _, f, a, b, root = prob
    x, n, conv = derived_step(f, a, b, XTOL, MAXFEV, sim=sim)
    assert conv, f"derived.{sim} did not converge"
    assert abs(x - root) <= 1e-8
    x2, n2, _ = pure_step(f, a, b, XTOL, MAXFEV, sim)
    assert math.isfinite(x2) and n2 <= MAXFEV


def test_counter_never_exceeds_the_budget():
    for budget in (3, 7, 25):
        for name, m in BASELINES.items():
            _, n, _ = m(lambda t: math.exp(-t) * (t - 1.0), -1.0, 40.0,
                        XTOL, budget)
            assert n <= budget, f"{name} spent {n} of {budget}"
        _, n, _ = sim_root(lambda t: math.exp(-t) * (t - 1.0), -1.0, 40.0,
                           XTOL, budget, prototypes="pair2d", guard=True)
        assert n <= budget


def test_root_already_at_an_endpoint_costs_two_evaluations():
    for name, m in BASELINES.items():
        x, n, conv = m(lambda t: t - 1.0, 1.0, 5.0, XTOL, MAXFEV)
        assert conv and abs(x - 1.0) <= 1e-8, name
        assert n <= 3, f"{name} spent {n} on a root sitting on the bracket end"