#!/usr/bin/env python3
"""
test_336 -- THE LINEARISATION OF F AT THE ROUND S^3: THE FACTORISATION ON TT, THE L^2 IDENTITY AT LEVEL 3, THE
            AUGMENTED SYMBOL, THE COUNT OF TEN AND THE DETERMINANT, AND THE SHELL DIPOLE'S DECOMPOSITION

gamma = |dx|^2/(1 + |x|^2/4)^2 (the unit round metric, stereographic coordinates, centred at the origin).

PRE-REGISTERED:
  P1 (Theorem E(2), as an operator identity) For tensors h that are transverse-traceless to the order that matters at
     the origin (tr h = 0 through degree 3 and div h = 0 through degree 2, as polynomial jets), the O(eps) part of
     F(gamma + eps h) at the origin equals Cod(k)(0) with k = (1/2)(-Lap h + 2 h), Lap the rough Laplacian, for three
     random such jets (exact rationals; the identity is linear, so random jets test it on the whole constrained space).
     Control: for a jet with tr h != 0 the two sides differ.
  P2 (Theorem E(4) at level 3) On S^3 = SU(2) with a left-invariant orthonormal frame, nabla_{E_a} E_b = eps_abc E_c
     (sectional curvature 1, checked), a constant symmetric trace-free h_ab is transverse-traceless with
     -Lap h = 6 h, |nabla h|^2 = 6 |h|^2 and |Cod h|^2 = 18 |h|^2 pointwise; so |Cod h|^2 = 2|nabla h|^2 + 6|h|^2 and,
     with P1, |DF h|^2 = 16 |Cod h|^2 = 288 |h|^2.
  P3 (Theorem E(5), the augmented symbol) For xi != 0 and symmetric H: xi ^ H = 0, H xi = 0 and (tr H) xi = 0 force
     H = 0; and on {H xi = 0, tr H = 0}, |xi ^ H|^2 = 2 |xi|^2 |H|^2 (sum over all a, b, c).
  P4 (Remark 6.1) sum_{n>=3} 2(n^2 - 4) lambda_n^{-s} continued to s = 0 is 10 for lambda_n = n(n^2 - 1) and for
     lambda_n = n^2 (n^2 - 1)^2/2 (continuation checked against direct sums at s = 2, 3); and for lambda_n = n(n^2 - 1)
     the determinant -Z'(0) computed by numerical differentiation of the continuation equals the closed form
     -6[zeta_H'(-2,3) - 4 zeta_H'(0,3)] - sum_{k>=1} (2/k)[zeta_H(2k-2,3) - 4 zeta_H(2k,3)] = -14.0371968...
  P5 (Proposition 6.2(3)) In gamma = dr^2 + sin^2 r (dtheta^2 + sin^2 theta dphi^2), h = 2 lambda cos(theta) dr^2 equals
     2 w gamma + L_eta gamma with w = (lambda + c sin r) cos(theta), eta_r = (c cos r + q0) cos(theta),
     eta_theta = -(c sin r + q0 sin r cos r + lambda sin^2 r) sin(theta), for all constants c, q0 (symbolic).

This script, with test_327 (P1-P2 for the kernel), replaces tests 79, 80, 83 and 98, the TT parts of test_315, and the
closed-range audit script for the paper's verification bundle.

FIRST ATTEMPT: stopped after 30 minutes with no check run. The linearisation was taken by carrying eps symbolically
through an inverse metric whose constant term depends on eps, which produced rational functions of eps. Rewritten to
carry eps as a dual number (truncated to first order after every product); the checks are unchanged.
SECOND RUN: P2-P5 passed; in P1 the three TT trials passed but the control did not differ. The control was badly
chosen, not the identity at fault: its trace 3 x0 x1 and its divergence have vanishing low-order jets at the origin,
which is all the difference of the two sides can see there. Control replaced by h = (x0 + x1 x2) times the identity
plus the same off-diagonal term, whose trace has a nonzero gradient at the origin.
"""
import sys
import random
import itertools
import sympy as sp
import mpmath as mp
from sympy import Rational as R

FAILURES = []


def check(name, ok, detail=""):
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}" + (f"  -- {detail}" if detail else ""))
    if not ok:
        FAILURES.append(name)


X = sp.symbols("x0:3", real=True)
eps = sp.Symbol("eps")
SYM = [(0, 0), (1, 1), (2, 2), (0, 1), (0, 2), (1, 2)]
CODIDX = [(a, b, c) for (a, b) in ((0, 1), (0, 2), (1, 2)) for c in range(3)]
o = {x: 0 for x in X}


def trunc(e, deg):
    e = sp.expand(e)
    if e == 0:
        return sp.Integer(0)
    P = sp.Poly(e, *X)
    return sum((cf * sp.prod([x**m for x, m in zip(X, mon)]) for mon, cf in P.terms() if sum(mon) <= deg), sp.Integer(0))


def tmat(M, deg):
    return M.applyfunc(lambda e: trunc(e, deg))


rr = X[0]**2 + X[1]**2 + X[2]**2
gam = tmat(sp.eye(3) * sp.series(1 / (1 + sp.Symbol("w") / 4)**2, sp.Symbol("w"), 0, 3).removeO().subs(sp.Symbol("w"), rr), 3)
gami = tmat(sp.eye(3) * (1 + rr / 4)**2, 3)
dgam = [[[sp.diff(gam[b, c], X[a]) for c in range(3)] for b in range(3)] for a in range(3)]
Gam = [[[trunc(sum(gami[i, e] * (dgam[c][e][b] + dgam[b][e][c] - dgam[e][b][c]) for e in range(3)) / 2, 2)
         for c in range(3)] for b in range(3)] for i in range(3)]


def te(e, deg):
    """truncate a polynomial in X and eps: total X-degree <= deg, eps-degree <= 1 (dual numbers)"""
    e = sp.expand(e)
    if e == 0:
        return sp.Integer(0)
    P = sp.Poly(e, *X, eps)
    return sum((cf * sp.prod([x**m for x, m in zip(X, mon[:3])]) * eps**mon[3] for mon, cf in P.terms()
                if sum(mon[:3]) <= deg and mon[3] <= 1), sp.Integer(0))


def tem(M, deg):
    return M.applyfunc(lambda e: te(e, deg))


def DF_at_origin(g, h):
    """The O(eps) part of F(g + eps h) at the origin, nine components C_abc (a < b). Every quantity is carried as a
    polynomial in X and eps, truncated to first order in eps after each product (no rational functions of eps)."""
    G = tem(g + eps * h, 3)
    G0 = G.subs(o)
    G0a, G0b = G0.subs(eps, 0), sp.diff(G0, eps)
    G0ai = G0a.inv()
    G0i = G0ai - eps * G0ai * G0b * G0ai
    A = G - G0
    gi, term = G0i, G0i
    for _ in range(3):
        term = tem(-G0i * A * term, 2)
        gi = gi + term
    gi = tem(gi, 2)
    dg = [[[sp.diff(G[b, c], X[a]) for c in range(3)] for b in range(3)] for a in range(3)]
    Gm = [[[te(sum(gi[i, e] * (dg[c][e][b] + dg[b][e][c] - dg[e][b][c]) for e in range(3)) / 2, 2)
            for c in range(3)] for b in range(3)] for i in range(3)]
    Riem = lambda a, b, c, d: te(sp.diff(Gm[a][d][b], X[c]) - sp.diff(Gm[a][c][b], X[d])
                                 + sum(Gm[a][c][e] * Gm[e][d][b] - Gm[a][d][e] * Gm[e][c][b] for e in range(3)), 1)
    Ric = sp.Matrix(3, 3, lambda b, d: sum(Riem(a, b, a, d) for a in range(3)))
    Rs = te(sum(gi[i, j] * Ric[i, j] for i in range(3) for j in range(3)), 1)
    M = tem(-gi * tem(Ric - Rs * G / 2, 1), 1)
    M0 = M.subs(o)
    M1 = [sp.diff(M, X[k]).subs(o) for k in range(3)]
    G0c = [[[Gm[i][k][e].subs(o) for e in range(3)] for k in range(3)] for i in range(3)]
    split = lambda Mx: (Mx.subs(eps, 0), sp.diff(Mx, eps))
    Ma, Mb = split(M0)
    Mai = Ma.inv()
    phi_a = sp.sqrt(Ma.det())
    phi = phi_a * (1 + eps * (Mai * Mb).trace() / 2)
    Minv = Mai - eps * Mai * Mb * Mai
    lin = lambda Mx: Mx.applyfunc(lambda e: sp.expand(e).coeff(eps, 0) + eps * sp.expand(e).coeff(eps, 1))
    Kt0 = lin(phi * Minv)
    dK = [lin(phi * ((Minv * M1[k]).trace() / 2 * Minv - Minv * M1[k] * Minv)) for k in range(3)]
    g1 = [sp.diff(G, X[k]).subs(o) for k in range(3)]
    Kl0 = lin(G0 * Kt0)
    dKl = [lin(g1[k] * Kt0 + G0 * dK[k]) for k in range(3)]
    nab = lambda a, b, c: dKl[a][b, c] - sum(G0c[e][a][b] * Kl0[e, c] + G0c[e][a][c] * Kl0[b, e] for e in range(3))
    return [sp.expand(nab(a, b, c) - nab(b, a, c)).coeff(eps, 1) for (a, b, c) in CODIDX]


def cov2(h):
    """nabla h for a symmetric 2-tensor field (polynomial matrix), truncated: (nabla_a h)_bc"""
    return [sp.Matrix(3, 3, lambda b, c: trunc(sp.diff(h[b, c], X[a]) - sum(Gam[e][a][b] * h[e, c] + Gam[e][a][c] * h[b, e]
                                                                       for e in range(3)), 2)) for a in range(3)]


def rough_laplacian(h):
    Dh = cov2(h)                                         # Dh[a][b,c], degree <= 2
    # nabla_d (nabla h)_abc
    def DDh(d, a, b, c):
        val = sp.diff(Dh[a][b, c], X[d])
        val -= sum(Gam[e][d][a] * Dh[e][b, c] + Gam[e][d][b] * Dh[a][e, c] + Gam[e][d][c] * Dh[a][b, e] for e in range(3))
        return trunc(val, 1)
    return sp.Matrix(3, 3, lambda b, c: trunc(sum(gami[d, a] * DDh(d, a, b, c) for d in range(3) for a in range(3)), 1))


def cod_at_origin(k):
    Dk = cov2(k)
    return [(Dk[a][b, c] - Dk[b][a, c]).subs(o) for (a, b, c) in CODIDX]


# ---------------------------------------------------------------- P1
print("P1  DF_gamma h = Cod (1/2)(-Lap + 2) h on transverse-traceless jets")
mons = [m for d in range(4) for m in itertools.product(range(d + 1), repeat=3) if sum(m) == d]
coef = {(w, m): sp.Symbol(f"c_{w}_{m[0]}{m[1]}{m[2]}") for w in range(6) for m in mons}
H = sp.zeros(3, 3)
for w, (i, j) in enumerate(SYM):
    e = sum(coef[(w, m)] * sp.prod([x**k for x, k in zip(X, m)]) for m in mons)
    H[i, j] = H[j, i] = e
trH = trunc(sum(gami[i, j] * H[i, j] for i in range(3) for j in range(3)), 3)
DH = cov2(H)
divH = [trunc(sum(gami[a, b] * DH[a][b, c] for a in range(3) for b in range(3)), 2) for c in range(3)]
eqs = []
for e in [trH] + divH:
    if e != 0:
        eqs += [cf for _, cf in sp.Poly(e, *X).terms()]
unknowns = list(coef.values())
sol = sp.solve(eqs, unknowns, dict=True)[0]
free = [u for u in unknowns if u not in sol]
print(f"      {len(unknowns)} jet coefficients, {len(eqs)} constraint equations, {len(free)} free")
random.seed(336)
ok1, det1 = True, []
for trial in range(3):
    vals = {u: R(random.randint(-9, 9), random.randint(1, 7)) for u in free}
    h = H.subs(sol).subs(vals)
    lhs = [sp.radsimp(v) for v in DF_at_origin(gam, h)]
    kk = tmat((-rough_laplacian(h) + 2 * h) / 2, 1)
    rhs = cod_at_origin(kk)
    same = all(sp.simplify(a - b) == 0 for a, b in zip(lhs, rhs))
    ok1 &= same and any(v != 0 for v in lhs)
    det1.append(f"trial {trial}: {'equal' if same else 'DIFFERENT'}")
# control: a jet that is not trace-free
hc = sp.Matrix(3, 3, lambda i, j: (X[0] + X[1] * X[2] if i == j else 0) + (X[2]**3 / 3 if (i, j) in ((0, 1), (1, 0)) else 0))
lhs_c = [sp.radsimp(v) for v in DF_at_origin(gam, hc)]
rhs_c = cod_at_origin(tmat((-rough_laplacian(hc) + 2 * hc) / 2, 1))
control = any(sp.simplify(a - b) != 0 for a, b in zip(lhs_c, rhs_c))
check("O(eps) of F(gamma + eps h) at 0 = Cod((1/2)(-Lap h + 2h))(0) on three random TT jets; differs off TT",
      ok1 and control, "; ".join(det1) + f"; control differs: {control}")

# ---------------------------------------------------------------- P2
print("P2  level 3 on SU(2): constant trace-free tensors in a left-invariant frame")
epsl = lambda a, b, c: sp.LeviCivita(a, b, c)
Cf = [[[epsl(a, b, c) for c in range(3)] for b in range(3)] for a in range(3)]   # nabla_{E_a} E_b = C[a][b][c] E_c
brk = lambda a, b: [Cf[a][b][c] - Cf[b][a][c] for c in range(3)]                # [E_a, E_b]
def Rvec(a, b, c):
    """R(E_a, E_b) E_c = nabla_a nabla_b E_c - nabla_b nabla_a E_c - nabla_[a,b] E_c (constant coefficients)"""
    nab_b_c = Cf[b][c]
    t1 = [sum(nab_b_c[d] * Cf[a][d][f] for d in range(3)) for f in range(3)]
    nab_a_c = Cf[a][c]
    t2 = [sum(nab_a_c[d] * Cf[b][d][f] for d in range(3)) for f in range(3)]
    br = brk(a, b)
    t3 = [sum(br[d] * Cf[d][c][f] for d in range(3)) for f in range(3)]
    return [t1[f] - t2[f] - t3[f] for f in range(3)]
sec_ok = all(Rvec(a, b, b)[a] == 1 for a in range(3) for b in range(3) if a != b)   # <R(E_a,E_b)E_b, E_a> = 1
hv = sp.symbols("h0:5")
h3 = sp.Matrix([[hv[0], hv[2], hv[3]], [hv[2], hv[1], hv[4]], [hv[3], hv[4], -hv[0] - hv[1]]])
def nab_const(Tm):
    """(nabla_a T)_bc for a constant-component symmetric T in the frame"""
    return [sp.Matrix(3, 3, lambda b, c: -sum(Cf[a][b][d] * Tm[d, c] + Cf[a][c][d] * Tm[b, d] for d in range(3))) for a in range(3)]
Dh3 = nab_const(h3)
div3 = [sp.expand(sum(Dh3[a][a, c] for a in range(3))) for c in range(3)]
# rough Laplacian: sum_a (nabla_a nabla_a h - nabla_{nabla_a E_a} h); nabla_a E_a = 0 here
def nab_3tensor_const(Ta, a):
    """nabla_a of the 3-tensor (nabla_x h)_bc (constant components Ta[x][b,c])"""
    return [sp.Matrix(3, 3, lambda b, c: -sum(Cf[a][x][d] * Ta[d][b, c] for d in range(3))
                      - sum(Cf[a][b][d] * Ta[x][d, c] + Cf[a][c][d] * Ta[x][b, d] for d in range(3))) for x in range(3)]
lap3 = sp.zeros(3, 3)
for a in range(3):
    lap3 += nab_3tensor_const(Dh3, a)[a]
norm2 = lambda Tm: sp.expand(sum(Tm[i, j]**2 for i in range(3) for j in range(3)))
grad2 = sp.expand(sum(norm2(Dh3[a]) for a in range(3)))
cod3 = sp.expand(sum((Dh3[a][b, c] - Dh3[b][a, c])**2 for a in range(3) for b in range(3) for c in range(3)))
ok2 = (sec_ok and all(v == 0 for v in div3) and sp.expand(lap3 + 6 * h3) == sp.zeros(3, 3)
       and sp.expand(grad2 - 6 * norm2(h3)) == 0 and sp.expand(cod3 - 18 * norm2(h3)) == 0
       and sp.expand(cod3 - 2 * grad2 - 6 * norm2(h3)) == 0 and 16 * 18 == 288)
check("sectional curvature 1; div h = 0; -Lap h = 6h; |nabla h|^2 = 6|h|^2; |Cod h|^2 = 18|h|^2 = 2|nabla h|^2 + 6|h|^2", ok2)

# ---------------------------------------------------------------- P3
print("P3  the augmented symbol")
xi = sp.Matrix(sp.symbols("y0:3", real=True))
Hs = sp.Matrix(3, 3, lambda i, j: sp.Symbol(f"H{min(i, j)}{max(i, j)}"))
wedge = [xi[a] * Hs[b, c] - xi[b] * Hs[a, c] for a in range(3) for b in range(3) for c in range(3)]
ok3 = True
for xv in ([1, 0, 0], [R(3, 5), R(4, 5), 0], [R(2, 3), R(1, 3), R(2, 3)], [1, -2, 3]):
    sub = dict(zip(xi, xv))
    eqs3 = [w.subs(sub) for w in wedge] + list((Hs * xi).subs(sub)) + [(Hs.trace() * xi[k]).subs(sub) for k in range(3)]
    syms = sorted(Hs.free_symbols, key=str)
    A3 = sp.Matrix([[sp.diff(e, s) for s in syms] for e in eqs3])
    ok3 &= A3.rank() == 6
    # the norm identity on {H xi = 0, tr H = 0}
    cons = list((Hs * xi).subs(sub)) + [Hs.trace()]
    solc = sp.solve(cons, syms, dict=True)[0]
    lhs = sum((w.subs(sub).subs(solc))**2 for w in wedge)
    rhs = 2 * sum(v**2 for v in xv) * sum((Hs[i, j].subs(solc))**2 for i in range(3) for j in range(3))
    ok3 &= sp.expand(lhs - rhs) == 0
check("(xi ^ H, H xi, (tr H) xi) injective in H; |xi ^ H|^2 = 2|xi|^2|H|^2 on {H xi = 0, tr H = 0}", ok3)

# ---------------------------------------------------------------- P4
print("P4  the count of ten and the determinant")
mp.mp.dps = 70
hz = lambda x: mp.zeta(x, 3)


def Z(s, p, scale=1, K=200):
    """sum_{n>=3} 2(n^2-4) (n^p (n^2 - 1)/scale)^{-s}: n^(p+2)(1 - n^-2), expanded binomially"""
    tot = mp.mpf(0)
    for k in range(K):
        tot += mp.rf(s, k) / mp.factorial(k) * 2 * (hz((p + 2) * s + 2 * k - 2) - 4 * hz((p + 2) * s + 2 * k))
    return mp.power(scale, s) * tot


ok4 = True
for s in (2, 3):
    d3 = mp.nsum(lambda n: 2 * (n * n - 4) * (n * (n * n - 1)) ** (-s), [3, mp.inf])
    d13 = mp.nsum(lambda n: 2 * (n * n - 4) * (n * n * (n * n - 1)**2 / 2) ** (-s), [3, mp.inf])
    ok4 &= abs(Z(mp.mpf(s), 1) - d3) < mp.mpf(10)**-25
    # n^2 (n^2-1)^2/2 = n^6 (1 - n^-2)^2 / 2: expand (1 - n^-2)^{-2s} as a binomial in n^-2
    Z13 = lambda ss: mp.power(2, ss) * sum(mp.rf(2 * ss, k) / mp.factorial(k) * 2 * (hz(6 * ss + 2 * k - 2) - 4 * hz(6 * ss + 2 * k))
                                           for k in range(200))
    ok4 &= abs(Z13(mp.mpf(s)) - d13) < mp.mpf(10)**-25
s0 = mp.mpf(10)**-40
z3, z13 = Z(s0, 1), Z13(s0)
ok4 &= abs(z3 - 10) < 1e-30 and abs(z13 - 10) < 1e-30
lndet_num = -mp.diff(lambda ss: Z(ss, 1), 0)
closed = -(6 * (mp.zeta(-2, 3, 1) - 4 * mp.zeta(0, 3, 1))
           + mp.nsum(lambda k: 2 / k * (hz(2 * k - 2) - 4 * hz(2 * k)), [1, mp.inf]))
ok4 &= abs(lndet_num - closed) < mp.mpf(10)**-25 and abs(closed - mp.mpf("-14.0371968")) < 1e-7
check("zeta(0) = 10 for both multipliers; ln det_zeta(n(n^2-1)) on TT: numerical = closed form = -14.0371968...", ok4,
      f"zeta(0) = {mp.nstr(z3, 12)}, {mp.nstr(z13, 12)}; ln det = {mp.nstr(lndet_num, 15)} vs {mp.nstr(closed, 15)}")

# ---------------------------------------------------------------- P5
print("P5  the shell dipole's decomposition")
r, th, ph = sp.symbols("r theta phi", positive=True)
lam, c, q0 = sp.symbols("lambda c q0")
Y = [r, th, ph]
gS = sp.diag(1, sp.sin(r)**2, sp.sin(r)**2 * sp.sin(th)**2)
giS = gS.inv()
GS = [[[sum(giS[i, e] * (sp.diff(gS[e, b], Y[cc]) + sp.diff(gS[e, cc], Y[b]) - sp.diff(gS[b, cc], Y[e])) for e in range(3)) / 2
        for cc in range(3)] for b in range(3)] for i in range(3)]
w = (lam + c * sp.sin(r)) * sp.cos(th)
eta = [(c * sp.cos(r) + q0) * sp.cos(th), -(c * sp.sin(r) + q0 * sp.sin(r) * sp.cos(r) + lam * sp.sin(r)**2) * sp.sin(th), 0]
Leta = sp.Matrix(3, 3, lambda a, b: sp.diff(eta[b], Y[a]) + sp.diff(eta[a], Y[b]) - 2 * sum(GS[e][a][b] * eta[e] for e in range(3)))
hS = sp.zeros(3, 3)
hS[0, 0] = 2 * lam * sp.cos(th)
ok5 = sp.simplify(2 * w * gS + Leta - hS) == sp.zeros(3, 3)
check("2 lambda cos(theta) dr^2 = 2 w gamma + L_eta gamma, w = (lambda + c sin r) cos(theta), all c, q0", ok5)

print()
print(f"test_336: {'ALL PASS' if not FAILURES else 'FAILURES: ' + ', '.join(FAILURES)}")
sys.exit(1 if FAILURES else 0)