import numpy as np
import pytest

import certified_robustness as cr
import diffusion_memorization as dm
import double_descent as dd
import high_dim as hd
import randomized_response as rr


# ---------------------------------------------------------------- high_dim
@pytest.mark.parametrize("d", [5, 50, 500])
def test_cosine_std_is_one_over_sqrt_d(d):
    c = hd.cosine_samples(d, n_pairs=20000, seed=d)
    assert abs(c.mean()) < 4 / np.sqrt(20000 * d)
    assert c.std() == pytest.approx(1 / np.sqrt(d), rel=0.03)


def test_shell_fraction_formula_matches_monte_carlo():
    for d in (2, 3, 4):
        mc, _ = hd.shell_fraction_mc(d, 0.2, n=100000, seed=d)
        assert mc == pytest.approx(hd.shell_fraction(d, 0.2), abs=0.01)


def test_pair_violation_probability_matches_simulation():
    d = 300
    c = hd.cosine_samples(d, n_pairs=40000, seed=1)
    assert np.mean(np.abs(c) >= 0.1) == pytest.approx(hd.pair_violation_prob(d), abs=0.01)


def test_greedy_packing_valid_and_grows():
    assert hd.greedy_packing(50, attempts=1000) < hd.greedy_packing(400, attempts=1000)


# ---------------------------------------------------------- double descent
def test_double_descent_peak_and_ridge():
    x, y = dd.make_data()
    grid = [8, 12, 1000]
    m0, _, _ = dd.test_error_curve(x, y, p_grid=grid, draws=30)
    m1, _, _ = dd.test_error_curve(x, y, p_grid=grid, ridge=1e-3, draws=30)
    assert m0[1] > 5 * m0[0] and m0[1] > 5 * m0[2]   # peak at p = n
    assert m1[1] < 2 * m1[0]                          # ridge removes it


def test_min_norm_interpolates_when_overparameterised():
    x, y = dd.make_data()
    rng = np.random.default_rng(0)
    w, b = rng.normal(0, dd.SCALE, 200), rng.uniform(0, 2 * np.pi, 200)
    F = dd.features(x, w, b)
    np.testing.assert_allclose(F @ dd.fit(F, y), y, atol=1e-6)


# ------------------------------------------------------ randomized response
def test_warner_estimator_unbiased_and_variance_formula():
    pi, p, n = 0.3, 0.75, 500
    means, sds = rr.simulate_spread(pi, [n], p, reps=20000, seed=0)
    assert means[0] == pytest.approx(pi, abs=0.002)
    assert sds[0] == pytest.approx(rr.theoretical_sd(pi, n, p), rel=0.03)


def test_warner_individual_level_and_epsilon():
    rng = np.random.default_rng(1)
    truth = rng.random(200000) < 0.2
    ans = rr.respond(truth, 0.8, rng)
    est, se = rr.estimate(ans.sum(), len(ans), 0.8)
    assert abs(est - truth.mean()) < 4 * se
    # P(yes | in A) / P(yes | not in A) = p/(1-p) = e^eps
    ratio = ans[truth].mean() / ans[~truth].mean()
    assert np.log(ratio) == pytest.approx(rr.epsilon(0.8), abs=0.05)


def test_ci_coverage():
    rng = np.random.default_rng(2)
    pi, p, n = 0.4, 0.75, 400
    lam = p * pi + (1 - p) * (1 - pi)
    hits = [rr.confidence_interval(k, n, p)[0] <= pi <= rr.confidence_interval(k, n, p)[1]
            for k in rng.binomial(n, lam, 400)]
    assert np.mean(hits) >= 0.93


def test_poll_cli(tmp_path, capsys):
    f = tmp_path / "poll.csv"
    assert rr.main(["make-example", str(f), "--n", "300", "--pi", "0.5"]) == 0
    assert rr.main(["poll", str(f), "--column", "answer"]) == 0
    assert "debiased estimate" in capsys.readouterr().out


# ------------------------------------------------------------- diffusion
def test_exact_score_matches_finite_difference():
    rng = np.random.default_rng(0)
    train = dm.sample_ring(5, rng)
    x = rng.standard_normal((3, 2))
    sigma = 0.3

    def logp(z):
        d2 = ((z[:, None] - train[None]) ** 2).sum(-1)
        return np.log(np.exp(-d2 / (2 * sigma**2)).sum(1))

    h = 1e-6
    num = np.stack([(logp(x + h * e) - logp(x - h * e)) / (2 * h) for e in np.eye(2)], 1)
    np.testing.assert_allclose(dm.exact_score(x, sigma, train), num, rtol=1e-5, atol=1e-5)


def test_exact_score_sampler_memorises_blur_does_not():
    rng = np.random.default_rng(0)
    train = dm.sample_ring(20, rng)
    s_exact = dm.reverse_sample(dm.exact_score, 300, rng, 200, train)
    s_blur = dm.reverse_sample(dm.blurred_score(0.15), 300, rng, 200, train)
    assert np.mean(dm.nn_dist(s_exact, train) < 0.03) > 0.9
    assert np.mean(dm.nn_dist(s_blur, train) < 0.03) < 0.2


# ---------------------------------------------------- certified robustness
@pytest.fixture(scope="module")
def models():
    return cr.build_models(quick=True)


def test_lipschitz_bound_is_valid(models):
    _, _, lip = models
    rng = np.random.default_rng(0)
    a, b = rng.uniform(-2, 2, (5000, 2)), rng.uniform(-2, 2, (5000, 2))
    ratio = np.abs(lip(a) - lip(b)) / np.linalg.norm(a - b, axis=1)
    assert ratio.max() <= lip.lipschitz_bound() + 1e-9 <= 1 + 1e-9


def test_certified_radius_never_violated_by_attack(models):
    (_, _, Xte, yte), wiggly, lip = models
    for net in (lip, wiggly):
        correct = yte * net(Xte) > 0
        rad = cr.lipschitz_radius(net, Xte)
        m, xadv = cr.pgd_attack(net, Xte, yte, 0.999 * rad, steps=40, restarts=2)
        assert np.all(m[correct] > 0)
        assert np.all(np.linalg.norm(xadv - Xte, axis=1) <= rad + 1e-9)


def test_attack_works_beyond_certificate(models):
    (_, _, Xte, yte), wiggly, _ = models
    m, _ = cr.pgd_attack(wiggly, Xte, yte, 0.5, steps=40, restarts=1)
    assert np.mean(m < 0) > 0.3  # big perturbations do flip many points


def test_randomized_smoothing_certificate(models):
    (_, _, Xte, yte), wiggly, _ = models
    X = Xte[:20]
    preds, radii = cr.smoothed_certify(wiggly, X, 0.25, n=2000, seed=0)
    assert np.all(radii >= 0) and np.all((radii > 0) == (preds != 0))
    assert np.all(radii < 0.25 * 3.3)  # Phi^-1 of a CP bound with n=2000 is < 3.3
    # Monte Carlo check of the guarantee at a point on the certified sphere
    rng = np.random.default_rng(1)
    for x, c, r in zip(X, preds, radii):
        if r > 0:
            u = rng.standard_normal(2)
            x2 = x + 0.99 * r * u / np.linalg.norm(u)
            votes = np.sign(wiggly(x2 + 0.25 * rng.standard_normal((4000, 2))))
            assert np.mean(votes == c) > 0.5