#!/usr/bin/env python3
"""
test_353 -- MB VERSION 2: THE QUARTIC CONFORMAL COEFFICIENT, AXIS AVOIDANCE FOR THE CORRUGATED PROFILES, THE PREIMAGE
            COUNT BEHIND THE EXACT EXCESS, THE POLES OF THE RADIAL GRAPHS, AND THE ALGEBRA OF THEOREM 5.5'S PROOF

For MB (the companion of MA), version 2. An external review of version 1 asked for five things to be shown in the
text rather than left to the scripts or to "every amplitude tested". This script gates each of them. Conventions as in
MA and MB: on the unit S^3, Delta = -nabla* nabla, and for a homogeneous polynomial F of degree d on R^4 restricted to
S^3, with P = I - x x^T,
    p = grad f = grad F - d F x,    H = Hess f = P (D^2 F) P - d F P,    q = Delta f = tr H.
Integrals over S^3 are exact: int x^alpha = 2 prod Gamma((alpha_i+1)/2) / Gamma((|alpha|+4)/2) for even alpha, else 0.

PRE-REGISTERED:
  P1 (the quartic coefficient, MB Proposition 5.7(2)) For a general symmetric 3x3 H and p in R^3, q = tr H, the
     t-expansion of sqrt(det[(1 - t q) I + t H - t^2 p p^T]) has
        c1 = -q,  c2 = (sigma_2(H) - |p|^2)/2,  c3 = (det H + H(p,p))/2,
        c4 = q (det H + H(p,p))/2 - (sigma_2(H) + |p|^2)^2/8 - |Hp|^2/2
     (symbolic, exact).
  P2 (the table of Proposition 5.7(3), through c4) int_{S^3} c4 dV / (2 pi^2) is -1/24, -3/160, -25/384, -19/40 for
     f = x0 x1, x0 x1 x2, x0 x1 + x0 x2/2, x0^3 + x1^2 x3, and 0 for the zonal f = x0, x0^2, x0^3; and int c1, c2, c3
     vanish for all seven (exact monomial moments). This evaluates the quartic coefficient by a route independent of
     test_348, which expanded J directly.
  P3 (Example 5.3, axis avoidance) For psi = sigma + A sin 2k sigma (k >= 2 an integer), rho(sigma) = int_0^sigma cos psi:
     cos psi(pi - s) = -cos psi(s) (symbolic, k = 2..6), so rho(pi - sigma) = rho(sigma) and rho(pi) = 0. Since
     |cos psi - cos sigma| <= |A|, rho(sigma) >= sin sigma - |A| sigma, which is > 0 on (0, pi/2] when |A| < 2/pi. Sampled:
     for k in {2, 3, 5, 8} and A in {0.2, 0.4, 0.6, 0.63}, on 4000 points of (0, pi/2], rho >= sin - |A| sigma - 1e-9 and
     rho > 0; and every amplitude MB uses (A <= 0.4) is below 2/pi.
  P4 (Theorem 6.4(3), the preimage count) For r = 1 + eps cos 2chi and eps in {1/4, 3/10, 2/5, 9/10}:
     sin beta = k_sph r sin chi with r - r' cot chi = 1 - eps + 6 eps cos^2 chi > 0, so beta stays in (0, pi) on (0, pi)
     (sampled at 20000 points); beta' changes sign exactly twice, at chi = a and pi - a; and for 400 normals b spread
     over (0, pi), the number of preimages is 3 inside the belt (beta(pi - a), beta(a)), exactly one of them on the
     decreasing branch (det K < 0), and 1 outside it.
  P5 (Proposition 7.4, the poles of Example 6.2) With c = cos chi: the axial coordinate is Z = (1 - eps + 2 eps c^2) c and
     u = rho^2 = r^2 (1 - c^2), with r = 1 - eps + 2 eps c^2, polynomials in c; dZ/dc = 1 - eps + 6 eps c^2 > 0 for
     0 <= eps < 1; and at c = +-1, u = 0 and du/dc = -+2 (1 + eps)^2, so du/dZ != 0 there (symbolic).
  P6 (the algebra of Theorem 5.5(2)'s proof) (a) For symmetric M1 with tr M1 = 0 and any M2,
     sqrt(det(I + t M1 + t^2 M2)) = 1 + t^2 (tr M2/2 - tr(M1^2)/4) + O(t^3); for traceless h, sqrt(det(I + t h)) =
     1 - t^2 |h|^2/4 + O(t^3) (symbolic). (b) On TT h at the unit S^3, with G = -gamma, G' = nabla* nabla h / 2:
     d^2/dt^2 int R dV = -(mu/2 + 2)|h|^2, so int R^(2) dV = -(mu + 4)/4 + 3/2 = -(mu - 2)/4 per unit |h|^2.
     (c) Q = (int R^(2) - |M^(1)|^2 - |h|^2)/4 with M^(1) = -(mu + 2) h/2 gives -(mu + 2)(mu + 3)/16 (symbolic in mu).
"""
import sys
import math
import time
import sympy as sp

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)


T0 = time.time()

# ------------------------------------------------------------------------------------------------------------- P1
print("P1  the Taylor coefficients of the conformal integrand")
t = sp.symbols("t")
hs = sp.symbols("h11 h12 h13 h22 h23 h33")
pv = sp.Matrix(sp.symbols("p1 p2 p3"))
Hm = sp.Matrix([[hs[0], hs[1], hs[2]], [hs[1], hs[3], hs[4]], [hs[2], hs[4], hs[5]]])
qq = Hm.trace()
Dt = sp.expand(((1 - t * qq) * sp.eye(3) + t * Hm - t**2 * pv * pv.T).det())
ser = sp.expand(sp.series(sp.sqrt(Dt), t, 0, 5).removeO())
c = [ser.coeff(t, k) for k in range(5)]
s2 = sp.expand((qq**2 - (Hm * Hm).trace()) / 2)
s3 = sp.expand(Hm.det())
pp = sp.expand((pv.T * pv)[0])
Hpp = sp.expand((pv.T * Hm * pv)[0])
H2pp = sp.expand((pv.T * Hm * Hm * pv)[0])
want = [1, -qq, (s2 - pp) / 2, (s3 + Hpp) / 2,
        qq * (s3 + Hpp) / 2 - (s2 + pp)**2 / 8 - H2pp / 2]
ok1 = all(sp.expand(ci - wi) == 0 for ci, wi in zip(c, want))
check("c1 = -q, c2 = (sigma2 - |p|^2)/2, c3 = (det H + H(p,p))/2, "
      "c4 = q(det H + H(p,p))/2 - (sigma2 + |p|^2)^2/8 - |Hp|^2/2", ok1)

# ------------------------------------------------------------------------------------------------------------- P2
print("P2  the quartic coefficients of Proposition 5.7(3), through c4 and exact moments")
X = sp.symbols("x0:4")
GENS = X


def moment(alpha):
    if any(a % 2 for a in alpha):
        return sp.Integer(0)
    num = sp.prod([sp.gamma(sp.Rational(a + 1, 2)) for a in alpha])
    return 2 * num / sp.gamma(sp.Rational(sum(alpha) + 4, 2))


def integrate_S3(poly):
    tot = sp.Integer(0)
    for mon, coef in poly.terms():
        if any(e % 2 for e in mon):
            continue
        tot += coef * moment(mon)
    return sp.nsimplify(sp.simplify(tot))


def P(e):
    return sp.Poly(e, *GENS, domain="QQ")


def invariants(F, d):
    xs = [P(xi) for xi in X]
    Fp = P(F)
    grad = [P(sp.diff(F, xi)) for xi in X]
    D2 = [[P(sp.diff(F, xi, xj)) for xj in X] for xi in X]
    one = P(1)
    zero = P(0)
    Pm = [[(one if i == j else zero) - xs[i] * xs[j] for j in range(4)] for i in range(4)]

    def mm(A, B):
        return [[sum((A[i][k] * B[k][j] for k in range(4)), zero) for j in range(4)] for i in range(4)]

    PD = mm(mm(Pm, D2), Pm)
    H4 = [[PD[i][j] - Fp * d * Pm[i][j] for j in range(4)] for i in range(4)]
    p = [grad[i] - Fp * d * xs[i] for i in range(4)]
    H2 = mm(H4, H4)
    H3 = mm(H2, H4)
    tr1 = sum((H4[i][i] for i in range(4)), zero)
    tr2 = sum((H2[i][i] for i in range(4)), zero)
    tr3 = sum((H3[i][i] for i in range(4)), zero)
    sig2 = (tr1 * tr1 - tr2) * sp.Rational(1, 2)
    sig3 = (tr1 * tr1 * tr1 - tr1 * tr2 * 3 + tr3 * 2) * sp.Rational(1, 6)
    Hp = [sum((H4[i][j] * p[j] for j in range(4)), zero) for i in range(4)]
    pp_ = sum((p[i] * p[i] for i in range(4)), zero)
    Hpp_ = sum((p[i] * Hp[i] for i in range(4)), zero)
    H2pp_ = sum((Hp[i] * Hp[i] for i in range(4)), zero)
    return tr1, sig2, sig3, pp_, Hpp_, H2pp_


x0, x1, x2, x3 = X
cases = [
    (x0 * x1, 2, sp.Rational(-1, 24)),
    (x0 * x1 * x2, 3, sp.Rational(-3, 160)),
    (x0 * x1 + x0 * x2 / 2, 2, sp.Rational(-25, 384)),
    (x0**3 + x1**2 * x3, 3, sp.Rational(-19, 40)),
    (x0, 1, sp.Integer(0)),
    (x0**2, 2, sp.Integer(0)),
    (x0**3, 3, sp.Integer(0)),
]
ok2 = True
vals = []
for F, d, target in cases:
    q_, s2_, s3_, pp_, Hpp_, H2pp_ = invariants(F, d)
    c1_ = -q_
    c2_ = (s2_ - pp_) * sp.Rational(1, 2)
    c3_ = (s3_ + Hpp_) * sp.Rational(1, 2)
    c4_ = q_ * (s3_ + Hpp_) * sp.Rational(1, 2) - (s2_ + pp_) * (s2_ + pp_) * sp.Rational(1, 8) - H2pp_ * sp.Rational(1, 2)
    i1, i2, i3, i4 = (integrate_S3(ci) for ci in (c1_, c2_, c3_, c4_))
    quart = sp.simplify(i4 / (2 * sp.pi**2))
    ok = i1 == 0 and i2 == 0 and i3 == 0 and sp.simplify(quart - target) == 0
    ok2 &= ok
    vals.append(f"{sp.sstr(F)}: {quart}")
check("int c1 = int c2 = int c3 = 0, and int c4/(2 pi^2) = -1/24, -3/160, -25/384, -19/40; 0 for zonal x0, x0^2, x0^3",
      ok2, "; ".join(vals))

# ------------------------------------------------------------------------------------------------------------- P3
print("P3  Example 5.3: the corrugated profiles stay off the axis for |A| < 2/pi")
sg, kk, AA = sp.symbols("sigma k A", real=True)
ok3 = True
for k in range(2, 7):
    psi = lambda s: s + AA * sp.sin(2 * k * s)
    ok3 &= sp.simplify(sp.expand_trig(sp.cos(psi(sp.pi - sg)) + sp.cos(psi(sg)))) == 0
sym_ok = ok3


def rho_grid(k, A, n):
    # cumulative composite Simpson on [0, pi/2] with n (even) panels
    hstep = (math.pi / 2) / n
    f = [math.cos(i * hstep + A * math.sin(2 * k * i * hstep)) for i in range(n + 1)]
    out = [0.0]
    acc = 0.0
    for i in range(0, n, 2):
        acc += hstep / 3 * (f[i] + 4 * f[i + 1] + f[i + 2])
        out.append(acc)
    return [(2 * j * hstep, out[j]) for j in range(len(out))]


worst = 1e9
for k in (2, 3, 5, 8):
    for A in (0.2, 0.4, 0.6, 0.63):
        for s, r in rho_grid(k, A, 8000)[1:]:
            ok3 &= r > 0 and r >= math.sin(s) - abs(A) * s - 1e-9
            worst = min(worst, r)
used = [0.2, 0.25, 0.3, 0.4, 0.15]
ok3 &= max(used) < 2 / math.pi
check("cos psi(pi - s) = -cos psi(s) (k = 2..6); rho >= sin sigma - |A| sigma > 0 on (0, pi/2] for |A| <= 0.63 < 2/pi",
      ok3 and sym_ok, f"min rho on the grids = {worst:.4f}; 2/pi = {2 / math.pi:.4f}; largest amplitude used in MB 0.4")

# ------------------------------------------------------------------------------------------------------------- P4
print("P4  Theorem 6.4(3): beta stays in (0, pi), folds once, and the belt has exactly one det K < 0 preimage")
ch, ep = sp.symbols("chi epsilon", positive=True)
r_ = 1 + ep * sp.cos(2 * ch)
num_sph = sp.simplify(sp.expand_trig(r_ - sp.diff(r_, ch) * sp.cos(ch) / sp.sin(ch)))
ok4 = sp.simplify(num_sph - (1 - ep + 6 * ep * sp.cos(ch)**2)) == 0
details = []
for eps in (0.25, 0.3, 0.4, 0.9):
    n = 20000
    chis = [math.pi * (i + 0.5) / n for i in range(n)]

    def beta(x):
        r = 1 + eps * math.cos(2 * x)
        rp = -2 * eps * math.sin(2 * x)
        return x - math.atan(rp / r)

    bs = [beta(x) for x in chis]
    ok4 &= all(0 < b < math.pi for b in bs)
    diffs = [bs[i + 1] - bs[i] for i in range(n - 1)]
    changes = [i for i in range(len(diffs) - 1) if (diffs[i] > 0) != (diffs[i + 1] > 0)]
    ok4 &= len(changes) == 2
    ia, ib = changes
    b_hi, b_lo = bs[ia + 1], bs[ib + 1]          # beta(a) (local max) and beta(pi - a) (local min)
    ok4 &= b_hi > b_lo
    for m in range(400):
        b = math.pi * (m + 0.5) / 400
        cross = [i for i in range(n - 1) if (bs[i] - b) * (bs[i + 1] - b) < 0]
        neg = [i for i in cross if diffs[i] < 0]
        inside = b_lo < b < b_hi
        if inside:
            ok4 &= len(cross) == 3 and len(neg) == 1
        elif abs(b - b_lo) > 1e-6 and abs(b - b_hi) > 1e-6:
            ok4 &= len(cross) == 1 and len(neg) == 0
    details.append(f"eps={eps}: belt ({b_lo:.4f}, {b_hi:.4f})")
check("k_sph numerator 1 - eps + 6 eps cos^2 chi > 0; beta in (0, pi); two turning points; "
      "3 preimages (one with det K < 0) inside the belt, 1 outside", ok4, "; ".join(details))

# ------------------------------------------------------------------------------------------------------------- P5
print("P5  Proposition 7.4: the radial graphs are smooth at the poles in the (zeta, u) description")
cc = sp.symbols("c", real=True)
rr = 1 - ep + 2 * ep * cc**2
ok5 = sp.expand(sp.expand_trig(1 + ep * sp.cos(2 * ch)) - (1 + ep * (2 * sp.cos(ch)**2 - 1))) == 0
Z = sp.expand(rr * cc)
u = sp.expand(rr**2 * (1 - cc**2))
dZ = sp.diff(Z, cc)
ok5 &= sp.expand(dZ - (1 - ep + 6 * ep * cc**2)) == 0
ok5 &= all(sp.expand(u.subs(cc, s)) == 0 for s in (1, -1))
ok5 &= sp.expand(sp.diff(u, cc).subs(cc, 1) + 2 * (1 + ep)**2) == 0
ok5 &= sp.expand(sp.diff(u, cc).subs(cc, -1) - 2 * (1 + ep)**2) == 0
# dZ/dc > 0 on [-1, 1] for 0 <= eps < 1: its minimum over c is 1 - eps at c = 0
ok5 &= sp.simplify(dZ.subs(cc, 0) - (1 - ep)) == 0 and sp.Poly(dZ, cc).coeff_monomial(cc**2) == 6 * ep
check("Z = (1 - eps + 2 eps c^2) c, u = r^2 (1 - c^2) polynomial in c; dZ/dc = 1 - eps + 6 eps c^2 > 0; "
      "u(+-1) = 0, du/dc(+-1) = -+2(1 + eps)^2", ok5)

# ------------------------------------------------------------------------------------------------------------- P6
print("P6  the algebra of Theorem 5.5(2)'s proof")
m = sp.symbols("m11 m12 m13 m22 m23")
M1 = sp.Matrix([[m[0], m[1], m[2]], [m[1], m[3], m[4]], [m[2], m[4], -m[0] - m[3]]])
n2 = sp.symbols("n11 n12 n13 n22 n23 n33")
M2 = sp.Matrix([[n2[0], n2[1], n2[2]], [n2[1], n2[3], n2[4]], [n2[2], n2[4], n2[5]]])
e6 = sp.expand(sp.series(sp.sqrt((sp.eye(3) + t * M1 + t**2 * M2).det()), t, 0, 3).removeO())
ok6 = sp.expand(e6 - (1 + t**2 * (M2.trace() / 2 - (M1 * M1).trace() / 4))) == 0
e7 = sp.expand(sp.series(sp.sqrt((sp.eye(3) + t * M1).det()), t, 0, 3).removeO())
ok6 &= sp.expand(e7 - (1 - t**2 * (M1 * M1).trace() / 4)) == 0
mu = sp.symbols("mu")
d2S = -(mu / 2 + 2)                               # -<G', h> + 2<G, h o h> with G = -gamma, G' = nabla*nabla h/2
intR2 = sp.Rational(1, 2) * d2S + sp.Rational(6, 4)  # t^2 coefficient of int R dV_t, minus R_0 times dV's t^2 term
ok6 &= sp.expand(intR2 + (mu - 2) / 4) == 0
Q = (intR2 - (mu + 2)**2 / 4 - 1) / 4
ok6 &= sp.expand(Q + (mu + 2) * (mu + 3) / 16) == 0
check("sqrt det expansions; int R^(2) = -(mu - 2)/4; Q = -(mu + 2)(mu + 3)/16", ok6)

print(f"\n[{time.time() - T0:.0f}s]")
if FAILURES:
    print("test_353: FAILURES:", FAILURES)
    sys.exit(1)
print("test_353: ALL PASS")