"""The credit safeguard: the budget earned instead of scheduled.

The scheduled projection needs a multiplier, and the argument for 2 was
circular.  This one banks the surplus a step delivers beyond a halving and
lends only what has been banked, so the only constant left is the endowment,
and one bit is the smallest integer for which the rule is not degenerate.

What must hold, and is asserted here: the stronger guarantee is real, the
endowment is minimal, and the cost of the stronger guarantee is stated rather
than hidden.
"""
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.pure import _Safeguard, pure_root

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


def problems():
    out = []
    for d in _APS_TESTS_DICTS:
        a, b = d["bracket"]
        args = d.get("args", ())
        out.append(((lambda ff, aa: lambda x: ff(x, *aa))(d["f"], args),
                    float(a), float(b), float(d["root"])))
    for d in _CHANDRUPATLA_TESTS_DICTS:
        a, b = d["bracket"]
        out.append(((lambda ff: lambda x: ff(x))(d["f"]),
                    float(a), float(b), float(d["root"])))
    return out


ALL = problems()


def sweep(guard):
    total = solved = over = 0
    for f, a, b, root in ALL:
        try:
            x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                                   scale="derived3", guard=guard)
        except Exception:
            x, n, conv = float("nan"), MAXFEV, False
        ok = bool(conv and math.isfinite(x) and abs(x - root) <= ACC)
        total += n if ok else MAXFEV
        solved += ok
        if ok and n > math.ceil(math.log2(abs(b - a) / XTOL)) + 4:
            over += 1
    return total, solved, over


def test_the_credit_guard_meets_the_near_minimal_bound_everywhere():
    """bisection + one step, on all 199 problems.  The scheduled projection
    is allowed twice that and uses it on 19 of them."""
    _t, solved, over = sweep("credit")
    assert solved == 193, solved
    assert over == 0, over


def test_the_scheduled_projection_really_does_use_its_larger_budget():
    """If this ever drops to zero the two guards have converged and the
    multiplier stopped mattering - which would be worth knowing."""
    _t, solved, over = sweep("project")
    assert solved == 193, solved
    assert over > 10, over


def test_the_stronger_guarantee_costs_and_the_cost_is_stated():
    """No free lunch: the paper must not claim the credit rule is both
    cheaper and safer."""
    t_credit, _s, _o = sweep("credit")
    t_sched, _s2, _o2 = sweep("project")
    assert t_credit > t_sched, (t_credit, t_sched)
    assert (t_credit - t_sched) / t_sched < 0.15, (t_credit, t_sched)


def test_a_zero_endowment_is_degenerate():
    """Why one bit is the minimum: with no endowment the allowance on the
    first step is exactly half the bracket, which forces the midpoint."""
    sg = _Safeguard("credit", 1.0, XTOL)
    assert sg.credit == 1.0
    sg.credit = 0.0
    assert sg.project_credit(0.9, 0.0, 1.0) == pytest.approx(0.5)
    # with the endowment the same step is left alone
    sg.credit = 1.0
    assert sg.project_credit(0.9, 0.0, 1.0) == pytest.approx(0.9)


def test_credit_never_goes_negative_and_banks_surplus():
    sg = _Safeguard("credit", 1.0, XTOL)
    sg.credit = 0.0
    sg.note_shrink(1.0, 1.0)          # no shrink at all
    assert sg.credit == 0.0           # clamped, never a debt
    sg.note_shrink(1.0, 0.125)        # three bits of shrink
    assert sg.credit == pytest.approx(2.0)


@pytest.mark.parametrize("c0,steps_over", [(0.25, 1), (0.5, 1), (1.0, 1),
                                           (2.0, 2), (3.0, 3)])
def test_the_guarantee_is_bisection_plus_ceil_of_the_endowment(c0, steps_over):
    """The relation that makes the endowment determined rather than chosen."""
    worst = 0
    for f, a, b, root in ALL:
        try:
            x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                                   scale="derived3", guard="credit",
                                   endowment=c0)
        except Exception:
            continue
        if conv and math.isfinite(x) and abs(x - root) <= ACC:
            half = math.ceil(math.log2(abs(b - a) / XTOL))
            worst = max(worst, (n - 2) - half)
    assert worst == steps_over, (c0, worst, steps_over)


def test_one_bit_is_the_cheapest_endowment_that_promises_one_step():
    """Why C0 = 1 is forced and not picked: every C0 in (0,1] gives the same
    one-step guarantee, and 1 is the most permissive of them, hence the
    cheapest.  If a smaller endowment ever became cheaper, the argument that
    the constant is determined would collapse."""
    def cost(c0):
        t = 0
        for f, a, b, root in ALL:
            try:
                x, n, conv = pure_root(f, a, b, XTOL, MAXFEV,
                                       similarity="ratio", scale="derived3",
                                       guard="credit", endowment=c0)
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            ok = conv and math.isfinite(x) and abs(x - root) <= ACC
            t += n if ok else MAXFEV
        return t

    c_one = cost(1.0)
    for smaller in (0.1, 0.25, 0.5):
        assert cost(smaller) > c_one, smaller


def test_the_scheduled_multiplier_has_no_such_determination():
    """The contrast that decides which safeguard the paper should propose.
    The multiplicative family has no smallest non-trivial member: 1 is
    degenerate and everything above it is admissible with a different cost, so
    the multiplier is genuinely free."""
    def cost(fac):
        t = 0
        for f, a, b, root in ALL:
            try:
                x, n, conv = pure_root(f, a, b, XTOL, MAXFEV,
                                       similarity="ratio", scale="derived3",
                                       guard="project", budget_factor=fac)
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            ok = conv and math.isfinite(x) and abs(x - root) <= ACC
            t += n if ok else MAXFEV
        return t

    degenerate = cost(1.0)
    assert degenerate > cost(2.0)          # 1 is degenerate
    assert cost(1.25) < cost(2.0)          # and 2 is not the best either


def test_the_endowment_cannot_be_earned_and_so_cannot_be_computed():
    """Why C0 is a licence rather than a quantity to be derived from the run.

    Credit is earned only by a step that beats a halving, and such a step is
    permitted only when credit already exists.  Starting from zero the
    mechanism is locked: the allowance is exactly half the bracket, which
    forces the midpoint, which delivers exactly one bit, which banks nothing.
    The method is then bisection and stays bisection - so no amount of running
    can produce the evidence that would justify a larger endowment."""
    sg = _Safeguard("credit", 1.0, XTOL)
    sg.credit = 0.0
    probe = None
    for _ in range(30):
        probe = sg.project_credit(0.99, 0.0, 1.0)   # a bold step is proposed
        assert probe == pytest.approx(0.5)          # ...and always refused
        sg.note_shrink(1.0, max(probe, 1.0 - probe))
        assert sg.credit == 0.0                     # nothing is ever banked

    # end to end: zero endowment reproduces bisection's cost
    from simroot.baselines import bisection
    t_zero = t_bis = 0
    for f, a, b, root in ALL:
        for meth, acc in ((lambda: pure_root(f, a, b, XTOL, MAXFEV,
                                             similarity="ratio",
                                             scale="derived3", guard="credit",
                                             endowment=0.0), "z"),
                          (lambda: bisection(f, a, b, XTOL, MAXFEV), "b")):
            try:
                x, n, conv = meth()
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            ok = conv and math.isfinite(x) and abs(x - root) <= ACC
            if acc == "z":
                t_zero += n if ok else MAXFEV
            else:
                t_bis += n if ok else MAXFEV
    assert abs(t_zero - t_bis) / t_bis < 0.01, (t_zero, t_bis)


def test_a_large_credit_cannot_overflow_the_exponent():
    """Found by the algorithm audit.  The allowance is w * 2^(C-1) and C
    accumulates the log of every shrink without bound, so the exponentiation
    overflowed a double at C ~ 1025 and the OverflowError escaped - only the
    budget exception is caught in pure_root.  It is reachable: sixty
    nine-decade steps reach C = 1012, and a bracket of width 1e300 with a
    tolerance of 1e-8 has about 1023 bits of shrink available."""
    sg = _Safeguard("credit", 1.0, XTOL)
    for c in (1.0, 100.0, 1024.0, 1100.0, 1e6):
        sg.credit = c
        assert sg.project_credit(0.9, 0.0, 1.0) == 0.9, c
    # and the accumulation that gets there must not raise either
    sg.credit = 1.0
    w = 1.0
    for _ in range(60):
        sg.note_shrink(w, w * 1e-9)
        w *= 1e-9
    assert sg.credit > 1000.0
    assert sg.project_credit(0.9, 0.0, 1.0) == 0.9


def test_a_very_wide_bracket_is_solved_end_to_end():
    """The same defect, reached through the public entry point."""
    x, n, conv = pure_root(lambda t: t - 1.0, -1e150, 1e150, 1e-12, 4000,
                           similarity="ratio", scale="derived3",
                           guard="credit")
    assert conv and abs(x - 1.0) < 1e-6, (x, n, conv)


def test_credit_needs_no_tolerance_at_all():
    """Unlike the scheduled version it reads only widths, so there is no
    n_max, no eps, and nothing fixed in advance from the initial bracket."""
    f = lambda t: t * t - 2.0
    x, _n, conv = pure_root(f, 0.0, 2.0, XTOL, MAXFEV, similarity="ratio",
                            scale="derived3", guard="credit")
    assert conv and abs(x - math.sqrt(2.0)) <= 1e-8