"""Tests for the projection safeguard.

This is the guard the paper will propose, so its guarantee has to be checked
where it matters - on every problem of both benchmark sets, not on a sample -
and the order it was built to preserve has to be shown to survive.
"""
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 PURE_SIMS, pure_root
from stage5_three_point import PROBLEMS, order_of

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

PROJECTING = ("project", "safe")


def _all_problems():
    out = []
    for d in _APS_TESTS_DICTS:
        a, b = d["bracket"]
        args = d.get("args", ())
        out.append((d["ID"], (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((d["ID"], (lambda ff: lambda x: ff(x))(d["f"]),
                    float(a), float(b), float(d["root"])))
    return out


ALL = _all_problems()


def bound(a, b):
    """The guarantee the paper may state.

    The budget gives w_{j+1} <= 2^(n_max - j) * eps with n_max = 2*half, so in
    exact arithmetic the width reaches xtol = 2*eps at step j = n_max - 1 and
    the run costs n_max steps plus the two endpoint evaluations.

    In floating point the bracket can land exactly ON the budget boundary -
    the instrumented invariant in stage6b shows the worst ratio at 1.000000 -
    and the loop condition `w > xtol` then demands one further step.  The
    honest statement is therefore n_max + 1 steps, i.e.

        2 * ceil(log2(w / xtol)) + 3   evaluations,

    which is still at most twice bisection plus one step.  The `+1` is stated
    rather than hidden: it is a property of the termination test at the
    boundary, not slack that could absorb a real violation."""
    return 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 3


@pytest.mark.parametrize("guard", PROJECTING)
@pytest.mark.parametrize("scale", ["derived", "derived3"])
def test_projection_respects_the_bound_on_all_199_problems(guard, scale):
    """The claim the paper makes about the projection guard.  A single
    violation invalidates it, so every problem is checked."""
    bad = []
    for pid, f, a, b, root in ALL:
        x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                               scale=scale, guard=guard)
        if conv and abs(x - root) <= ACC and n > bound(a, b):
            bad.append((pid, n, bound(a, b)))
    assert not bad, f"{guard}/{scale} exceeded the bound on {bad[:5]}"


@pytest.mark.parametrize("guard", PROJECTING)
def test_projection_solves_as_much_as_the_unguarded_rule(guard):
    """A safeguard that bought its guarantee by failing more often would be
    no bargain."""
    solved_guard = solved_none = 0
    for _pid, f, a, b, root in ALL:
        for g, acc in ((guard, "g"), ("none", "n")):
            x, _n, conv = pure_root(f, a, b, XTOL, MAXFEV, similarity="ratio",
                                    scale="derived3", guard=g)
            ok = conv and abs(x - root) <= ACC
            if acc == "g":
                solved_guard += ok
            else:
                solved_none += ok
    assert solved_guard >= solved_none, (solved_guard, solved_none)


def test_projection_keeps_the_order_the_other_guards_destroy():
    """The whole point: `bracket` and `both` collapse the three-point rule
    from order ~1.94 to ~1.00; `project` and `safe` must not."""
    med = {}
    for guard in ("none", "bracket", "both", "project", "safe", "step"):
        ps = [order_of((lambda g: lambda f_, a_, b_, xt, mx: pure_root(
            f_, a_, b_, XTOL, MAXFEV, scale="derived3", guard=g))(guard),
            f, a, b, root) for _n, f, a, b, root in PROBLEMS]
        ps = [p for p in ps if math.isfinite(p) and 0.5 < p < 5.0]
        med[guard] = sorted(ps)[len(ps) // 2] if ps else float("nan")
    assert med["bracket"] < 1.2, med          # the known collapse
    assert med["both"] < 1.2, med
    assert med["project"] > 1.7, med          # the fix
    assert med["safe"] > 1.7, med
    assert med["safe"] >= med["bracket"] + 0.5, med


def test_zero_slack_would_be_degenerate():
    """Guards against someone 'tightening' the budget back to the exact
    bisection bound: with no slack the projection radius on a unit bracket is
    a twentieth of it from the very first step, and the method becomes
    bisection.  The doubling is load-bearing."""
    from simroot.pure import _Safeguard
    sg = _Safeguard("project", 1.0, XTOL)
    r_doubled = 2.0 ** (sg.n_half - 0) * sg.eps - 0.5
    half = sg.n_half // 2
    r_zero_slack = 2.0 ** half * sg.eps - 0.5
    assert r_zero_slack < 0.06, r_zero_slack
    assert r_doubled > 1e6, r_doubled


@pytest.mark.parametrize("sim", sorted(PURE_SIMS))
def test_every_similarity_works_under_the_projection(sim):
    f = lambda t: t * t - 2.0
    x, _n, conv = pure_root(f, 0.0, 2.0, XTOL, MAXFEV, similarity=sim,
                            scale="derived3", guard="safe")
    assert conv and abs(x - math.sqrt(2.0)) <= 1e-8


def test_projection_needs_a_tolerance_and_a_real_bracket():
    with pytest.raises(ValueError):
        from simroot.pure import _Safeguard
        _Safeguard("project", 1.0, 0.0)
    with pytest.raises(ValueError):
        from simroot.pure import _Safeguard
        _Safeguard("safe", 0.0, XTOL)