#!/usr/bin/env python3
"""
test_337 -- SHELL METRICS AT FIRST ORDER: THE THREE FAMILIES, THE PREFACTOR D, AND THE KERNEL

Shell metric g0 = A(r)^2 dr^2 + Phi(r)^2 (dtheta^2 + sin^2 theta dphi^2) on the positive branch
Q = (A^2 - Phi'^2)/(A^2 Phi^2) > 0, P = (A' Phi' - A Phi'')/(A^3 Phi) > 0. Deformation A -> A(1 + eps Y), Y an axisymmetric
spherical harmonic of degree l, Lambda = l(l+1). C = the O(eps) part of the Codazzi tensor of K~ = sqrt(det M) M^{-1},
computed with the full nonlinear chain in exact jets (ratjets.py) about rational points; no closed form is used upstream
of a comparison. Z' = sqrt(A^2 - Phi'^2), T_ab[Y] = (hess Y - (1/2) gamma^ Lap Y)_ab on the unit S^2.

PRE-REGISTERED:
  P1 Background: M0 = diag(Q, P, P) and K~0 = diag(P/sqrt Q, sqrt Q, sqrt Q) at every configuration.
  P2 First family: C_{r theta r} = -((Lambda - 2)/2) A^3/(Phi Z') d_theta Y   (l = 2, 3).
  P3 Second family: C_{theta phi phi} = -((Lambda - 2)/2) Z'/(A Phi P) sin^2(theta) d_theta Y   (l = 2, 3).
  P4 Third family: C_{r theta theta} = D T_{theta theta}[Y] and C_{r phi phi} = D T_{phi phi}[Y], with
         D = -d_r(sqrt Q/P) + (Phi'/Phi)(sqrt Q/P) - A Phi'/Z'   (the closed form)
           = (1/k_mer) d_r ln(k_mer/k_sph),  k_sph = sqrt Q, k_mer = P/sqrt Q   (the curvature form);
     D measured is the same for l = 2 and l = 3 and at two angular points.
  P5 The other five components (an odd number of phi indices) vanish.
  P6 Kernel: for l = 0 and l = 1 every component vanishes, on every base; for l = 2, 3 the first family does not.
  P7 Two further first-order kernel directions, on non-round bases: (a) the l = 1 deformation with an r-dependent
     amplitude, A -> A(1 + eps y(r) cos theta), y = 1 + r + r^2; (b) the off-diagonal term
     g_{r theta} = eps beta(r) sin theta, beta = 1 + r^2/3, alone. Every O(eps) component vanishes in both.
  P8 Symbolic, general base: the closed form of D equals the curvature form; D = 0 on the round base; D = 0
     identically on the constant-ratio family Phi'^2 = A^2 (1 - C0^2 Phi^(2c)).
  P9 D = 0 measured on a constant-ratio base (A = 2, Phi = 2r - r^2, k_mer/k_sph = 1/2), and D != 0 on the others.

FIRST RUN: P1-P7 and P9 passed. P8 failed for two implementation reasons, the claims unchanged: the substitution of
symbols for A(r), Phi(r) and their derivatives was passed as a dict, and A(r) was replaced before its derivatives (so
they became 0); and on the round base sympy left cot r - cos r |sin r|/sin^2 r, which is 0 only for 0 < r < pi.
Re-run with an ordered substitution (derivatives first) and the round base evaluated at four points of (0, pi).

BASES: round (A = 1, Phi = sin r at cos r0 = 3/5), B2 (A = 1 + r^2/4, Phi = r - r^3/12), B3 (A = 2 + r^2,
Phi = r + r^2/5), B4 (A = 1 + r^2/2, Phi = r - r^3/8), BC (A = 2, Phi = 2r - r^2).

LIBRARY CHANGE (29 Oct, after the first passing run): ratjets.py and this script no longer use sympy's nsimplify,
which can replace an algebraic number by a different one agreeing to ~15 digits (found in test_340). Values are
now kept exact (radsimp only). Re-run after the change: identical verdicts.
"""
import sys
import time
from pathlib import Path
import sympy as sp
from sympy import Rational as R

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "libraries"))
import ratjets as J

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)


r = sp.Symbol("r", positive=True)
BASES = {
    "round": (sp.Integer(1), sp.sin(r), sp.atan(R(4, 3))),
    "B2": (1 + r**2 / 4, r - r**3 / 12, R(1, 2)),
    "B2'": (1 + r**2 / 4, r - r**3 / 12, R(3, 4)),
    "B3": (2 + r**2, r + r**2 / 5, R(1, 2)),
    "B4": (1 + r**2 / 2, r - r**3 / 8, R(1, 2)),
    "BC": (sp.Integer(2), 2 * r - r**2, R(1, 2)),
}
P1PT, P2PT = (R(3, 5), R(4, 5)), (R(5, 13), R(12, 13))


def harmonic(ell):
    c = J.cos_u()
    if ell == 0:
        return J.const(1), 0
    if ell == 1:
        return c, 2
    if ell == 2:
        return J.smul(R(1, 2), J.sub(J.smul(3, J.mul(c, c)), J.const(1))), 6
    if ell == 3:
        return J.smul(R(1, 2), J.sub(J.smul(5, J.mul(c, J.mul(c, c))), J.smul(3, c))), 12
    raise ValueError(ell)


def base_data(Af, Ff, r0):
    A0 = sp.radsimp(sp.simplify(Af.subs(r, r0)))
    F0 = sp.radsimp(sp.simplify(Ff.subs(r, r0)))
    Ap = sp.radsimp(sp.simplify(sp.diff(Af, r).subs(r, r0)))
    Fp = sp.radsimp(sp.simplify(sp.diff(Ff, r).subs(r, r0)))
    Fpp = sp.radsimp(sp.simplify(sp.diff(Ff, r, 2).subs(r, r0)))
    Q = (A0**2 - Fp**2) / (A0**2 * F0**2)
    P = (Ap * Fp - A0 * Fpp) / (A0**3 * F0)
    return dict(A0=A0, F0=F0, Fp=Fp, Q=Q, P=P, Zp=sp.sqrt(A0**2 - Fp**2))


def D_forms(Af, Ff):
    sQ = sp.sqrt((Af**2 - sp.diff(Ff, r)**2) / (Af**2 * Ff**2))
    Pf = (sp.diff(Af, r) * sp.diff(Ff, r) - Af * sp.diff(Ff, r, 2)) / (Af**3 * Ff)
    u = sQ / Pf
    closed = -sp.diff(u, r) + (sp.diff(Ff, r) / Ff) * u - Af * sp.diff(Ff, r) / sp.sqrt(Af**2 - sp.diff(Ff, r)**2)
    kmer, ksph = Pf / sQ, sQ
    curv = sp.diff(sp.log(kmer / ksph), r) / kmer
    return closed, curv


def run(Af, Ff, r0, Yjet=None, ell=None, ydep=None, beta=None, point=P1PT):
    """O(eps) Codazzi components (all nine) for A -> A(1 + eps Y) [or y(r) Y], or the off-diagonal beta term."""
    J.set_point(*point)
    Ax, Fx = J.taylor(Af, r, r0), J.taylor(Ff, r, r0)
    A_, F_ = J.poly_x(Ax), J.poly_x(Fx)
    sin_t = J.sin_u()
    g = [[J.zero() for _ in range(3)] for _ in range(3)]
    A2 = J.mul(A_, A_)
    if Yjet is not None:
        amp = Yjet if ydep is None else J.mul(Yjet, J.poly_x(J.taylor(ydep, r, r0)))
        onep = J.add(J.const(1), J.at_eps(amp, 1))
        g[0][0] = J.mul(A2, J.mul(onep, onep))
    else:
        g[0][0] = A2
    if beta is not None:
        bt = J.at_eps(J.mul(J.poly_x(J.taylor(beta, r, r0)), sin_t), 1)
        g[0][1] = g[1][0] = bt
    F2 = J.mul(F_, F_)
    g[1][1] = F2
    g[2][2] = J.mul(F2, J.mul(sin_t, sin_t))
    ch = J.chain(g)
    comps = {idx: sp.radsimp(sp.simplify(J.val(ch["cod"](*idx), 1))) for idx in J.CODIDX}
    M0 = [[J.val(ch["M"][i][j], 0) for j in range(3)] for i in range(3)]
    K0 = [J.val(ch["Kmix"][i][i], 0) for i in range(3)]
    return comps, M0, K0


t_start = time.time()
ok1, ok2, ok3, ok4, ok5, ok6 = True, True, True, True, True, True
det = {k: [] for k in range(1, 10)}
Dmeas = {}
for name, (Af, Ff, r0) in BASES.items():
    bd = base_data(Af, Ff, r0)
    assert bd["Q"] > 0 and bd["P"] > 0, f"{name}: off the positive branch"
    Dclosed, Dcurv = D_forms(Af, Ff)
    Dc = sp.radsimp(sp.simplify(Dclosed.subs(r, r0)))
    Dk = sp.radsimp(sp.simplify(Dcurv.subs(r, r0)))
    for ell in (0, 1, 2, 3):
        for point in ((P1PT, P2PT) if (name == "B2" and ell == 2) else (P1PT,)):
            J.set_point(*point)
            Y, Lam = harmonic(ell)
            dY = J.du(Y)
            T = J.add(J.du(dY), J.smul(R(Lam, 2), Y))           # T_thth = Y'' + (Lam/2) Y on the unit S^2
            Y0, dY0, T0 = J.val(Y, 0), J.val(dY, 0), J.val(T, 0)
            s2 = point[1]**2
            comps, M0, K0 = run(Af, Ff, r0, Yjet=Y, point=point)
            ok1 &= (sp.simplify(M0[0][0] - bd["Q"]) == 0 and sp.simplify(M0[1][1] - bd["P"]) == 0
                    and sp.simplify(M0[2][2] - bd["P"]) == 0 and all(M0[i][j] == 0 for i in range(3) for j in range(3) if i != j)
                    and sp.simplify(K0[0] - bd["P"] / sp.sqrt(bd["Q"])) == 0 and sp.simplify(K0[1] - sp.sqrt(bd["Q"])) == 0)
            others = [(0, 1, 2), (0, 2, 0), (0, 2, 1), (1, 2, 0), (1, 2, 1)]
            ok5 &= all(comps[k] == 0 for k in others)
            if ell in (0, 1):
                ok6 &= all(v == 0 for v in comps.values())
                continue
            f1 = -R(Lam - 2, 2) * bd["A0"]**3 / (bd["F0"] * bd["Zp"]) * dY0
            f2 = -R(Lam - 2, 2) * bd["Zp"] / (bd["A0"] * bd["F0"] * bd["P"]) * s2 * dY0
            ok2 &= sp.simplify(comps[(0, 1, 0)] - f1) == 0 and comps[(0, 1, 0)] != 0
            ok3 &= sp.simplify(comps[(1, 2, 2)] - f2) == 0
            Dm = sp.radsimp(sp.simplify(comps[(0, 1, 1)] / T0))
            ok4 &= (sp.simplify(Dm - Dc) == 0 and sp.simplify(Dm - Dk) == 0
                    and sp.simplify(comps[(0, 2, 2)] + s2 * comps[(0, 1, 1)]) == 0)
            Dmeas.setdefault(name, set()).add(Dm)
            ok6 &= comps[(0, 1, 0)] != 0
    det[4].append(f"{name}: D = {Dc}")
ok4 &= all(len(v) == 1 for v in Dmeas.values())
check("P1 background M0 = diag(Q, P, P), K~0 = diag(P/sqrt Q, sqrt Q, sqrt Q) on every base", ok1)
check("P2 first family C_{r th r} = -((Lam-2)/2) A^3/(Phi Z') dY  (l = 2, 3)", ok2)
check("P3 second family C_{th ph ph} = -((Lam-2)/2) Z'/(A Phi P) sin^2 th dY  (l = 2, 3)", ok3)
check("P4 third family = D T[Y]; D measured = closed form = curvature form; l- and angle-independent", ok4,
      "; ".join(det[4]))
check("P5 the five components with an odd number of phi indices vanish", ok5)
check("P6 kernel: l = 0 and l = 1 give C = 0 on every base; l = 2, 3 do not", ok6)

print("P7  two further first-order kernel directions")
ok7, det7 = True, []
for name in ("B2", "B3", "B4"):
    Af, Ff, r0 = BASES[name]
    J.set_point(*P1PT)
    ca, _, _ = run(Af, Ff, r0, Yjet=J.cos_u(), ydep=1 + r + r**2)
    cb, _, _ = run(Af, Ff, r0, beta=1 + r**2 / 3)
    za, zb = all(v == 0 for v in ca.values()), all(v == 0 for v in cb.values())
    ok7 &= za and zb
    det7.append(f"{name}: {za}, {zb}")
check("P7 (a) l = 1 with r-dependent amplitude, (b) off-diagonal beta(r) sin(theta): all O(eps) components vanish",
      ok7, "; ".join(det7))

print("P8  symbolic, general base")
Af_, Ff_ = sp.Function("A")(r), sp.Function("Phi")(r)
Dcl, Dcv = D_forms(Af_, Ff_)
# compare on the positive branch: write A, Phi and derivatives as positive symbols with Q, P > 0
syms = sp.symbols("a0 a1 a2 f0 f1 f2 f3 a3", real=True)
a0, a1, a2, f0, f1, f2, f3, a3 = syms
rep = [(sp.Derivative(Af_, (r, 2)), a2), (sp.Derivative(Ff_, (r, 3)), f3), (sp.Derivative(Ff_, (r, 2)), f2),
       (sp.Derivative(Af_, r), a1), (sp.Derivative(Ff_, r), f1), (Af_, a0), (Ff_, f0)]     # derivatives first
diffD = (Dcl - Dcv).subs(rep)
vals = [{a0: R(3, 2), a1: R(1, 3), a2: R(-1, 5), f0: R(7, 10), f1: R(1, 2), f2: R(-2, 3), f3: R(1, 7)},
        {a0: 2, a1: R(-1, 4), a2: R(1, 2), f0: R(1, 3), f1: R(3, 2), f2: R(-1, 2), f3: R(-2, 5)}]
ok8 = all(abs(sp.N(diffD.subs(v), 40)) < 1e-35 for v in vals)
Dround = Dcl.subs({Af_: 1, Ff_: sp.sin(r)}).doit()
ok8 &= all(abs(sp.N(Dround.subs(r, rv), 40)) < 1e-35 for rv in (R(1, 3), R(1), R(2), R(3)))   # 0 < r < pi
c, C0s = sp.symbols("c C0", positive=True)
# constant ratio family: Phi' = A sqrt(1 - C0^2 Phi^(2c)); differentiate to get Phi'', Phi''' in terms of A, A', Phi
fp_rel = Af_ * sp.sqrt(1 - C0s**2 * Ff_**(2 * c))
fpp_rel = sp.diff(fp_rel, r).subs(sp.Derivative(Ff_, r), fp_rel)
fppp_rel = sp.diff(sp.diff(fp_rel, r), r)
fppp_rel = fppp_rel.subs(sp.Derivative(Ff_, (r, 2)), fpp_rel).subs(sp.Derivative(Ff_, r), fp_rel)
Dcr = Dcv.subs(sp.Derivative(Ff_, (r, 3)), fppp_rel).subs(sp.Derivative(Ff_, (r, 2)), fpp_rel).subs(sp.Derivative(Ff_, r), fp_rel)
Dcr = Dcr.subs(rep)
vals_cr = [{a0: R(3, 2), a1: R(1, 3), a2: R(-1, 5), f0: R(1, 2), c: R(1, 2), C0s: R(3, 4)},
           {a0: 2, a1: R(-1, 4), a2: R(1, 2), f0: R(1, 3), c: R(3, 2), C0s: R(1, 2)}]
ok8 &= all(abs(sp.N(Dcr.subs(v), 40)) < 1e-35 for v in vals_cr)
check("P8 closed form = curvature form (general base); D = 0 on the round base and on the constant-ratio family", ok8)

print("P9  D on a constant-ratio base and on the others")
ok9 = Dmeas["BC"] == {0} and all(0 not in Dmeas[k] for k in ("B2", "B2'", "B3", "B4")) and Dmeas["round"] == {0}
check("P9 D = 0 on the round base and BC (k_mer/k_sph = 1/2); D != 0 on B2, B3, B4", ok9,
      ", ".join(f"{k}: {list(v)[0]}" for k, v in Dmeas.items()))

print(f"\n[{time.time() - t_start:.0f}s]")
print(f"test_337: {'ALL PASS' if not FAILURES else 'FAILURES: ' + ', '.join(FAILURES)}")
sys.exit(1 if FAILURES else 0)