#!/usr/bin/env python3
"""
test_351 -- THE CONVEX BRANCH: SUPPORT FUNCTIONS, THE EQUIVARIANT CONTRACTION, HUISKEN'S MODE RATES, THE
            CROSS-CURVATURE THRESHOLD ON BIANCHI IX, AND THE REVOLUTION CLASS

For MB §7 (the companion of MA). Conventions as in MA. The engine is test_348's: S^3 is the unit quaternions, the
left-invariant fields X_a(x) = x e_a are linear on R^4 and orthonormal for gamma, [X_a, X_b] = 2 eps_abc X_c; the
Hessian of a function in this frame is X_a X_b f - eps_abc X_c f.

PRE-REGISTERED:
  P1 (the support-function parametrisation) For h a polynomial on S^3 and nu in S^3, put
     x_h(nu) = h(nu) nu + sum_a (X_a h)(nu) X_a(nu) (so x_h = h nu + grad h) and L_h = Hess h + h Id. Then, as an
     identity of R^4-valued polynomials, X_b(x_h) = sum_a (L_h)_ab X_a(nu), for h = x0 x1 + x2^3 - 2 x1 x3 + 5. So x_h
     has Gauss map nu, induced metric gamma(L_h., L_h.) in the frame, and shape operator L_h^{-1}.
  P2 (Gauss for g_h, through the engine) For h = 1 + t Y with Y = x0 x1 + x2^3, the metric g_h = L_h L_h (frame
     components) has M(g_h) = L_h / det L_h = adj(L_h^{-1}) through t^2, exactly.
  P3 (the straight-line contraction) L_{(1-s)h + s} = (1-s) L_h + s Id, symbolically; for the h of P1 scaled so that
     L_h > 0 at 500 random points, L stays positive definite along s in {0, 1/4, 1/2, 3/4, 1}; the centring
     int h x dV = 0 is preserved (exact integrals: int x dV = 0); and the map commutes with h -> h o A, A in O(4),
     because the constant 1 is O(4)-invariant (checked for a random rational rotation A on the h of P1).
  P4 (Huisken's linearisation) For R > 0 and Y_l = x0, x0 x1, x0 x1 x2, Re(z1^4) (z1 = x0 + i x1), of degrees
     l = 1, 2, 3, 4: Delta Y_l = -l(l+2) Y_l, and the mean curvature of the support function R + eps Y_l is
     tr(L^-1) = 3/R + eps (l(l+2) - 3) Y_l / R^2 + O(eps^2) (exact series). So d(eps)/dt = -[l(l+2) - 3] eps / R^2
     under h_t = Hbar - H: every l >= 2 decays, l = 1 is neutral.
  P5 (the cross-curvature threshold) For the left-invariant metric diag(1 + e, 1, 1 - e) in the unit sphere's frame
     (Milnor constants lambda_i = 2 sqrt(a_i / (a_j a_k))), M > 0 iff |e| < (sqrt2 - 1)/2: the first eigenvalue vanishes
     exactly at e = (sqrt2 - 1)/2 (symbolic), and M > 0 at e = 0.2, not at e = 0.21.
  P6 (the revolution class) Write an embedded hypersurface of revolution whose profile is a graph over the axis as
     (rho(z) omega, z), z in [-1, 1] after an affine change, and u = rho^2. For u0 = 1 - z^2 (the unit sphere) and u1
     the profile of the radial graph r = 1 + (3/10) cos 2chi (computed), u_s = (1 - s) u0 + s u1 is positive on (-1, 1)
     with nonzero derivative at z = +-1 for s in {0, 1/4, 1/2, 3/4, 1}; rho_s = sqrt(u_s) is convex in z somewhere
     (a neck, so K indefinite) at s = 1 and nowhere at s = 0.

FIRST RUN, AND THE CHANGES:
  P5 failed on its normalisation clause only: sp.N(M_i(0)) == 1 compared a float to 1 and failed. The threshold
  clauses held (M_1(e*) = 0 exactly, M > 0 at 0.2, M_1 < 0 at 0.21). Now compared exactly, sp.simplify(M_i(0)) == 1.
  P6 passed for a wrong reason: rho_zz was taken by numerical differentiation of a profile found by bisection, and
  the noise gave 4.9e30 at s = 1. Re-posed: rho_zz from u1 and its z-derivatives computed by the chain rule from the
  closed forms in chi, z only located by bisection; the slope at the poles likewise. The claim is unchanged.
  Before the bundle's final run, sympy's nsimplify was removed (it was applied to an already exact rational rotation,
  and to a printed value); the checks are unchanged.
"""
import sys
import random
from fractions import Fraction as Fr
import sympy as sp
import mpmath as mp
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 D(e, a):
    return sum(sp.diff(e, x[i]) * VF[a][i] for i in range(4))


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


def trunc(e, n):
    e = sp.expand(e)
    if 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, 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) * (D(g[b, c], a) + D(g[a, c], b) - D(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 [D(Gu[b][c][e], a) - D(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)


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 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)


# ------------------------------------------------------------------------------------------------------------- P1
print("P1  the support-function parametrisation")
h = x[0] * x[1] + x[2]**3 - 2 * x[1] * x[3] + 5
H, p = hess(h)
L = H + h * sp.eye(3)
xh = [sp.expand(h * x[i] + sum(p[a] * VF[a][i] for a in R3)) for i in range(4)]
ok1 = True
for b in R3:
    lhs = [sp.expand(D(xh[i], b)) for i in range(4)]
    rhs = [sp.expand(sum(L[a, b] * VF[a][i] for a in R3)) for i in range(4)]
    ok1 &= all(sp.expand(l - r) == 0 for l, r in zip(lhs, rhs))
check("X_b(x_h) = sum_a (Hess h + h Id)_ab X_a(nu): metric L^2, shape operator L^-1, Gauss map nu", ok1)

# ------------------------------------------------------------------------------------------------------------- P2
print("P2  Gauss for g_h through the curvature engine")
Y = x[0] * x[1] + x[2]**3
HY, _ = hess(Y)
Lt = ((1 + t * Y) * sp.eye(3) + t * HY).applyfunc(sp.expand)
g = (Lt * Lt).applyfunc(lambda e: trunc(e, 2))
M = curvature(g, 2, series_inv(g, 2))
detL = trunc(Lt.det(), 2)
inv_det = trunc(1 - (detL - 1) + (detL - 1)**2, 2)
want = (Lt * inv_det).applyfunc(lambda e: trunc(e, 2))
check("M(g_h) = L_h / det L_h through t^2, exactly", (M - want).applyfunc(sp.expand) == sp.zeros(3, 3))

# ------------------------------------------------------------------------------------------------------------- P3
print("P3  the straight-line contraction")
s_ = sp.Symbol("s")
hs = (1 - s_) * h + s_
Hs, _ = hess(hs)
ok3 = ((Hs + hs * sp.eye(3)) - ((1 - s_) * L + s_ * sp.eye(3))).applyfunc(sp.expand) == sp.zeros(3, 3)
random.seed(351)
pts = []
for _ in range(500):
    v = [random.gauss(0, 1) for _ in range(4)]
    nv = sum(w * w for w in v) ** 0.5
    pts.append([w / nv for w in v])
Lf = sp.lambdify(x, L, "mpmath")
mp.mp.dps = 30
okp = True
for sv in (0, 0.25, 0.5, 0.75, 1):
    for q in pts:
        Lq = (1 - sv) * mp.matrix(Lf(*q)) + sv * mp.eye(3)
        okp &= min(mp.eig(Lq)[0], key=lambda z: mp.re(z)).real > 0
ok3 &= okp
ok3 &= all(sphere_int(x[i]) == 0 for i in range(4))
# equivariance: a random rational rotation A of R^4 (Cayley), h o A, and the constant 1
sk = [Fr(random.randint(-3, 3), random.randint(1, 3)) for _ in range(6)]
S = sp.Matrix([[0, sk[0], sk[1], sk[2]], [-sk[0], 0, sk[3], sk[4]], [-sk[1], -sk[3], 0, sk[5]], [-sk[2], -sk[4], -sk[5], 0]])
Arot = ((sp.eye(4) - S) * (sp.eye(4) + S).inv()).applyfunc(sp.Rational)          # exact rationals already
okA = (Arot.T * Arot).applyfunc(sp.simplify) == sp.eye(4)
y = Arot * sp.Matrix(x)
hA = h.subs({x[i]: y[i] for i in range(4)}, simultaneous=True)
lhs = sp.expand(((1 - s_) * h + s_).subs({x[i]: y[i] for i in range(4)}, simultaneous=True))
ok3 &= okA and sp.expand(lhs - ((1 - s_) * hA + s_)) == 0
check("L along the line is (1-s) L_h + s Id; positive at 500 points for s in [0, 1]; centring kept; O(4)-equivariant",
      ok3)

# ------------------------------------------------------------------------------------------------------------- P4
print("P4  Huisken's linearisation at the round sphere")
Rr, ep = sp.symbols("R epsilon", positive=True)
z1 = x[0] + sp.I * x[1]
Ys = {1: x[0], 2: x[0] * x[1], 3: x[0] * x[1] * x[2], 4: sp.expand(sp.re(sp.expand(z1**4)))}
ok4 = True
for l, Yl in Ys.items():
    HY, _ = hess(Yl)
    ok4 &= sp.expand(HY.trace() + l * (l + 2) * Yl) == 0
    Lm = (Rr + ep * Yl) * sp.eye(3) + ep * HY
    # L^-1 = R^-1 (Id + (eps/R) B)^-1, B = Y Id + Hess Y: first order
    Bm = Yl * sp.eye(3) + HY
    trinv = 3 / Rr - ep / Rr**2 * Bm.trace()
    ok4 &= sp.expand(trinv - (3 / Rr + ep * (l * (l + 2) - 3) * Yl / Rr**2)) == 0
    # and against the exact inverse, to first order in eps at a point
    pt = {x[0]: sp.Rational(1, 2), x[1]: sp.Rational(1, 2), x[2]: sp.Rational(1, 2), x[3]: sp.Rational(1, 2)}
    exact = sp.series(Lm.subs(pt).inv().trace(), ep, 0, 2).removeO()
    ok4 &= sp.simplify(exact - trinv.subs(pt)) == 0
check("Delta Y_l = -l(l+2) Y_l and tr L^-1 = 3/R + eps (l(l+2) - 3) Y_l/R^2, l = 1..4: rates -[l(l+2)-3]/R^2", ok4)

# ------------------------------------------------------------------------------------------------------------- P5
print("P5  the cross-curvature threshold on Bianchi IX")
e = sp.Symbol("e", real=True)
a = [1 + e, 1, 1 - e]
lam = [2 * sp.sqrt(a[i] / (a[(i + 1) % 3] * a[(i + 2) % 3])) for i in R3]
mu = [sum(lam) / 2 - lv for lv in lam]
Mi = [mu[i] * lam[i] - mu[(i + 1) % 3] * mu[(i + 2) % 3] for i in R3]
es = (sp.sqrt(2) - 1) / 2
ok5 = sp.simplify(Mi[0].subs(e, es)) == 0 and all(sp.N(m.subs(e, es)) > 0 for m in Mi[1:])
ok5 &= all(sp.N(m.subs(e, sp.Rational(1, 5))) > 0 for m in Mi) and sp.N(Mi[0].subs(e, sp.Rational(21, 100))) < 0
ok5 &= all(sp.simplify(m.subs(e, 0)) == 1 for m in Mi)
check("M > 0 iff |e| < (sqrt2 - 1)/2: zero exactly there, positive at 0.2, negative at 0.21", ok5,
      f"M at e*: {[sp.radsimp(sp.simplify(m.subs(e, es))) for m in Mi]}")

# ------------------------------------------------------------------------------------------------------------- P6
print("P6  the revolution class is convex in u = rho^2")
ev = mp.mpf(3) / 10
X1 = lambda c: (1 + ev * mp.cos(2 * c)) * mp.cos(c)
RHO = lambda c: (1 + ev * mp.cos(2 * c)) * mp.sin(c)


def u1(zz):                                   # zz in [-1, 1]; x1 = (1 + eps) zz; solve x1(chi) = that
    target = (1 + ev) * zz
    lo, hi = mp.mpf(0), mp.pi                 # x1 decreasing on [0, pi]
    for _ in range(120):
        mid = (lo + hi) / 2
        lo, hi = (mid, hi) if X1(mid) > target else (lo, mid)
    return (RHO((lo + hi) / 2) / (1 + ev))**2


u0 = lambda zz: 1 - zz**2
ok6 = True
zgrid = [mp.mpf(k) / 50 for k in range(-49, 50)]
u1v = {zz: u1(zz) for zz in zgrid}
for sv in (0, mp.mpf(1) / 4, mp.mpf(1) / 2, mp.mpf(3) / 4, 1):
    us = lambda zz: (1 - sv) * u0(zz) + sv * (u1v[zz] if zz in u1v else u1(zz))
    ok6 &= all(us(zz) > 0 for zz in zgrid)


Zc = lambda c: X1(c) / (1 + ev)               # the profile in chi, normalised to z in [-1, 1]
Pc = lambda c: RHO(c) / (1 + ev)


def chi_of(zz):
    lo, hi = mp.mpf(0), mp.pi
    for _ in range(200):
        mid = (lo + hi) / 2
        lo, hi = (mid, hi) if Zc(mid) > zz else (lo, mid)
    return (lo + hi) / 2


def u1_jet(zz):                                  # u1, du1/dz, d2u1/dz2 from closed forms in chi (chain rule)
    c = chi_of(zz)
    Zp = mp.diff(Zc, c)
    du = lambda cc: 2 * Pc(cc) * mp.diff(Pc, cc) / mp.diff(Zc, cc)
    return Pc(c)**2, du(c), mp.diff(du, c) / Zp


def neck(sv):                                    # max of rho_zz near the equator, rho = sqrt(u_s)
    vals = []
    for zz in [mp.mpf(k) / 20 for k in range(-8, 9)]:
        a1, b1, c1 = u1_jet(zz)
        u = (1 - sv) * u0(zz) + sv * a1
        du = (1 - sv) * (-2 * zz) + sv * b1
        d2u = (1 - sv) * (-2) + sv * c1
        vals.append(d2u / (2 * mp.sqrt(u)) - du**2 / (4 * u**mp.mpf(1.5)))
    return max(vals)


for sv in (0, mp.mpf(1) / 4, mp.mpf(1) / 2, mp.mpf(3) / 4, 1):    # slope at the poles, from the chain rule
    for zend, sgn in ((mp.mpf(1) - mp.mpf("1e-8"), -1), (-mp.mpf(1) + mp.mpf("1e-8"), 1)):
        du_end = (1 - sv) * (-2 * zend) + sv * u1_jet(zend)[1]
        ok6 &= sgn * du_end > mp.mpf("0.1")
n0, n1 = neck(mp.mpf(0)), neck(mp.mpf(1))
ok6 &= n0 < 0 < n1
check("u_s > 0 inside, nonzero slope at the poles, for s in [0, 1]; no neck at s = 0, a neck at s = 1", ok6,
      f"max rho_zz near the equator: s=0 {mp.nstr(n0, 4)}, s=1 {mp.nstr(n1, 4)}")

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