"""Exact exponential-model identities (Theorem 3 ingredients) checked numerically on random inputs."""
import numpy as np
from scipy import stats
from pamle.jackknife import vandermonde_weights, leading_bias_coefficient


def test_gamma_inverse_moments():
    rng = np.random.default_rng(1)
    for n in (5, 17, 64):
        theta0 = 2.5
        S = rng.gamma(n, 1 / theta0, size=400000)
        th = n / S
        assert abs(th.mean() - theta0 * n / (n - 1)) < 0.01 * theta0
        assert abs(th.var() - theta0 ** 2 * n ** 2 / ((n - 1) ** 2 * (n - 2))) < 0.05 * theta0 ** 2 * n ** 2 / ((n - 1) ** 2 * (n - 2))


def test_newton_identity():
    rng = np.random.default_rng(2)
    for _ in range(1000):
        x = rng.uniform(0.1, 5); xbar = rng.uniform(0.1, 5); th = 1 / xbar
        g = 1 / x - xbar; H = -1 / x ** 2
        newton = x - g / H
        assert abs(newton - (2 * x - x * x * xbar)) < 1e-12
        assert abs((newton - th) + (x - th) ** 2 / th) < 1e-12


def test_newton_basin():
    th = 1.3
    for x in (0.01, 0.5, 1.0, 2.0, 2.59):
        e = x - th
        for _ in range(60):
            e = -e * e / th
        assert abs(e) < 1e-6
    for x in (-0.1, 2.61, 3.0):
        e = x - th
        for _ in range(12):
            e = -e * e / th
        assert abs(e) > 1
    e = 0.0 + (2 * th - th)                      # x = 2 theta_hat -> maps to 0 -> stays at 0 (error -theta_hat)
    e = -e * e / th
    assert abs(e + th) < 1e-12


def test_exact_bias_of_split_estimators():
    """E[J_r] - theta0 = theta0 * sum_j w_j s_j/(m - s_j) exactly, leading term beta_r theta0 / m^r."""
    theta0 = 1.0
    for s in [(1, 2), (1, 2, 4), (1, 2, 4, 8)]:
        w = vandermonde_weights(s); r = len(s)
        for m in (64, 256, 1024):
            bias = sum(wj * theta0 * m / (m - sj) for wj, sj in zip(w, s)) - theta0
            if s == (1, 2):
                assert abs(bias + 2 * theta0 / ((m - 1) * (m - 2))) < 1e-14
            if s == (1, 2, 4):
                assert abs(bias - 8 * theta0 / ((m - 1) * (m - 2) * (m - 4))) < 1e-14
            assert abs(bias * m ** r / (float(leading_bias_coefficient(s)) * theta0) - 1) < 20 / m


def test_boundary_constant_numeric():
    """Deterministic version of Theorem 3: plug e0 = beta_r theta0 c^r N^{-1/2^(q+1)} into the q-fold Newton identity."""
    theta0, c = 1.7, 0.8
    for s, q in [((1, 2), 1), ((1, 2, 4), 1), ((1, 2), 2), ((1, 2, 4, 8), 1)]:
        r = len(s); beta = float(leading_bias_coefficient(s))
        for N in (1e12, 1e16):
            alpha = 1 - 1 / (2 ** (q + 1) * r); m = N ** (1 - alpha) / c
            e = beta * theta0 * m ** (-r)
            for _ in range(q):
                e = -e * e / theta0
            lim = -theta0 * (beta * c ** r) ** (2 ** q)
            assert abs(np.sqrt(N) * e - lim) < 1e-8 * abs(lim)
        assert abs((-theta0 * ((-2) * c ** 2) ** 2) - (-4 * theta0 * c ** 4)) < 1e-14


def test_exact_variance_of_half_block_jackknife():
    """Var(J2) = theta0^2 m^2 (m^2 - 6m + 17)/((m-4)(m-2)^2(m-1)^2) (SymPy, theory/symbolic_checks.py S6) vs simulation."""
    rng = np.random.default_rng(5); m = 32; B = 400000
    halves = rng.gamma(m // 2, 1.0, size=(B, 2))
    full = m / halves.sum(axis=1); half = (m / 2) / halves
    J2 = 2 * full - half.mean(axis=1)
    exact = m ** 2 * (m ** 2 - 6 * m + 17) / ((m - 4) * (m - 2) ** 2 * (m - 1) ** 2)
    assert abs(J2.var() / exact - 1) < 0.02
    assert abs(J2.mean() - (1 - 2 / ((m - 1) * (m - 2)))) < 0.002