"""Pins for PREREG_extrapolate_in_method_2026-09-21.md: step extrapolation (A)."""
from __future__ import annotations

import math

import pytest

from simroot.pure import pure_root

SMOOTH = (
    (lambda t: t ** 3 - 2.0 * t - 5.0, 1.0, 4.0),
    (lambda t: math.exp(t) - 2.0, 0.0, 3.0),
    (lambda t: math.sin(t) - 0.5, 0.0, 1.5),
)
FLAT = (
    (lambda t: 6.0 * (t - 2.0) ** 5, 0.0, 10.0),
    (lambda t: t ** 9, -1.0, 2.0),
    (lambda t: (t - 3.0) ** 3, 0.0, 10.0),
)
XTOL, MAXFEV = 1e-12, 200


@pytest.mark.parametrize("f,a,b", SMOOTH + FLAT)
@pytest.mark.parametrize("guard", ("project", "safe", "credit"))
def test_default_is_the_published_method_bit_for_bit(f, a, b, guard):
    ref = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3", guard=guard)
    off = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3", guard=guard,
                    extrapolate="none")
    assert ref == off


@pytest.mark.parametrize("f,a,b", SMOOTH + FLAT)
@pytest.mark.parametrize("guard", ("project", "credit"))
def test_the_guarantee_survives_extrapolation(f, a, b, guard):
    x, n, ok = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3",
                         guard=guard, extrapolate="steps")
    assert ok and min(a, b) <= x <= max(a, b)
    half = math.ceil(math.log2(abs(b - a) / XTOL))
    bound = 2 * half + 3 if guard == "project" else half + 4
    assert n <= bound, (n, bound)


@pytest.mark.parametrize("f,a,b", FLAT[1:])
def test_extrapolation_breaks_the_flat_cycle(f, a, b):
    """The mechanism the pre-registration rests on: on a pure power the similar successive
    undershoots are replaced by a longer step, and the run gets cheaper."""
    _x, n_off, _ = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3",
                             guard="project")
    _x, n_on, _ = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3",
                            guard="project", extrapolate="steps")
    assert n_on < n_off, (n_on, n_off)


def test_on_a_budget_bound_run_extrapolation_changes_the_path_not_the_count():
    """MEASURED, pinned as measured: on 6(x-2)^5 over [0, 10] both runs cost 90 evaluations
    against a projection guarantee of 2*44 + 3 = 91 - the budget schedule, not the step,
    decides when the run ends.  The paths differ from the fourth evaluation on (87
    positions), so the mode is active; its gain is absorbed by the safeguard here.  The
    expectation written first (strictly fewer evaluations) was wrong for this problem."""
    f, a, b = FLAT[0]
    ev = {}
    for mode in ("none", "steps"):
        seen = []

        def g(t, _s=seen):
            _s.append(t)
            return f(t)
        _x, n, ok = pure_root(g, a, b, XTOL, MAXFEV, similarity="ratio", scale="derived3",
                              guard="project", extrapolate=mode)
        assert ok
        ev[mode] = (n, seen)
    assert ev["steps"][0] <= ev["none"][0]
    assert ev["steps"][1] != ev["none"][1]


def test_unknown_mode_is_refused():
    with pytest.raises(ValueError):
        pure_root(lambda t: t, -1.0, 1.0, XTOL, MAXFEV, extrapolate="aitken")