#!/usr/bin/env python3
"""
test_340 -- SECOND ORDER AT l = 1 ON NON-ROUND BASES: THE DRIFT FAMILY, ITS ROOT, AND THE MULTIPLICITIES

Base g0 = A(r)^2 dr^2 + Phi(r)^2 (dtheta^2 + sin^2 theta dphi^2) on the positive branch. A drift of the shell centres,
x = (c(r) e_3 + Phi n, Z), with radial modulation 2 eps A^2 cos(theta), has c' = eps A^2/Phi' and, in shell coordinates,
    g_rr = A^2 (1 + 2 eps cos theta) + eps^2 mu,   g_{r theta} = eps beta sin theta,
    beta_pred = -Phi A^2/Phi',   mu_pred = A^4/Phi'^2.
Write beta = s beta_pred. Two parametrisations are run:
  (F) beta(r) = s beta_pred(r), mu(r) = mu_pred(r) as functions of r: at s = 1 this is the drifting hypersurface;
  (C) beta and mu frozen at their values at r0 (the parametrisation of the earlier exploratory runs).
f_X(s) is the O(eps^2) part, at the point, of the Codazzi component X of K~ = sqrt(det M) M^{-1}, computed with the full
nonlinear chain in exact jets (ratjets.py). The four components: C_{r th r} and C_{th ph ph} (the (Lambda - 2) families),
C_{r th th} and C_{r ph ph} (the D family).

PRE-REGISTERED, for (F) and for (C) separately:
  P1 For every s the O(eps) part vanishes, and f_X(s) is quadratic in s: f(3) = 3 f(2) - 3 f(1) + f(0).
  P2 f_X(1) = 0 for all four components, at eight base-radius pairs (B2, B3, B4 at r0 = 1/2, 1/3, 3/4, those on the
     positive branch) and, at B2, r0 = 1/2, at three angular points.
  P3 The (Lambda - 2) families have a double root at s = 1 (f(0) = f(2) != 0); the D family a simple one
     (f(0) != f(2)), so the zero is transversal.
  P4 The second root of the D family, s2 = 2 f(0)/(f(0) + f(2)), takes different values on different bases.
  P5 On constant-ratio bases (D == 0): BC (A = 2, Phi = 2r - r^2) at r0 = 1/4, 1/3, 1/2, 3/4 and BC2 (A = 4r,
     Phi = 2 r^2 - r^4) at 1/2, 3/4: the D slots still have a simple root at s = 1 and the (Lambda - 2) slots a double
     root.
  P6 (C only) Rescaling mu by 2, 1/2 or 0 at B2, r0 = 1/2, leaves every f_X(s) unchanged.

FIRST RUN, AND THE CHANGES. (C): P2-P6 passed; P1 failed. (F): P2 passed; P1, P3, P4 and P5 failed.
  - P1 failed in both modes for an implementation reason: values were passed through sympy's nsimplify, which can
    replace an algebraic number by a different one agreeing to ~15 digits, so the exact third difference came out
    ~1e-17 instead of 0 at some radii (B2 and B4 at r0 = 3/4). nsimplify is removed here and from ratjets.py; a value
    counts as zero if it simplifies to 0 or, failing that, is below 1e-45 at 60 digits (exact expressions, evaluated).
  - P3-P5 failed in mode (F) for a mathematical reason: with beta and mu as functions of r (the drift itself), the D
    family has a DOUBLE root at s = 1, like the other two (s2 = 1 everywhere). The simple root, the transversal zero
    and the base-dependent second root belong to the frozen parametrisation (C), which drops beta'(r0); beta' enters
    only the D family. Re-posed:
  P3' (F) All four components have a double root at s = 1 (f(0) = f(2) != 0) at every configuration.  (C) as P3.
  P4' (C only) as P4.
  P5' (F) On the constant-ratio bases all four components have a double root at s = 1.  (C) as P5.
  P7  (added with the re-posing) The (Lambda - 2) components agree exactly between (F) and (C) for every s; the D
      components differ for s != 1.
SECOND RUN: everything passed except P7, which was mis-registered: at s = 0 both parametrisations have beta = 0 and
coincide (mu does not enter, P6), so the D components cannot differ there. A diagnostic run (B3 1/2, B4 1/3, BC 1/2,
BC2 1/2; s = 0, 2, 3) showed exact agreement of the (Lambda - 2) components everywhere and of the D components at s = 0
only. Re-posed:
  P7' The (Lambda - 2) components agree between (F) and (C) for every s; the D components differ for s = 2 and s = 3,
      at every configuration (constant-ratio bases included).

This script replaces test 105, t107d-t107h and t109 for the paper's verification bundle. Those were records (one
configuration per run, results written to files, no exit code) in the frozen parametrisation (C).
"""
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)


def is0(e):
    e = sp.sympify(e)
    if e == 0 or sp.simplify(e) == 0:
        return True
    return abs(sp.N(e, 60)) < sp.Float("1e-45")


r = sp.Symbol("r", positive=True)
BASES = {"B2": (1 + r**2 / 4, r - r**3 / 12), "B3": (2 + r**2, r + r**2 / 5), "B4": (1 + r**2 / 2, r - r**3 / 8),
         "BC": (sp.Integer(2), 2 * r - r**2), "BC2": (4 * r, 2 * r**2 - r**4)}
COMPS = {"C_rthr": (0, 1, 0), "C_thphph": (1, 2, 2), "C_rthth": (0, 1, 1), "C_rphph": (0, 2, 2)}
LAM2, DFAM = ("C_rthr", "C_thphph"), ("C_rthth", "C_rphph")
PT1, PT2, PT3 = (R(3, 5), R(4, 5)), (R(5, 13), R(12, 13)), (R(8, 17), R(15, 17))


def on_branch(Af, Ff, r0):
    A0, F0, Ap = Af.subs(r, r0), Ff.subs(r, r0), sp.diff(Af, r).subs(r, r0)
    Fp, Fpp = sp.diff(Ff, r).subs(r, r0), sp.diff(Ff, r, 2).subs(r, r0)
    return (A0**2 - Fp**2) > 0 and (Ap * Fp - A0 * Fpp) / (A0**3 * F0) > 0


def fvals(Af, Ff, r0, s, mode, point, muscale=1):
    J.set_point(*point)
    bpred = -Ff * Af**2 / sp.diff(Ff, r)
    mpred = Af**4 / sp.diff(Ff, r)**2
    if mode == "F":
        bj = J.poly_x(J.taylor(s * bpred, r, r0))
        mj = J.poly_x(J.taylor(muscale * mpred, r, r0), 2)
    else:
        bj = J.const(sp.radsimp(s * bpred.subs(r, r0)))
        mj = J.zero()
        mj[2][0][0] = sp.radsimp(muscale * mpred.subs(r, r0))
    A_, F_ = J.poly_x(J.taylor(Af, r, r0)), J.poly_x(J.taylor(Ff, r, r0))
    cos_t, sin_t = J.cos_u(), J.sin_u()
    A2 = J.mul(A_, A_)
    g = [[J.zero() for _ in range(3)] for _ in range(3)]
    g[0][0] = J.add(J.add(A2, J.at_eps(J.smul(2, J.mul(A2, cos_t)), 1)), mj)
    g[0][1] = g[1][0] = J.at_eps(J.mul(bj, sin_t), 1)
    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)
    o1 = all(sp.simplify(J.val(ch["cod"](*idx), 1)) == 0 for idx in J.CODIDX)
    o2 = {name: sp.radsimp(sp.simplify(J.val(ch["cod"](*idx), 2))) for name, idx in COMPS.items()}
    return o1, o2


t0 = time.time()
CONFIGS = []
for nm in ("B2", "B3", "B4"):
    for r0 in (R(1, 2), R(1, 3), R(3, 4)):
        if on_branch(*BASES[nm], r0):
            CONFIGS.append((nm, r0, PT1))
CONFIGS += [("B2", R(1, 2), PT2), ("B2", R(1, 2), PT3)]
CR = [("BC", r0, PT1) for r0 in (R(1, 4), R(1, 3), R(1, 2), R(3, 4))] + [("BC2", r0, PT1) for r0 in (R(1, 2), R(3, 4))]
npairs = len({(c[0], c[1]) for c in CONFIGS})
print(f"base-radius pairs on the positive branch: {npairs}; constant-ratio configurations: {len(CR)}")

results = {}
for mode in ("F", "C"):
    for (nm, r0, pt) in CONFIGS + CR:
        assert on_branch(*BASES[nm], r0), (nm, r0)
        results[(mode, nm, r0, pt)] = {s: fvals(*BASES[nm], r0, s, mode, pt) for s in (0, 1, 2, 3)}
    print(f"  mode {mode} done [{time.time() - t0:.0f}s]")

for mode in ("F", "C"):
    ok1 = ok2 = ok3 = ok5 = True
    s2vals, det2 = {}, []
    for (m, nm, r0, pt), fs in results.items():
        if m != mode:
            continue
        o1 = all(fs[s][0] for s in fs)
        f = {s: fs[s][1] for s in fs}
        quad = all(is0(f[3][X] - (3 * f[2][X] - 3 * f[1][X] + f[0][X])) for X in COMPS)
        root = all(is0(f[1][X]) for X in COMPS)
        dbl = lambda fams: all(is0(f[0][X] - f[2][X]) and not is0(f[0][X]) for X in fams)
        smp = lambda fams: all(not is0(f[0][X] - f[2][X]) for X in fams)
        ok1 &= o1 and quad
        ok2 &= root
        if mode == "F":
            good = dbl(LAM2) and dbl(DFAM)
        else:
            good = dbl(LAM2) and smp(DFAM)
        if nm in ("BC", "BC2"):
            ok5 &= good and root
        else:
            ok3 &= good
            if pt == PT1 and mode == "C":
                X = "C_rthth"
                s2vals[(nm, r0)] = sp.radsimp(sp.simplify(2 * f[0][X] / (f[0][X] + f[2][X])))
        if not (root and quad):
            det2.append(f"{nm} r0={r0} pt={pt}: root {root}, quadratic {quad}")
    check(f"P1 ({mode}) O(eps) zero for every s; f_X(s) quadratic", ok1, "; ".join(det2[:3]))
    check(f"P2 ({mode}) f_X(1) = 0 for all four components, {npairs} base-radius pairs and three angular points", ok2)
    if mode == "F":
        check("P3' (F) all four families have a double root at s = 1, every configuration", ok3)
        check("P5' (F) constant-ratio bases: all four families double root at s = 1", ok5)
    else:
        check("P3 (C) (Lambda-2) families double root, D family simple root", ok3)
        s2num = {k: sp.N(v, 30) for k, v in s2vals.items()}
        distinct = all(not is0(s2vals[a] - s2vals[b]) for a in s2vals for b in s2vals if a < b)
        check("P4' (C) the D family's second root differs between bases", distinct,
              ", ".join(f"{k[0]}@{k[1]}: {float(v):.6f}" for k, v in s2num.items()))
        check("P5 (C) constant-ratio bases: D slots simple root, (Lambda-2) slots double root, root at s = 1", ok5)

print("P7  the two parametrisations: beta' enters only the D family")
ok7 = True
for (m, nm, r0, pt), fs in results.items():
    if m != "F":
        continue
    fc = results[("C", nm, r0, pt)]
    for sv in (0, 1, 2, 3):
        ok7 &= all(is0(fs[sv][1][X] - fc[sv][1][X]) for X in LAM2)
        if sv not in (0, 1):
            ok7 &= all(not is0(fs[sv][1][X] - fc[sv][1][X]) for X in DFAM)
check("P7' (Lambda-2) components agree between (F) and (C); D components differ for s = 2, 3 (every configuration)", ok7)

print("P6  (C) the eps^2 radial companion does not enter")
ok6 = True
nm, r0 = "B2", R(1, 2)
for ms in (2, R(1, 2), 0):
    for s in (0, 1, 2):
        _, f_ms = fvals(*BASES[nm], r0, s, "C", PT1, muscale=ms)
        ok6 &= all(is0(f_ms[X] - results[("C", nm, r0, PT1)][s][1][X]) for X in COMPS)
check("P6 (C) mu scaled by 2, 1/2, 0: every f_X(s) unchanged", ok6)

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