#!/usr/bin/env python3
"""
test_335 -- THE ALGEBRAIC REDUCTION, AND THE PRINCIPAL SYMBOL CHECKED AGAINST THE FULL OPERATOR

M = -G (Einstein endomorphism), K~ = sqrt(det M) M^{-1}, F(g) = Cod K~. The symbol formula of the paper is
    sigma(xi) h0 = xi ^ (D adj_K)^{-1} (sigma_2(DM)(xi) h0),   sigma_2(DM)(xi) h0 = -(1/2) adj_2(P h0 P),
with P = 1 - xi xi^T (|xi| = 1), adj_2(Y) = (tr Y) P - Y, D adj_K(X) = KX + XK - sigma_1 X - (tr X) K + (sigma_1 tr X - tr KX) 1,
K = K~(g)(p) in a frame orthonormal at p, and (xi ^ L)_abc = xi_a L_bc - xi_b L_ac (the formal symbol, d -> xi).

PRE-REGISTERED:
  P1 (Theorem A(1)) det adj K = (det K)^2 for every symmetric 3x3 K (symbolic).
  P2 (Theorem A(2)) adj(sqrt(det M) M^{-1}) = M for M = diag(m1, m2, m3), m1 m2 m3 > 0 (symbolic, with
     m1 = a, m2 = b, m3 = c/(a b) and c > 0, a, b of either sign), hence for every symmetric M with det M > 0.
  P3 (Theorem A(3)) K~(adj K) = sign(det K) K for every invertible symmetric K: symbolic in an eigenframe with every
     sign pattern of (k1, k2, k3), and exactly at four rational K in a rational rotated frame.
  P4 (Theorem A, convexity) adj K is positive definite iff K is definite (all eight sign patterns).
  P5 (the symbol against F) At the origin of a metric g with g(0) = 1 and prescribed M(0), for h = ((xi.x)^3/6) h0,
     the O(eps) part of F(g + eps h) at 0 equals sigma(xi) h0 computed from the formula above, exactly (all nine
     components), for M(0) = diag(1, 2, 3) (convex branch) and M(0) = diag(2, -1, -3) (indefinite branch, det > 0),
     two unit covectors and two h0 each; and sigma(xi) h0 != 0 there. (Only third derivatives of h are nonzero at
     the origin, so the O(eps) part of F there is exactly the principal symbol.)

FIRST RUN: all passed. Then, before logging, a check that the constructed metric has the prescribed M(0) was added to
P5 (it had been assumed); the second run is identical and confirms it.

This script replaces test 96 and test_227 for the paper's verification bundle.

BUNDLE EDIT (29 October, before the bundle's final run; claims unchanged): sympy's nsimplify removed from P5. On a
number with no free symbols nsimplify evaluates numerically and can return a different algebraic number agreeing to
about 15 digits; the values there involve sqrt(det M), so the risk was real. The check now compares exact
expressions directly. The run log is of the edited script.
"""
import sys
import itertools
import sympy as sp
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)


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


def symmat(v):
    S = sp.zeros(3, 3)
    for x, (i, j) in zip(v, SYM):
        S[i, j] = S[j, i] = x
    return S


def Kt(M):
    return sp.sqrt(M.det()) * M.inv()


# ---------------------------------------------------------------- P1-P4
print("P1  det adj K = (det K)^2")
Kg = symmat(sp.symbols("q0:6", real=True))
check("det adj K = (det K)^2, general symmetric K", sp.expand(Kg.adjugate().det() - Kg.det()**2) == 0)

print("P2  adj K~ = M on det M > 0")
a, b = sp.symbols("a b", real=True, nonzero=True)
c = sp.symbols("c", positive=True)
Md = sp.diag(a, b, c / (a * b))
ok2 = sp.simplify(Kt(Md).adjugate() - Md) == sp.zeros(3, 3)
check("adj(sqrt(det M) M^{-1}) = M, M = diag(a, b, c/(ab)), c > 0, a, b any sign", ok2)

print("P3  K~(adj K) = sign(det K) K")


def rat_rot(x, y, z):
    A = sp.Matrix([[0, -z, y], [z, 0, -x], [-y, x, 0]])
    return (sp.eye(3) - A) * (sp.eye(3) + A).inv()


ok3 = True
kp = sp.symbols("k1:4", positive=True)
for signs in itertools.product((1, -1), repeat=3):
    Kd = sp.diag(*[s * k for s, k in zip(signs, kp)])
    sg = signs[0] * signs[1] * signs[2]
    ok3 &= sp.simplify(Kt(Kd.adjugate()) - sg * Kd) == sp.zeros(3, 3)
Rr = rat_rot(R(2, 7), R(-3, 5), R(5, 11))
for kd in ((1, 2, 3), (2, -1, -3), (-1, -2, -5), (3, -1, 2)):
    Kr = Rr * sp.diag(*kd) * Rr.T
    ok3 &= sp.simplify(Kt(Kr.adjugate()) - sp.sign(Kr.det()) * Kr) == sp.zeros(3, 3)
check("K~(adj K) = sign(det K) K, all sign patterns (symbolic) and four rotated rational K", ok3)

print("P4  adj K > 0 iff K definite")
ok4 = True
for signs in itertools.product((1, -1), repeat=3):
    ev = sp.diag(*[s * k for s, k in zip(signs, kp)]).adjugate()
    pos = all(sp.ask(sp.Q.positive(ev[i, i])) for i in range(3))
    definite = len(set(signs)) == 1
    ok4 &= pos == definite
check("adj K positive definite iff all k_i of one sign (eight sign patterns)", ok4)

# ---------------------------------------------------------------- P5 machinery: truncated jets about the origin
X = sp.symbols("x0:3", real=True)
eps = sp.Symbol("eps")


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


def F_at_origin(g):
    """F(g)(0) for a polynomial metric g with g(0) = identity + O(eps), coefficients polynomial in eps.
    Returns the nine components C_abc, a < b, as expressions in eps (exact)."""
    g = tmat(g, 3)
    o = {x: 0 for x in X}
    g0 = g.subs(o)
    g0i = g0.inv()
    A = g - g0
    gi, term = g0i, g0i
    for _ in range(3):
        term = tmat(-g0i * A * term, 2)
        gi = gi + term
    gi = tmat(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)]
    Gam = [[[trunc(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_: trunc(sp.diff(Gam[a_][d_][b_], X[c_]) - sp.diff(Gam[a_][c_][b_], X[d_])
                                        + sum(Gam[a_][c_][e] * Gam[e][d_][b_] - Gam[a_][d_][e] * Gam[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 = trunc(sum(gi[i, j] * Ric[i, j] for i in range(3) for j in range(3)), 1)
    G = tmat(Ric - Rs * g / 2, 1)
    M = tmat(-gi * G, 1)
    M0 = M.subs(o)
    M1 = [sp.diff(M, X[k]).subs(o) for k in range(3)]
    G0 = [[[Gam[i][k][e].subs(o) for e in range(3)] for k in range(3)] for i in range(3)]
    # K~ (endomorphism) and its first derivatives at 0, differentiated along coordinate lines
    t = sp.Symbol("t")
    Kt0 = sp.sqrt(M0.det()) * M0.inv()
    dK = []
    for k in range(3):
        Mt = M0 + t * M1[k]
        dK.append((sp.sqrt(Mt.det()) * Mt.inv()).diff(t).subs(t, 0))
    # lower with g: K_low = g K (endomorphism K^i_j -> K_ij = g_im K^m_j); derivatives of the product
    g1 = [sp.diff(g, X[k]).subs(o) for k in range(3)]
    Kl0 = g0 * Kt0
    dKl = [g1[k] * Kt0 + g0 * dK[k] for k in range(3)]
    nab = lambda a_, b_, c_: dKl[a_][b_, c_] - sum(G0[e][a_][b_] * Kl0[e, c_] + G0[e][a_][c_] * Kl0[b_, e] for e in range(3))
    return [nab(a_, b_, c_) - nab(b_, a_, c_) for (a_, b_, c_) in CODIDX], M0


def metric_with_M(Mdiag, cubic):
    """g = 1 - (1/3) R_ikjl x^k x^l + cubic, with the curvature tensor of M(0) = diag(m1, m2, m3):
    sectional curvatures K_23 = m1, K_13 = m2, K_12 = m3 (in 3D, M_ii is the sectional curvature of the plane
    orthogonal to e_i)."""
    sec = {(1, 2): Mdiag[0], (0, 2): Mdiag[1], (0, 1): Mdiag[2]}
    Rt = {}
    for (i, j), kk in sec.items():
        for (p, q, s, u, sg) in ((i, j, i, j, 1), (j, i, j, i, 1), (i, j, j, i, -1), (j, i, i, j, -1)):
            Rt[(p, q, s, u)] = sg * kk                     # R_pqsu with R_ijij = K_ij
    Rf = lambda p, q, s, u: Rt.get((p, q, s, u), 0)
    g = sp.Matrix(3, 3, lambda i, j: (1 if i == j else 0)
                  - sp.Rational(1, 3) * sum(Rf(i, k, j, l) * X[k] * X[l] for k in range(3) for l in range(3)))
    return g + cubic


def sigma(K, xi, h0):
    P = sp.eye(3) - xi * xi.T
    Y = P * h0 * P
    s2 = -(Y.trace() * P - Y) / 2
    s1 = K.trace()
    z = sp.symbols("z0:6")
    Z = symmat(z)
    DA = K * Z + Z * K - s1 * Z - Z.trace() * K + (s1 * Z.trace() - (K * Z).trace()) * sp.eye(3)
    sol = sp.solve([(DA - s2)[i, j] for (i, j) in SYM], z, dict=True)[0]
    L = Z.subs(sol)
    return [xi[a_] * L[b_, c_] - xi[b_] * L[a_, c_] for (a_, b_, c_) in CODIDX]


print("P5  the symbol formula against the full operator")
cubic = sp.Matrix(3, 3, lambda i, j: 0)
cubic[0, 1] = cubic[1, 0] = X[0] * X[1] * X[2] / 7 + X[2]**3 / 11
cubic[0, 0] = X[1]**3 / 5
cubic[2, 2] = X[0]**2 * X[1] / 9
cubic[1, 2] = cubic[2, 1] = X[0]**3 / 13
XIS = [sp.Matrix([R(3, 5), R(4, 5), 0]), sp.Matrix([R(2, 3), R(1, 3), R(2, 3)])]
H0S = [symmat([1, -2, R(1, 2), R(1, 3), 0, -1]), symmat([0, 1, 1, R(-1, 2), 2, R(1, 4)])]
ok5, det5 = True, []
for Mdiag in ((1, 2, 3), (2, -1, -3)):
    g = metric_with_M(Mdiag, cubic)
    # confirm M(0)
    base, M0chk = F_at_origin(g)
    ok0 = M0chk == sp.diag(*Mdiag)
    ok5 &= ok0
    det5.append(f"M(0) = diag{Mdiag}: {ok0}")
    o = {x: 0 for x in X}
    for xi in XIS:
        for h0 in H0S:
            h = ((xi.T * sp.Matrix(X))[0])**3 / 6 * h0
            Fe, _ = F_at_origin(g + eps * h)
            lin = [sp.radsimp(sp.simplify(sp.diff(v, eps).subs(eps, 0))) for v in Fe]
            K0 = Kt(sp.diag(*Mdiag))
            sg = [sp.radsimp(sp.simplify(v)) for v in sigma(K0, xi, h0)]
            same = all(sp.simplify(u - v) == 0 for u, v in zip(lin, sg))
            ok5 &= same and any(v != 0 for v in sg)
            det5.append(f"M0 = diag{Mdiag}, xi = {list(xi)}: {'equal' if same else 'DIFFERENT'}")
check("O(eps) part of F(g + eps (xi.x)^3/6 h0) at 0 = sigma(xi) h0, both branches, nonzero", ok5, "; ".join(det5))

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