#!/usr/bin/env python3
"""Symbolic verification of the algebraic identities used in WP28 (SymPy).

Every check raises AssertionError on failure. Run:  python3 theory/symbolic_checks.py
Items (numbering follows theory/THEOREM_AUDIT.md):
  S1  Vandermonde extrapolation weights for several split patterns; moment constraints.
  S2  Exponential(rate) model: Gamma inverse moments, E[1/Xbar_n], Var, remainder R_n moments.
  S3  Exact bias of the split estimators J_r in the exponential model and its leading coefficient beta_r.
  S4  Exact Newton identity N(x) = 2x - x^2 Xbar and the error identity N(x) - theta_hat = -(x - theta_hat)^2 / theta_hat.
  S5  q-fold Newton error identity and the boundary limit -theta0 (beta_r c^r)^(2^q).
  S6  Exact variance of the half-block jackknife J_2 in the exponential model (nested dependence handled exactly).
  S7  Newton basin of attraction in the exponential model: convergence iff 0 < x < 2 theta_hat.
  S8  Remainder-amplification constant A_r = sum |w_j| s_j and the sufficient exponent table alpha(r,q).
"""
import sympy as sp
from fractions import Fraction

def vandermonde_weights(s):
    """Weights w solving sum w_j = 1, sum w_j s_j^a = 0 (a=1..r-1); exact rationals."""
    r = len(s)
    M = sp.Matrix([[sp.Integer(sj)**a for sj in s] for a in range(r)])
    rhs = sp.Matrix([1] + [0]*(r-1))
    return list(M.LUsolve(rhs))

# ---------- S1: weights ----------
checks = {}
for s in [(1,2), (1,2,4), (1,2,4,8), (1,2,3), (1,3,9), (1,2,4,8,16)]:
    w = vandermonde_weights(s)
    for a in range(len(s)):
        tot = sum(wj * sp.Integer(sj)**a for wj, sj in zip(w, s))
        assert tot == (1 if a == 0 else 0), (s, a, tot)
    checks[s] = w
assert checks[(1,2)] == [2, -1]
assert checks[(1,2,4)] == [sp.Rational(8,3), -2, sp.Rational(1,3)]
print("S1 OK  weights:", {k: [str(x) for x in v] for k, v in checks.items()})

# ---------- S2: exponential model moments ----------
n, th = sp.symbols('n theta0', positive=True)
# S_n ~ Gamma(n, rate theta0): E[S^-k] = theta0^k * Gamma(n-k)/Gamma(n)
def Einv(k):  # E[(theta0*Xbar_n)^(-k)] = E[(S_n/n)^-k] * theta0^-k ... compute E[Y^-k] with Y = theta0*Xbar_n, nY ~ Gamma(n,1)
    # E[(nY)^-k] = Gamma(n-k)/Gamma(n)  =>  E[Y^-k] = n^k Gamma(n-k)/Gamma(n)
    return sp.simplify(n**k * sp.gamma(n-k)/sp.gamma(n))
EY1 = sp.simplify(sp.gammasimp(Einv(1)))
EY2 = sp.simplify(sp.gammasimp(Einv(2)))
assert sp.simplify(EY1 - n/(n-1)) == 0
assert sp.simplify(EY2 - n**2/((n-1)*(n-2))) == 0
# E[theta_hat_n] = theta0 * n/(n-1); Var = theta0^2 n^2/((n-1)^2 (n-2))
var = sp.simplify(th**2*(EY2 - EY1**2))
assert sp.simplify(var - th**2*n**2/((n-1)**2*(n-2))) == 0
# remainder R_n = theta0 (Y^-1 - 2 + Y),  E Y = 1, E Y^2 = 1 + 1/n
ER = sp.simplify(th*(EY1 - 2 + 1))
assert sp.simplify(ER - th/(n-1)) == 0
ER2 = sp.simplify(th**2*(EY2 - 4*EY1 + 6 - 4*1 + (1 + 1/n)))
assert sp.simplify(ER2 - th**2*(3*n+2)/(n*(n-1)*(n-2))) == 0
# series: E R_n = theta0 (1/n + 1/n^2 + ...) => b_a = theta0 for all a
ser = sp.series(th/(n-1), n, sp.oo, 4).removeO()
assert sp.simplify(ser - th*(1/n + 1/n**2 + 1/n**3)) == 0
# uniform constant: sup_{n>=3} n^2 E R_n^2 / theta0^2
f = sp.lambdify(n, n**2*(3*n+2)/(n*(n-1)*(n-2)))
Csup = max(f(k) for k in range(3, 2000))
assert abs(Csup - 16.5) < 1e-12
print("S2 OK  E[theta_hat_n]=theta0 n/(n-1), Var=theta0^2 n^2/((n-1)^2(n-2)), E R_n=theta0/(n-1), E R_n^2=theta0^2(3n+2)/(n(n-1)(n-2)), sup n^2 E R_n^2/theta0^2 = 16.5 (n=3)")

# ---------- S3: exact bias of J_r and leading coefficient ----------
m = sp.symbols('m', positive=True)
def exact_bias(s):
    w = vandermonde_weights(s)
    # sub-block size m/s_j, E[theta_hat] = theta0 (m/s_j)/((m/s_j)-1) = theta0 m/(m - s_j)
    return sp.simplify(sum(wj*th*m/(m-sj) for wj, sj in zip(w, s)) - th)
b2 = exact_bias((1,2)); b3 = exact_bias((1,2,4))
assert sp.simplify(b2 + 2*th/((m-1)*(m-2))) == 0
assert sp.simplify(b3 - 8*th/((m-1)*(m-2)*(m-4))) == 0
betas = {}
for s in [(1,2),(1,2,4),(1,2,4,8),(1,2,3),(1,3,9)]:
    w = vandermonde_weights(s); r = len(s)
    beta = sum(wj*sp.Integer(sj)**r for wj, sj in zip(w, s))
    lead = sp.limit(exact_bias(s)*m**r, m, sp.oo)
    assert sp.simplify(lead - th*beta) == 0, (s, lead, beta)
    betas[s] = beta
assert betas[(1,2)] == -2 and betas[(1,2,4)] == 8 and betas[(1,2,4,8)] == -64
print("S3 OK  exact biases: J2 = -2 theta0/((m-1)(m-2)), J3 = 8 theta0/((m-1)(m-2)(m-4)); beta_r:", {k: str(v) for k, v in betas.items()})

# ---------- S4: Newton identity ----------
x, xbar = sp.symbols('x xbar', positive=True)
L = sp.log(x) - x*xbar              # average log-likelihood of Exp(rate x)
g = sp.diff(L, x); H = sp.diff(L, x, 2)
Nx = sp.simplify(x - g/H)
assert sp.simplify(Nx - (2*x - x**2*xbar)) == 0
that = 1/xbar
assert sp.simplify((Nx - that) + (x - that)**2/that) == 0
print("S4 OK  N(x) = 2x - x^2 Xbar,  N(x) - theta_hat = -(x - theta_hat)^2/theta_hat")

# ---------- S5: q-fold identity and boundary limit ----------
e0, thh, c, N = sp.symbols('e0 thetahat c N', positive=True)
def newton_err(e, q):
    for _ in range(q):
        e = -e**2/thh
    return sp.simplify(e)
for q in range(1, 4):
    assert sp.simplify(newton_err(e0, q) + e0**(2**q)/thh**(2**q-1)) == 0
# boundary: K = c N^alpha, alpha = 1 - 1/(2^(q+1) r)  =>  m = N/K = N^{1/(2^(q+1) r)}/c ;  e0 ~ beta_r theta0 m^{-r}
for s in [(1,2),(1,2,4),(1,2,4,8)]:
    r = len(s); beta = betas[s]
    for q in range(0, 3):
        alpha = 1 - sp.Rational(1, 2**(q+1)*r)
        mm = N**(1-alpha)/c
        e_pilot = beta*th*mm**(-r)
        if q == 0:
            lim = sp.limit(sp.sqrt(N)*e_pilot * N**(sp.Rational(1,2)) / N**(sp.Rational(1,2)), N, sp.oo)  # sqrt(N) e0 -> infinity: use N^{1/2} scaling only for q>=1
        err = newton_err(e_pilot, q).subs(thh, th)
        scaled = sp.simplify(sp.sqrt(N)*err)
        lim = sp.limit(scaled, N, sp.oo)
        expected = -th*(beta*c**r)**(2**q) if q >= 1 else sp.nan
        if q >= 1:
            assert sp.simplify(lim - expected) == 0, (s, q, lim, expected)
        else:
            # q = 0: the scaled bias N^{1/2} m^{-r} = c^r N^{1/2 - 1/2} -> beta theta0 c^r (alpha(r,0) = 1 - 1/(2r) makes sqrt(N) m^{-r} = c^r)
            assert sp.simplify(lim - beta*th*c**r) == 0, (s, q, lim)
assert sp.simplify(-th*((-2)*c**2)**2 + 4*th*c**4) == 0
print("S5 OK  N^q(x) - theta_hat = -(x-theta_hat)^(2^q)/theta_hat^(2^q-1); boundary limit sqrt(N)(theta^(q) - theta_hat) -> -theta0 (beta_r c^r)^(2^q) for q>=1 (r=2,q=1: -4 theta0 c^4) and beta_r theta0 c^r for q=0")

# ---------- S6: exact variance of J2 in the exponential model ----------
# S_A, S_B ~ Gamma(h, 1), h = m/2 (theta0 = 1 w.l.o.g.); S = S_A+S_B ~ Gamma(2h,1), W = S_A/S ~ Beta(h,h) indep. of S.
# theta_full = m/S = 2h/S ; theta_A = h/S_A = h/(S W) ; theta_B = h/(S (1-W)).
# J2 = 2*theta_full - (theta_A+theta_B)/2 = 4h/S - (h/2)(1/W + 1/(1-W))/S
h = sp.symbols('h', positive=True)
def EgammaInv(a, k):   # E[S^-k], S ~ Gamma(a,1)
    return sp.gamma(a-k)/sp.gamma(a)
def EbetaInv(a, b, k, l):  # E[W^-k (1-W)^-l], W ~ Beta(a,b)
    return sp.gamma(a-k)*sp.gamma(b-l)*sp.gamma(a+b)/(sp.gamma(a)*sp.gamma(b)*sp.gamma(a+b-k-l))
# J2 = (1/S) * [4h - (h/2) (1/W + 1/(1-W))] =: (1/S) * G(W)
# E J2 = E[S^-1] * E[G];  E J2^2 = E[S^-2] * E[G^2]
EG  = 4*h - (h/2)*(EbetaInv(h,h,1,0) + EbetaInv(h,h,0,1))
EG2 = 16*h**2 - 2*4*h*(h/2)*(EbetaInv(h,h,1,0)+EbetaInv(h,h,0,1)) + (h**2/4)*(EbetaInv(h,h,2,0) + EbetaInv(h,h,0,2) + 2*EbetaInv(h,h,1,1))
EJ2  = sp.gammasimp(EgammaInv(2*h,1)*EG)
EJ22 = sp.gammasimp(EgammaInv(2*h,2)*EG2)
EJ2_m  = sp.simplify(EJ2.subs(h, m/2))
VarJ2_m = sp.simplify((EJ22 - EJ2**2).subs(h, m/2))
assert sp.simplify(EJ2_m - (1 - 2/((m-1)*(m-2)))) == 0          # matches S3 (theta0 = 1)
# asymptotic expansion of the variance: 1/m + v2/m^2 + ...
serV = sp.series(VarJ2_m, m, sp.oo, 4).removeO()
v1 = sp.limit(VarJ2_m*m, m, sp.oo); v2 = sp.limit((VarJ2_m - 1/m)*m**2, m, sp.oo)
assert v1 == 1
# compare with the full-block MLE variance m^2/((m-1)^2 (m-2)) = 1/m + 3/m^2 + ...
v2_full = sp.limit((m**2/((m-1)**2*(m-2)) - 1/m)*m**2, m, sp.oo)
print("S6 OK  Var(J2) exact =", sp.factor(VarJ2_m), " ~ 1/m + (%s)/m^2 ;  Var(theta_hat_m) ~ 1/m + (%s)/m^2" % (v2, v2_full))
VARJ2_EXACT = VarJ2_m

# ---------- S7: Newton basin ----------
# e_{j+1} = -e_j^2/theta_hat ; |e_{j+1}|/theta_hat = (|e_j|/theta_hat)^2 -> 0 iff |e_0| < theta_hat iff 0 < x < 2 theta_hat
u = sp.symbols('u', positive=True)
assert sp.simplify(sp.Abs(-(u*thh)**2/thh)/thh - u**2) == 0
print("S7 OK  Newton basin: |x - theta_hat| < theta_hat  <=>  0 < x < 2 theta_hat ; x <= 0 or x >= 2 theta_hat diverges (x = 2 theta_hat maps to 0 then stays at 0)")

# ---------- S8: amplification constants and exponent table ----------
def A_const(s):
    w = vandermonde_weights(s); return sum(abs(wj)*sj for wj, sj in zip(w, s))
def W_sum_abs(s):
    w = vandermonde_weights(s); return sum(abs(wj) for wj in w)
tab = []
for s in [(1,), (1,2), (1,2,4), (1,2,4,8), (1,2,3), (1,3,9)]:
    tab.append((s, A_const(s), W_sum_abs(s)))
assert A_const((1,2)) == 4 and A_const((1,2,4)) == 8
alpha = {(r,q): sp.Rational(1) - sp.Rational(1, 2**(q+1)*r) for r in (1,2,3,4) for q in (0,1,2,3)}
assert alpha[(1,0)] == sp.Rational(1,2) and alpha[(1,1)] == sp.Rational(3,4) and alpha[(2,1)] == sp.Rational(7,8) and alpha[(3,1)] == sp.Rational(11,12) and alpha[(1,2)] == sp.Rational(7,8) and alpha[(2,0)] == sp.Rational(3,4) and alpha[(3,0)] == sp.Rational(5,6)
print("S8 OK  A_r = sum|w_j| s_j:", [(s, str(a), str(b)) for s, a, b in tab])
print("       alpha(r,q) = 1 - 1/(2^(q+1) r):", {k: str(v) for k, v in alpha.items()})
print("ALL SYMBOLIC CHECKS PASSED")