"""Regression suite for simroot.pure and for the benchmark harness.

Each audited defect is pinned as a test, then the contracts the method is
allowed to claim in the paper - and only those.
"""
from __future__ import annotations

import math
import os
import sys

import pytest
from scipy.optimize._tstutils import (_APS_TESTS_DICTS,
                                      _CHANDRUPATLA_TESTS_DICTS)

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

from simroot.baselines import bisection, itp
from simroot.pure import GUARDS, PURE_SIMS, pure_root

XTOL = 1e-12
MAXFEV = 200
ACC = 1e-6

SMALL = [
    ("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),
    ("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),
]

# the guards that are ALLOWED to claim a worst-case bound
SOUND_GUARDS = ("bracket", "both")


def _problems(dicts, chandrupatla=False):
    out = []
    for d in dicts:
        a, b = float(d["bracket"][0]), float(d["bracket"][1])
        args = () if chandrupatla else d.get("args", ())
        f = (lambda ff, aa: lambda x: ff(x, *aa))(d["f"], args)
        out.append((d["ID"], f, a, b, float(d["root"])))
    return out


APS = _problems(_APS_TESTS_DICTS)
CHA = _problems(_CHANDRUPATLA_TESTS_DICTS, chandrupatla=True)


def two_x_bisection_bound(a, b):
    return 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4


# --------------------------------------------------------------------------
# audited defects, pinned
# --------------------------------------------------------------------------
@pytest.mark.parametrize("guard", SOUND_GUARDS)
def test_bug_pending_flag_gave_three_x_not_two_x(guard):
    """The bracket guard cleared `pending` on the halving branch and then
    returned early, so TWO similarity steps were granted after every halving
    instead of one - a 3x bisection guarantee dressed up as 2x.  Chandrupatla
    fun4.3 caught it: 139 evaluations against a bound of 114."""
    d = next(x for x in _CHANDRUPATLA_TESTS_DICTS if x["ID"] == "fun4.3")
    a, b = float(d["bracket"][0]), float(d["bracket"][1])
    x, n, conv = pure_root(d["f"], a, b, XTOL, MAXFEV,
                           similarity="ratio", scale="derived", guard=guard)
    assert conv and abs(x - float(d["root"])) <= ACC
    assert n <= two_x_bisection_bound(a, b), f"{n} evaluations, 3x regression"


def test_step_guard_makes_no_worst_case_claim():
    """The step guard is cheap but unsound: a sequence creeping towards a
    non-root with halving steps never trips it.  The suite must NOT assert a
    bound for it - this test exists so that nobody adds one later without
    noticing that the benchmark reports real violations."""
    violations = 0
    for _pid, f, a, b, root in CHA:
        x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                               scale="derived", guard="step")
        if conv and abs(x - root) <= ACC and n > two_x_bisection_bound(a, b):
            violations += 1
    assert violations > 0, ("the step guard no longer violates the 2x bound; "
                            "if that is a real improvement, promote it to "
                            "SOUND_GUARDS and prove it")


# --------------------------------------------------------------------------
# contracts
# --------------------------------------------------------------------------
@pytest.mark.parametrize("guard", SOUND_GUARDS)
def test_sound_guards_respect_the_bound_on_every_benchmark_problem(guard):
    """The only worst-case claim the paper may make."""
    bad = []
    for pid, f, a, b, root in APS + CHA:
        x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                               scale="derived", guard=guard)
        if conv and abs(x - root) <= ACC and n > two_x_bisection_bound(a, b):
            bad.append((pid, n, two_x_bisection_bound(a, b)))
    assert not bad, f"{guard} exceeded the 2x bound on {bad[:5]}"


@pytest.mark.parametrize("sim", sorted(PURE_SIMS))
@pytest.mark.parametrize("prob", SMALL, ids=[p[0] for p in SMALL])
def test_derived_scale_with_a_sound_guard_always_converges(sim, prob):
    _, f, a, b, root = prob
    x, n, conv = pure_root(f, a, b, XTOL, 400, similarity=sim,
                           scale="derived", guard="both")
    assert conv, f"pure.{sim}.both did not converge"
    assert abs(x - root) <= 1e-8


@pytest.mark.parametrize("guard", GUARDS)
def test_never_leaves_the_bracket(guard):
    for _pid, f, a, b, _root in SMALL:
        x, _n, _c = pure_root(f, a, b, XTOL, MAXFEV, guard=guard)
        assert min(a, b) - 1e-9 <= x <= max(a, b) + 1e-9


def test_rejects_unknown_settings_loudly():
    f = lambda t: t * t - 2.0
    with pytest.raises(ValueError):
        pure_root(f, 0.0, 2.0, XTOL, MAXFEV, similarity="no-such")
    with pytest.raises(ValueError):
        pure_root(f, 0.0, 2.0, XTOL, MAXFEV, scale="no-such")
    with pytest.raises(ValueError):
        pure_root(f, 0.0, 2.0, XTOL, MAXFEV, guard="no-such")
    with pytest.raises(ValueError):
        pure_root(f, 3.0, 4.0, XTOL, MAXFEV)      # no sign change


def test_budget_is_never_exceeded():
    for budget in (0, 1, 2, 5, 30):
        for guard in GUARDS:
            _x, n, _c = pure_root(lambda t: math.exp(-t) * (t - 1.0),
                                  -1.0, 40.0, XTOL, budget, guard=guard)
            assert n <= budget


# --------------------------------------------------------------------------
# the competitor must be sound too, or the benchmark is worthless
# --------------------------------------------------------------------------
@pytest.mark.parametrize("n0", [0, 1, 2, 5])
def test_itp_respects_its_own_published_guarantee(n0):
    """ITP's theorem: it never needs more iterations than bisection plus n0.
    Checking the theorem is the strongest validation available for the
    implementation without a reference to diff against.

    The permitted count is n_half + n0 + 1, not n_half + n0: the induction
    gives bracket_j <= 2^(n_max-j+1)*eps, so the width first drops to 2*eps at
    j = n_max, and the iterations actually run j = 0 .. n_max, which is
    n_max + 1 of them.  Stating it as n_max flagged 3 of 154 problems by
    exactly one iteration - the signature of an indexing mistake in the test
    rather than of a broken implementation."""
    bad = []
    for pid, f, a, b, root in APS:
        if not a < b:
            continue
        allowed = math.ceil(math.log2((b - a) / XTOL)) + n0 + 1
        x, n, conv = itp(f, a, b, XTOL, 400, kappa1=0.2 / (b - a), n0=n0)
        if conv and abs(x - root) <= ACC and (n - 2) > allowed:
            bad.append((pid, n - 2, allowed))
    assert not bad, f"ITP broke its own bound on {bad[:5]}"


def test_bisection_is_the_reference_ceiling_everyone_beats():
    """Sanity check on the harness: if the new method were not actually
    solving anything, this is the comparison that would expose it."""
    tot_pure = tot_bis = 0
    for _pid, f, a, b, root in APS:
        _x, n1, c1 = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                               scale="derived", guard="both")
        _y, n2, c2 = bisection(f, a, b, XTOL, MAXFEV)
        tot_pure += n1 if c1 else MAXFEV
        tot_bis += n2 if c2 else MAXFEV
    assert tot_pure < tot_bis, (tot_pure, tot_bis)