#!/usr/bin/env python3
"""
test_339 -- SECOND ORDER AT l = 1 ON THE ROUND BASE: THE OBSTRUCTION AND ITS TWO REMOVALS

Round base: A = 1, Phi = sin r (the unit S^3, gamma = dr^2 + sin^2 r (dtheta^2 + sin^2 theta dphi^2)). The full nonlinear
chain g -> M = -G -> K~ = sqrt(det M) M^{-1} -> Cod K~ is run in exact jets (ratjets.py) about points with rational
(cos r0, sin r0) and rational (cos theta0, sin theta0), carrying eps to second order. No linearisation is used.

PRE-REGISTERED:
  P1 (the obstruction) For A = 1 + eps cos(theta) exactly (g_rr = (1 + eps cos theta)^2), the O(eps) Codazzi tensor
     vanishes and the O(eps^2) part is
        C_{r th r}   = -2 sin th cos th cos^2 r / sin^4 r,     C_{r th th} = -sin^2 th cos r / sin^3 r,
        C_{r ph ph}  =  sin^4 th cos r / sin^3 r,              C_{th ph ph} = -2 sin^3 th cos th / tan^2 r,
     with the other five components zero.
  P2 (route 1) With A = 1 + eps cos(theta) + eps^2 cot^2 r [1/6 + (1/3) P_2(cos theta)], every component vanishes
     through O(eps^2).
  P3 (route 2) With g_rr = 1 + 2 eps cos(theta) + eps^2 mu(r) and g_{r theta} = eps beta(r) sin(theta), the shells
     unchanged, and beta = -s tan r: the O(eps) part vanishes for every s; the O(eps^2) part of C_{r th r} equals
     -(1 - s)^2 2 sin th cos th / sin^4 r (s = 0, 1, 2, 3), a perfect square with its double root at s = 1; at s = 1
     with mu = sec^2 r every component vanishes; and every O(eps^2) component is the same for mu = sec^2 r and
     mu = 0 (s = 0, 2).
Checked at (cos r0, sin r0) = (3/5, 4/5) and (5/13, 12/13), and (cos th0, sin th0) = (3/5, 4/5) and (8/17, 15/17).

This script, with test_338 P4, replaces tests 100 and 103 and the second-order audit scripts for the paper's bundle.

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)
RPTS = [sp.atan(R(4, 3)), sp.atan(R(12, 5))]
TPTS = [(R(3, 5), R(4, 5)), (R(8, 17), R(15, 17))]
OTHERS = [(0, 1, 2), (0, 2, 0), (0, 2, 1), (1, 2, 0), (1, 2, 1)]


def jet_of(expr, r0, eps_order=0):
    return J.poly_x(J.taylor(expr, r, r0), eps_order)


def run(r0, grr_jet, grt_jet=None):
    F_ = jet_of(sp.sin(r), r0)
    sin_t = J.sin_u()
    g = [[J.zero() for _ in range(3)] for _ in range(3)]
    g[0][0] = grr_jet
    if grt_jet is not None:
        g[0][1] = g[1][0] = grt_jet
    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)
    out = {}
    for idx in J.CODIDX:
        Cj = ch["cod"](*idx)
        out[idx] = (sp.radsimp(sp.simplify(J.val(Cj, 1))), sp.radsimp(sp.simplify(J.val(Cj, 2))))
    return out


t0 = time.time()
ok1, ok2, ok3a, ok3b, ok3c, ok3d = True, True, True, True, True, True
det1 = []
for r0 in RPTS:
    cr, sr = sp.radsimp(sp.cos(r0)), sp.radsimp(sp.sin(r0))
    for (ct, st) in TPTS:
        J.set_point(ct, st)
        cos_t, sin_t = J.cos_u(), J.sin_u()
        # P1: A = 1 + eps cos th
        A = J.add(J.const(1), J.at_eps(cos_t, 1))
        c = run(r0, J.mul(A, A))
        want = {(0, 1, 0): -2 * st * ct * cr**2 / sr**4, (0, 1, 1): -st**2 * cr / sr**3,
                (0, 2, 2): st**4 * cr / sr**3, (1, 2, 2): -2 * st**3 * ct * cr**2 / sr**2}
        ok1 &= all(v[0] == 0 for v in c.values()) and all(sp.simplify(c[k][1] - w) == 0 for k, w in want.items())
        ok1 &= all(c[k][1] == 0 for k in OTHERS)
        det1.append(f"r0 = atan({sr / cr}), th0 = acos({ct}): C_rthr = {c[(0, 1, 0)][1]}")
        # P2: route 1
        P2l = J.smul(R(1, 2), J.sub(J.smul(3, J.mul(cos_t, cos_t)), J.const(1)))
        corr = J.mul(jet_of(sp.cot(r)**2, r0), J.add(J.const(R(1, 6)), J.smul(R(1, 3), P2l)))
        A1 = J.add(A, J.at_eps(corr, 2))
        c1 = run(r0, J.mul(A1, A1))
        ok2 &= all(v[0] == 0 and v[1] == 0 for v in c1.values())
        # P3: route 2
        base_rr = J.add(J.const(1), J.at_eps(J.smul(2, cos_t), 1))
        res = {}
        for mu_expr, mu_tag in ((1 / sp.cos(r)**2, "sec2"), (sp.Integer(0), "zero")):
            for s in (0, 1, 2, 3):
                grr = J.add(base_rr, jet_of(mu_expr, r0, 2)) if mu_tag == "sec2" else base_rr
                grt = J.at_eps(J.mul(jet_of(-s * sp.tan(r), r0), sin_t), 1)
                res[(mu_tag, s)] = run(r0, grr, grt)
        for (mu_tag, s), cc in res.items():
            ok3a &= all(v[0] == 0 for v in cc.values())
            ok3b &= sp.simplify(cc[(0, 1, 0)][1] + (1 - s)**2 * 2 * st * ct / sr**4) == 0
        ok3c &= all(v[1] == 0 for v in res[("sec2", 1)].values())
        ok3d &= all(sp.simplify(res[("sec2", s)][k][1] - res[("zero", s)][k][1]) == 0 for s in (0, 2) for k in J.CODIDX)
check("P1 the O(eps^2) obstruction: the four displayed components, the rest zero, O(eps) zero", ok1, "; ".join(det1[:2]))
check("P2 route 1 (the tilt's second-order correction) removes every component through O(eps^2)", ok2)
check("P3 route 2: O(eps) zero for every s", ok3a)
check("P3 route 2: C_{r th r} = -(1 - s)^2 2 sin th cos th / sin^4 r, double root at s = 1", ok3b)
check("P3 route 2: s = 1, mu = sec^2 r: every component vanishes", ok3c)
check("P3 route 2: the O(eps^2) residual does not depend on mu (sec^2 r against 0, s = 0, 2)", ok3d)

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