#!/usr/bin/env python3
"""
test_348 -- J AT THE ROUND SPHERE: THE SECOND VARIATION BY DIRECT EVALUATION, ITS RADICAL, THE CONFORMAL SECTOR, AND
            THE DEFICIT ALONG MA's CONFORMAL PATH

For MB §§5.2-5.4 (the companion of MA). Conventions as in MA: M = -G = adj K; J(g) = int sqrt(det M) dV_g; gamma the
unit round metric on S^3; TT levels n >= 3 with rough-Laplacian eigenvalue mu_n = n^2 - 3 (MA §2.6).

THE ENGINE. S^3 is the unit quaternions, x = x0 + x1 i + x2 j + x3 k. The left-invariant fields X_a(x) = x e_a
(e_a = i, j, k) are linear vector fields on R^4, orthonormal for gamma, with [X_a, X_b] = 2 eps_abc X_c, so they act on
polynomials as polynomials. A metric g = g_ab theta^a theta^b with polynomial (or power-series-in-t) coefficients is
handled exactly: Koszul's formula for the connection, R(X_a, X_b) X_c from it, Ric_ab = sum_c <R(X_c, X_b) X_a, theta^c>,
M = (R/2) Id - g^-1 Ric, dV_g = sqrt(det g_ab) dV_gamma. Series in t are truncated; integrals over S^3 are exact,
monomial by monomial: int x^alpha dV = 2 prod Gamma((alpha_i + 1)/2) / Gamma((|alpha| + 4)/2) (all alpha_i even), so
int 1 = 2 pi^2. Nothing here uses the reduction formula of paper 2 or the conformal formula being tested.

PRE-REGISTERED:
  P1 (TT eigentensors) With z1 = x0 + i x1, z2 = x2 + i x3 and E = theta^2 - i theta^3, the real tensors
     h = Re(z1^a conj(z2)^d E (x) E), a + d = n - 3, are traceless, divergence-free and satisfy nabla* nabla h = (n^2 - 3) h
     exactly, at n = 3, 4, 5, 6, 7. The orientation-reversed tensor at n = 4 (the pullback by q -> conj(q)) is TT with
     the same eigenvalue, and not of the same form (its theta^1 theta^1 component is nonzero).
  P2 (the second variation, directly) Along gamma + t h_n, J = 2 pi^2 + t J1 + t^2 Q + O(t^3) with J1 = 0 and
     Q / ||h_n||^2 = -(mu_n + 2)(mu_n + 3)/16 exactly, i.e. -9/2, -15, -75/2, -315/4, -147 at n = 3..7
     (mu = 6, 13, 22, 33, 46); the same value for the mirror tensor at n = 4. The two ingredients of
     paper 2's reduction, at each n: the first-order M is M1 = -(mu_n + 2) h_n / 2 pointwise, and
     int R2 dV_gamma = -(mu_n - 2) ||h_n||^2 / 4 (R2 the t^2 coefficient of the scalar curvature).
  P3 (the radical) Q(L_Y gamma) = 0 for Y = (x0 x1, x2^2, x3) in the frame; Q(u gamma) = 0 for u = x0 x2 + x1^3;
     Q(h_4 + u gamma + L_Y gamma) = Q(h_4) exactly, and J1 = 0 along it.
  P4 (the conformal sector)
     (a) For g = e^{2u} gamma, e^{2u} M(g) = (1 - Delta u) Id + Hess u - du (x) du exactly (the exponential carried as a
         symbol E with X_a(E) = 2 E X_a(u)), for u = x0 x1, x0^3 + x1 x2, x2.
     (b) On a ray u = t f, with H = Hess f, p = grad f, q = Delta f, the integrand
         sqrt(det[Id + t(H - q Id) - t^2 p p^T]) has t-coefficients c1 = -q, c2 = (sigma2(H) - |p|^2)/2,
         c3 = (det H + H(p, p))/2, as identities in a symbolic symmetric H and vector p.
     (c) On the unit S^3, div T2(Hess f) = -T1(Hess f) grad f (T1 = q Id - H, T2 = adj H), for f = x0 x1 x2,
         x0^3 + x1^2 x3, x0^4 - x1 x2 x3.
     (d) int c1 = int c2 = int c3 = 0 for f = x0 x1, x0 x1 x2, x0 x1 + x0 x2/2, x0^3 + x1^2 x3.
     (e) The t^4 coefficient of J(e^{2tf} gamma)/(2 pi^2): -1/24 for x0 x1; -3/160 for x0 x1 x2; -25/384 for
         x0 x1 + x0 x2/2; 0 for the zonal x0^2 and x0^3 (and for x0^2 also at t^5 and t^6); 0 for the first harmonic x1.
  P5 (the deficit, through the full engine; a realizable control) For f = x0 x1, with the curvature engine itself (not
     P4(a)'s formula), J(e^{2tf} gamma)/(2 pi^2) = 1 - t^4/24 + O(t^5), and along the radial path
     e^{2tf}(gamma + t^2 df (x) df), which is the metric of the embedding x -> e^{t f(x)} x, J/(2 pi^2) = 1 + O(t^5).
     The conformal metric stays in U: at t = 1/5 the determinant in P4(a) is positive at 200 random points of S^3.
"""
import sys
import random
import time
import sympy as sp
from sympy import gamma as Gam_

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)


R3 = range(3)
x = sp.symbols("x0:4", real=True)
t = sp.Symbol("t")


def qmul(p, q):
    a1, b1, c1, d1 = p
    a2, b2, c2, d2 = q
    return (a1 * a2 - b1 * b2 - c1 * c2 - d1 * d2, a1 * b2 + b1 * a2 + c1 * d2 - d1 * c2,
            a1 * c2 - b1 * d2 + c1 * a2 + d1 * b2, a1 * d2 + b1 * c2 - c1 * b2 + d1 * a2)


UNITS = [(0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1)]
VF = [qmul(x, u) for u in UNITS]


def eps(i, j, k):
    return ((i - j) * (j - k) * (k - i)) // 2


def make_D(extra=None):
    def D(e, a):
        r = sum(sp.diff(e, x[i]) * VF[a][i] for i in range(4))
        for s, f in (extra or []):
            r += sp.diff(e, s) * f(a)
        return r
    return D


D = make_D()


def trunc(e, n):
    e = sp.expand(e)
    if n is None or not e.has(t):
        return e
    p = sp.Poly(e, t)
    return sum(c * t**m[0] for m, c in zip(p.monoms(), p.coeffs()) if m[0] <= n)


def curvature(g, Dd, n, ginv):
    T = lambda e: trunc(e, n)
    C = [[[sum(2 * eps(a, b, m) * g[m, c] for m in R3) for c in R3] for b in R3] for a in R3]
    Gl = [[[T(sp.Rational(1, 2) * (Dd(g[b, c], a) + Dd(g[a, c], b) - Dd(g[a, b], c)
                                   + C[a][b][c] - C[b][c][a] + C[c][a][b])) for c in R3] for b in R3] for a in R3]
    Gu = [[[T(sum(ginv[c, d] * Gl[a][b][d] for d in R3)) for c in R3] for b in R3] for a in R3]

    def Rv(a, b, c):
        return [Dd(Gu[b][c][e], a) - Dd(Gu[a][c][e], b)
                + sum(Gu[b][c][d] * Gu[a][d][e] - Gu[a][c][d] * Gu[b][d][e] for d in R3)
                - sum(2 * eps(a, b, m) * Gu[m][c][e] for m in R3) for e in R3]
    Ric = sp.Matrix(3, 3, lambda a, b: T(sum(Rv(c, b, a)[c] for c in R3)))
    Rmix = sp.Matrix(3, 3, lambda a, b: T(sum(ginv[a, c] * Ric[c, b] for c in R3)))
    Rs = T(Rmix.trace())
    return (Rs / 2 * sp.eye(3) - Rmix).applyfunc(T), Rs


def series_inv(g, n):
    A = (g - sp.eye(3)).applyfunc(sp.expand)
    inv, P = sp.eye(3), sp.eye(3)
    for _ in range(n):
        P = (-P * A).applyfunc(lambda e: trunc(e, n))
        inv = inv + P
    return inv.applyfunc(lambda e: trunc(e, n))


def sqrt_series(e, n):
    e = trunc(e, n)
    out, term = 0, 1
    for k in range(n + 1):
        out += sp.binomial(sp.Rational(1, 2), k) * term
        term = trunc(term * e, n)
    return trunc(out, n)


def sphere_int(expr):
    p = sp.Poly(sp.expand(expr), *x)
    tot = 0
    for m, c in zip(p.monoms(), p.coeffs()):
        if any(k % 2 for k in m):
            continue
        tot += c * 2 * sp.prod([Gam_(sp.Rational(k + 1, 2)) for k in m]) / Gam_(sp.Rational(sum(m) + 4, 2))
    return sp.simplify(tot)


def J_coeffs(g, N):
    """[int of the t^k coefficient of sqrt(det M) dV_g / dV_gamma, k = 0..N], and the scalar-curvature series."""
    M, Rs = curvature(g, D, N, series_inv(g, N))
    integ = trunc(sqrt_series(trunc(M.det(), N) - 1, N) * sqrt_series(trunc(g.det() - 1, N), N), N)
    P = sp.Poly(integ, t)
    return [sphere_int(P.coeff_monomial(t**k)) for k in range(N + 1)], M, Rs


def nab(c, hh):
    return sp.Matrix(3, 3, lambda a, b: D(hh[a, b], c) - sum(eps(c, a, m) * hh[m, b] + eps(c, b, m) * hh[a, m]
                                                          for m in R3))


def rough(h):
    lap = sp.zeros(3, 3)
    for c in R3:
        lap += nab(c, nab(c, h))
    return (-lap).applyfunc(sp.expand)


def div(h):
    return [sp.expand(sum(D(h[c, b], c) - sum(eps(c, b, m) * h[c, m] for m in R3) for c in R3)) for b in R3]


def norm2(h):
    return sphere_int(sum(h[i, j]**2 for i in R3 for j in R3))


I = sp.I
z1, z2 = x[0] + I * x[1], x[2] + I * x[3]
Amat = sp.diag(0, 1, -1)
Bmat = sp.Matrix([[0, 0, 0], [0, 0, 1], [0, 1, 0]])


def tt(a_, d_):
    phi = sp.expand(z1**a_ * sp.conjugate(z2)**d_)
    re = sp.expand(sp.expand_complex(phi + sp.conjugate(phi)) / 2)
    im = sp.expand(sp.expand_complex((phi - sp.conjugate(phi)) / (2 * I)))
    return (re * Amat + im * Bmat).applyfunc(sp.expand)


# ------------------------------------------------------------------------------------------------------------- P1
print("P1  TT eigentensors at levels 3-7, and a mirror")
T0 = time.time()
levels = {3: (0, 0), 4: (1, 0), 5: (1, 1), 6: (2, 1), 7: (2, 2)}
H_ = {n: tt(*ad) for n, ad in levels.items()}
ok1 = True
for n, h in H_.items():
    mu = n * n - 3
    ok1 &= h.trace() == 0 and all(v == 0 for v in div(h)) and (rough(h) - mu * h).applyfunc(sp.expand) == sp.zeros(3, 3)
check("h_n traceless, divergence-free, nabla* nabla h_n = (n^2 - 3) h_n, n = 3..7", ok1)
qbar = (x[0], -x[1], -x[2], -x[3])
Ad = sp.zeros(3, 3)
for a in R3:
    v = qmul(qmul(x, UNITS[a]), qbar)
    for c in R3:
        Ad[c, a] = sp.expand(v[c + 1])
hb = H_[4].subs({x[1]: -x[1], x[2]: -x[2], x[3]: -x[3]}, simultaneous=True)
hm = (Ad.T * hb * Ad).applyfunc(sp.expand)
okm = hm.trace() == 0 and all(v == 0 for v in div(hm)) and (rough(hm) - 13 * hm).applyfunc(sp.expand) == sp.zeros(3, 3)
check("the orientation-reversed h_4 is TT with eigenvalue 13, with a nonzero theta^1 theta^1 component", okm
      and hm[0, 0] != 0, f"{time.time() - T0:.0f}s")

# ------------------------------------------------------------------------------------------------------------- P2
print("P2  J along gamma + t h_n, directly")
ok2, ok2i = True, True
rows = []
for n, h in list(H_.items()) + [("4 mirror", hm)]:
    mu = 13 if n == "4 mirror" else n * n - 3
    c, M, Rs = J_coeffs(sp.eye(3) + t * h, 2)
    nh = norm2(h)
    ratio = sp.simplify(c[2] / nh)
    law = sp.Rational(-(mu + 2) * (mu + 3), 16)
    ok2 &= c[0] == 2 * sp.pi**2 and c[1] == 0 and ratio == law
    M1 = M.applyfunc(lambda e: sp.Poly(e, t).coeff_monomial(t) if e.has(t) else 0)
    ok2i &= (M1 + sp.Rational(mu + 2, 2) * h).applyfunc(sp.expand) == sp.zeros(3, 3)
    R2 = sp.Poly(Rs, t).coeff_monomial(t**2)
    ok2i &= sp.simplify(sphere_int(R2) + sp.Rational(mu - 2, 4) * nh) == 0
    rows.append(f"n={n}: {ratio}")
check("J1 = 0 and Q/||h||^2 = -(mu+2)(mu+3)/16 at n = 3..7 and for the mirror", ok2, "; ".join(rows))
check("ingredients: M1 = -(mu+2) h/2 pointwise; int R2 = -(mu-2) ||h||^2/4", ok2i, f"{time.time() - T0:.0f}s")

# ------------------------------------------------------------------------------------------------------------- P3
print("P3  the radical: Lie and conformal directions, and the mixed terms")
Y = [x[0] * x[1], x[2]**2, x[3]]
LY = sp.Matrix(3, 3, lambda a, b: sp.expand(D(Y[b], a) + D(Y[a], b)))
u = x[0] * x[2] + x[1]**3
cL = J_coeffs(sp.eye(3) + t * LY, 2)[0]
cU = J_coeffs(sp.eye(3) + t * u * sp.eye(3), 2)[0]
comb = (H_[4] + u * sp.eye(3) + LY).applyfunc(sp.expand)
cC = J_coeffs(sp.eye(3) + t * comb, 2)[0]
c4 = J_coeffs(sp.eye(3) + t * H_[4], 2)[0]
check("Q(L_Y gamma) = 0, Q(u gamma) = 0, Q(h_4 + u gamma + L_Y gamma) = Q(h_4), J1 = 0",
      cL[1] == 0 and cL[2] == 0 and cU[1] == 0 and cU[2] == 0 and cC[1] == 0 and sp.simplify(cC[2] - c4[2]) == 0,
      f"Q(h_4) = {c4[2]}; {time.time() - T0:.0f}s")

# ------------------------------------------------------------------------------------------------------------- P4
print("P4  the conformal sector")
Es = sp.Symbol("E", positive=True)


def hess(f):
    p = [D(f, a) for a in R3]
    H = sp.Matrix(3, 3, lambda a, b: sp.expand(D(p[b], a) - sum(eps(a, b, c) * p[c] for c in R3)))
    return H, p


ok4a = True
for uu in (x[0] * x[1], x[0]**3 + x[1] * x[2], x[2]):
    du = [D(uu, a) for a in R3]
    DE = make_D([(Es, lambda a, du=du: 2 * Es * du[a])])
    M, _ = curvature(Es * sp.eye(3), DE, None, sp.eye(3) / Es)
    H, _ = hess(uu)
    rhs = (1 - H.trace()) * sp.eye(3) + H - sp.Matrix(3, 3, lambda a, b: du[a] * du[b])
    ok4a &= (Es * M - rhs).applyfunc(sp.simplify) == sp.zeros(3, 3)
check("(a) e^{2u} M(e^{2u} gamma) = (1 - Delta u) Id + Hess u - du du, exactly", ok4a)
Hs = sp.Matrix(3, 3, lambda a, b: sp.Symbol(f"H{min(a, b)}{max(a, b)}"))
ps = sp.Matrix(sp.symbols("p0:3"))
qs = Hs.trace()
ser = sqrt_series(sp.expand((sp.eye(3) + t * (Hs - qs * sp.eye(3)) - t**2 * ps * ps.T).det()) - 1, 3)
Ps = sp.Poly(ser, t)
sig2 = sp.expand((qs**2 - (Hs * Hs).trace()) / 2)
ok4b = sp.expand(Ps.coeff_monomial(t) + qs) == 0
ok4b &= sp.expand(Ps.coeff_monomial(t**2) - (sig2 - (ps.T * ps)[0]) / 2) == 0
ok4b &= sp.expand(Ps.coeff_monomial(t**3) - (Hs.det() + (ps.T * Hs * ps)[0]) / 2) == 0
check("(b) c1 = -q, c2 = (sigma2(H) - |p|^2)/2, c3 = (det H + H(p,p))/2", ok4b)
ok4c = True
for f in (x[0] * x[1] * x[2], x[0]**3 + x[1]**2 * x[3], x[0]**4 - x[1] * x[2] * x[3]):
    H, p = hess(f)
    T2 = H.adjugate().applyfunc(sp.expand)
    T1 = H.trace() * sp.eye(3) - H
    ok4c &= all(sp.expand(dv + sum(T1[b, c] * p[c] for c in R3)) == 0 for b, dv in enumerate(div(T2)))
check("(c) div T2(Hess f) = -T1(Hess f) grad f on the unit S^3", ok4c)


def conf_coeffs(f, N):
    H, p = hess(f)
    q = H.trace()
    X = (t * (H - q * sp.eye(3)) - t**2 * sp.Matrix(3, 3, lambda a, b: p[a] * p[b])).applyfunc(sp.expand)
    P = sp.Poly(sqrt_series(trunc((sp.eye(3) + X).det(), N) - 1, N), t)
    return [sp.simplify(sphere_int(P.coeff_monomial(t**k)) / (2 * sp.pi**2)) for k in range(N + 1)]


fs = {"x0x1": x[0] * x[1], "x0x1x2": x[0] * x[1] * x[2], "x0x1 + x0x2/2": x[0] * x[1] + x[0] * x[2] / 2,
      "x0^3 + x1^2 x3": x[0]**3 + x[1]**2 * x[3]}
cc = {nm: conf_coeffs(f, 4) for nm, f in fs.items()}
check("(d) int c1 = int c2 = int c3 = 0 for four f, harmonic and not", all(v[1:4] == [0, 0, 0] for v in cc.values()),
      "; ".join(f"{k}: {v}" for k, v in cc.items()))
want = {"x0x1": sp.Rational(-1, 24), "x0x1x2": sp.Rational(-3, 160), "x0x1 + x0x2/2": sp.Rational(-25, 384)}
z2c = conf_coeffs(x[0]**2, 6)
z3c = conf_coeffs(x[0]**3, 4)
h1c = conf_coeffs(x[1], 4)
check("(e) quartic: -1/24, -3/160, -25/384; zonal x0^2 (through t^6) and x0^3, and the first harmonic x1: 0",
      all(cc[k][4] == v for k, v in want.items()) and z2c[1:] == [0] * 6 and z3c[1:] == [0] * 4
      and h1c[1:] == [0] * 4, f"x0^3+x1^2x3: {cc['x0^3 + x1^2 x3'][4]}; {time.time() - T0:.0f}s")

# ------------------------------------------------------------------------------------------------------------- P5
print("P5  the deficit through the full engine; the radial control")
f = x[0] * x[1]
N = 4
Ef = trunc(sum((2 * t * f)**k / sp.factorial(k) for k in range(N + 1)), N)
pf = [D(f, a) for a in R3]
g_rad = (Ef * (sp.eye(3) + t**2 * sp.Matrix(3, 3, lambda a, b: pf[a] * pf[b]))).applyfunc(lambda e: trunc(e, N))
g_conf = (Ef * sp.eye(3)).applyfunc(lambda e: trunc(e, N))
cr = [sp.simplify(v / (2 * sp.pi**2)) for v in J_coeffs(g_rad, N)[0]]
cf = [sp.simplify(v / (2 * sp.pi**2)) for v in J_coeffs(g_conf, N)[0]]
H, p = hess(f)
detf = ((1 - t * H.trace()) * sp.eye(3) + t * H - t**2 * sp.Matrix(3, 3, lambda a, b: p[a] * p[b])).det()
random.seed(348)
okU = True
for _ in range(200):
    v = [random.gauss(0, 1) for _ in range(4)]
    nv = sum(w * w for w in v) ** 0.5
    okU &= float(detf.subs({**{x[i]: v[i] / nv for i in range(4)}, t: 0.2})) > 0
check("conformal path: J/2pi^2 = 1 - t^4/24 + O(t^5); radial path: 1 + O(t^5); det > 0 at t = 1/5",
      cf == [1, 0, 0, 0, sp.Rational(-1, 24)] and cr == [1, 0, 0, 0, 0] and okU,
      f"conformal {cf}, radial {cr}; {time.time() - T0:.0f}s")

print()
if FAILURES:
    print(f"test_348: {len(FAILURES)} FAILED")
    sys.exit(1)
print("test_348: ALL PASS")