"""The bounds: against scipy, and by simulation (do they really fail at most delta of the time?)."""
import math

import numpy as np
import pytest
from scipy import stats

from router import bounds


def test_beta_cdf_matches_scipy():
    for a, b in [(1, 1), (3, 7), (40, 2), (120, 631), (700, 51)]:
        for x in np.linspace(0.01, 0.99, 25):
            assert abs(bounds.beta_cdf(x, a, b) - stats.beta.cdf(x, a, b)) < 1e-10


def test_clopper_pearson_matches_scipy():
    for n in (10, 100, 750):
        for k in (1, n // 3, n // 2, n - 1):
            for d in (0.1, 0.05, 0.01):
                ours = bounds.clopper_pearson_lower(k, n, d)
                ref = stats.beta.ppf(d, k, n - k + 1)
                assert abs(ours - ref) < 1e-9, (k, n, d, ours, ref)
    assert bounds.clopper_pearson_lower(0, 50, 0.05) == 0.0
    assert abs(bounds.clopper_pearson_lower(50, 50, 0.05) - 0.05 ** (1 / 50)) < 1e-12


@pytest.mark.parametrize("p", [0.6, 0.9, 0.97])
def test_clopper_pearson_coverage_by_simulation(p):
    """The lower bound is below the true p in at least 1 - delta of samples (one-sided, exact)."""
    rng = np.random.default_rng(0)
    n, delta, sims = 200, 0.05, 4000
    ks = rng.binomial(n, p, size=sims)
    cache = {}
    fails = 0
    for k in ks:
        if k not in cache:
            cache[k] = bounds.clopper_pearson_lower(int(k), n, delta)
        fails += cache[k] > p
    rate = fails / sims
    assert rate <= delta + 3 * math.sqrt(delta * (1 - delta) / sims), rate


def test_hoeffding_coverage_by_simulation():
    rng = np.random.default_rng(1)
    n, delta, sims = 150, 0.05, 3000
    scores = rng.beta(5, 2, size=(sims, n))       # scores in [0, 1], true mean 5/7
    lb = scores.mean(axis=1) - bounds.hoeffding_radius(n, delta, 1.0)
    assert (lb > 5 / 7).mean() <= delta                # Hoeffding is conservative: well under delta


def test_required_rate_is_the_smallest_certifiable_rate():
    for n in (100, 750, 3000):
        for t in (0.8, 0.9, 0.95):
            r = bounds.required_rate(t, n, 0.05, True)
            k = round(r * n)
            assert bounds.quality_lower_bound(k, n, 0.05, True) >= t
            assert bounds.quality_lower_bound(k - 1, n, 0.05, True) < t
    assert bounds.required_rate(0.99, 20, 0.05, True) is None   # 20 rows can never certify 99%
    r = bounds.required_rate(0.9, 500, 0.05, False)
    assert abs(r - (0.9 + math.sqrt(math.log(20) / 1000))) < 1e-12


def test_rows_to_certify():
    assert bounds.rows_to_certify(0.85, 0.9, 0.05, True) is None   # below target: hopeless
    n = bounds.rows_to_certify(0.92, 0.9, 0.05, True)
    assert bounds.quality_lower_bound(round(0.92 * n), n, 0.05, True) >= 0.9
    assert bounds.quality_lower_bound(round(0.92 * (n - 1)), n - 1, 0.05, True) < 0.9 or n == 1
    assert bounds.rows_to_certify(0.9000001, 0.9, 0.05, True) > 10_000_000