"""Round S3 geometry in arXiv:1212.3388v3 Appendix A conventions.

Run with Python 3.13.7, SymPy 1.13.3 and mpmath 1.3.0:
    python sphere_geometry.py
    python sphere_geometry.py --wrong-connection-sign

No repository imports, network access or data files. Exact symbolic checks
are separate from finite-step RK4 diagnostics. Indices in code start at zero.
The normal vector y and transported vectors are ambient four-columns;
Omega[a,b,c] is a connection coefficient in an orthonormal three-frame.
"""

import argparse
import json
from itertools import product

import mpmath as mp
import sympy as sp


class Checks:
    def __init__(self):
        self.counts = {}

    def equal(self, actual, expected, group):
        residual = actual - expected
        entries = list(residual) if isinstance(residual, sp.MatrixBase) else [residual]
        for entry in entries:
            assert sp.simplify(sp.trigsimp(entry)) == 0, (group, entry)
        # One grouped assertion, not one count per matrix entry.
        self.counts[group] = self.counts.get(group, 0) + 1


def ambient_frame():
    y = sp.Matrix(sp.symbols("y0:4", real=True))
    M = sp.Matrix([[-y[1], y[0], -y[3], y[2]],
                   [-y[2], y[3], y[0], -y[1]],
                   [-y[3], -y[2], y[1], y[0]]])
    return y, M


def curvature(omega, de):
    """Return F[a,b,c,d] = (dOmega+Omega wedge Omega)[a,b](E_c,E_d)."""
    return {(a, b, c, d): sp.simplify(sum(
        omega[a, b, f] * de[f, c, d]
        + omega[a, f, c] * omega[f, b, d]
        - omega[a, f, d] * omega[f, b, c]
        for f in range(3))) for a, b, c, d in product(range(3), repeat=4)}


def paper_ricci(F):
    # P uses R[c,d,a,b] = F[b,a,c,d] and Ric[a,c] = sum_b R[a,b,c,b].
    return sp.Matrix(3, 3, lambda a, c: sum(F[b, c, a, b] for b in range(3)))


def exact_geometry(wrong_sign=False):
    check = Checks()
    ell = sp.Symbol("ell", positive=True)
    y, M = ambient_frame()
    norm = (y.T * y)[0]
    check.equal(M * y, sp.zeros(3, 1), "global_frame")
    check.equal(M * M.T, norm * sp.eye(3), "global_frame")
    check.equal(M.T * M, norm * sp.eye(4) - y * y.T, "global_frame")

    de, bracket = {}, {}
    for a, b, c in product(range(3), repeat=3):
        ambient = sum((sp.diff(M[a, j], y[i]) - sp.diff(M[a, i], y[j]))
                      * M[b, i] * M[c, j]
                      for i, j in product(range(4), repeat=2)) / ell
        # Restriction to norm=1 is justified by the polynomial identity,
        # not by choosing one chart or a few sample points.
        de[a, b, c] = sp.cancel(ambient / norm)
        check.equal(ambient, 2 * sp.LeviCivita(a, b, c) * norm / ell,
                    "structure")
    bracket = {(c, a, b): -de[c, a, b]
               for c, a, b in product(range(3), repeat=3)}

    for a, b in product(range(3), repeat=2):
        lie = sp.Matrix(4, 1, lambda j, column: sum(
            M[a, i] * sp.diff(M[b, j], y[i])
            - M[b, i] * sp.diff(M[a, j], y[i]) for i in range(4)) / ell**2)
        reconstructed = sum((bracket[c, a, b] * M.row(c).T / ell
                             for c in range(3)), sp.zeros(4, 1))
        check.equal(lie, reconstructed, "structure")

    # Nine unknown coefficients: metric compatibility is antisymmetry in a,b.
    unknown = {(a, b, c): sp.Symbol(f"w{a}{b}{c}")
               for a in range(3) for b in range(a + 1, 3) for c in range(3)}
    omega = {(a, b, c): (0 if a == b else unknown[a, b, c] if a < b
                         else -unknown[b, a, c])
             for a, b, c in product(range(3), repeat=3)}
    torsion_equations = [de[a, b, c] + omega[a, c, b] - omega[a, b, c]
                        for a in range(3) for b in range(3)
                        for c in range(b + 1, 3)]
    solutions = sp.solve(torsion_equations, list(unknown.values()), dict=True)
    assert len(solutions) == 1 and set(solutions[0]) == set(unknown.values())
    omega = {key: sp.sympify(value).subs(solutions[0]) for key, value in omega.items()}
    for a, b, c in product(range(3), repeat=3):
        check.equal(omega[a, b, c], sp.LeviCivita(a, b, c) / ell, "connection")

    if wrong_sign:
        omega = {key: -value for key, value in omega.items()}
        residual = sp.simplify(de[0, 1, 2] + omega[0, 2, 1] - omega[0, 1, 2])
        scalar = sp.simplify(sp.trace(paper_ricci(curvature(omega, de))))
        check.equal(residual, 4 / ell, "deliberate_failure")
        check.equal(scalar, 18 / ell**2, "deliberate_failure")
        return {"mode": "expected failure detected", "torsion_123": str(residual),
                "wrong_scalar": str(scalar), "expected_scalar": str(-6 / ell**2),
                "meaning": "Negating Omega alone changes the connection; it is not a frame rotation."}

    F = curvature(omega, de)
    for a, b, c, d in product(range(3), repeat=4):
        check.equal(F[a, b, c, d],
                    (int(a == c) * int(b == d) - int(a == d) * int(b == c)) / ell**2,
                    "curvature")
    ricci = paper_ricci(F)
    check.equal(ricci, -2 * sp.eye(3) / ell**2, "curvature")
    check.equal(sp.trace(ricci), -6 / ell**2, "curvature")

    # Independent coordinates; this calculation does not use omega or de.
    chi, alpha, beta = sp.symbols("chi alpha beta", real=True)
    coordinates = [chi, alpha, beta]
    g = ell**2 * sp.diag(1, sp.cos(chi)**2, sp.sin(chi)**2)
    gi = g.inv()
    Gamma = {(a, b, c): sp.simplify(sum(gi[a, d] * (
        sp.diff(g[d, b], coordinates[c]) + sp.diff(g[d, c], coordinates[b])
        - sp.diff(g[b, c], coordinates[d])) / 2 for d in range(3)))
        for a, b, c in product(range(3), repeat=3)}
    def coordinate_R(a, b, c, d):
        return sp.diff(Gamma[d, b, c], coordinates[a]) - sp.diff(Gamma[d, a, c], coordinates[b]) + sum(
            Gamma[d, a, f] * Gamma[f, b, c] - Gamma[d, b, f] * Gamma[f, a, c]
            for f in range(3))
    ricci_coordinate = sp.Matrix(3, 3, lambda a, c: sp.trigsimp(sum(
        coordinate_R(a, b, c, b) for b in range(3))))
    check.equal(ricci_coordinate, -2 * g / ell**2, "coordinate_metric")
    volume = sp.integrate(ell**3 * sp.sin(chi) * sp.cos(chi), (chi, 0, sp.pi / 2)) * (2 * sp.pi)**2
    check.equal(volume, 2 * sp.pi**2 * ell**3, "coordinate_metric")

    # Spin lift; ordered a,b summation includes both antisymmetric pairs.
    gamma = [sp.diag(1, -1), sp.Matrix([[0, -1], [-1, 0]]), sp.Matrix([[0, sp.I], [-sp.I, 0]])]
    A = [sum((omega[a, b, c] * gamma[a] * gamma[b] / 4
              for a, b in product(range(3), repeat=2)), sp.zeros(2)) for c in range(3)]
    for a in range(3):
        check.equal(A[a], sp.I * gamma[a] / (2 * ell), "spin_lift")
    for a, b in product(range(3), repeat=2):
        spin_curvature = A[a] * A[b] - A[b] * A[a] - sum(
            (bracket[c, a, b] * A[c] for c in range(3)), sp.zeros(2))
        lifted = sum((F[c, d, a, b] * gamma[c] * gamma[d] / 4
                      for c, d in product(range(3), repeat=2)), sp.zeros(2))
        check.equal(spin_curvature, lifted, "spin_lift")
    return {"mode": "round geometry", "grouped_exact_assertions": sum(check.counts.values()),
            "assertions_by_group": check.counts,
            "de1_coefficient_of_e2_wedge_e3": str(de[0, 1, 2]),
            "Omega12_coefficient_of_e3": str(omega[0, 1, 2]),
            "curvature12_coefficient_of_e1_wedge_e2": str(F[0, 1, 0, 1]),
            "paper_Ricci_diagonal": [str(ricci[a, a]) for a in range(3)],
            "paper_scalar": str(sp.trace(ricci)), "volume": str(volume)}


def latitude_transport(cos_latitude, steps, digits=50):
    """Classical RK4 on ambient vectors, no connection/frame formula in RHS."""
    with mp.workdps(digits):
        c = mp.mpf(cos_latitude.p) / int(cos_latitude.q)
        s = mp.sqrt(1 - c**2)
        initial = mp.matrix([[-s, 0, 0], [c, 0, 0], [0, 1, 0], [0, 0, 1]])
        V = initial.copy()
        def rhs(t, vectors):
            y = mp.matrix([c, s * mp.cos(t), s * mp.sin(t), 0])
            velocity = mp.matrix([0, -s * mp.sin(t), s * mp.cos(t), 0])
            return -y * (velocity.T * vectors)
        dt = 2 * mp.pi / steps
        for n in range(steps):
            t = n * dt
            k1 = rhs(t, V)
            k2 = rhs(t + dt / 2, V + dt * k1 / 2)
            k3 = rhs(t + dt / 2, V + dt * k2 / 2)
            k4 = rhs(t + dt, V + dt * k3)
            V += dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6
        angle = -2 * mp.pi * c
        rotation = mp.matrix([[mp.cos(angle), -mp.sin(angle), 0],
                              [mp.sin(angle), mp.cos(angle), 0], [0, 0, 1]])
        target = initial * rotation
        normal = mp.matrix([c, s, 0, 0])
        maxabs = lambda array: max(abs(entry) for entry in array)
        return {"steps": steps,
                "max_ambient_error": mp.nstr(maxabs(V - target), 12),
                "max_tangency_error": mp.nstr(maxabs(normal.T * V), 12),
                "max_orthonormality_error": mp.nstr(maxabs(V.T * V - mp.eye(3)), 12)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--wrong-connection-sign", action="store_true")
    parser.add_argument("--cos-latitude", default="3/5", help="rational c, -1<c<1")
    parser.add_argument("--steps", type=int, nargs="+", default=[64, 128, 256])
    parser.add_argument("--digits", type=int, default=50)
    args = parser.parse_args()
    c = sp.Rational(args.cos_latitude)
    if not -1 < c < 1 or args.digits < 25 or any(n < 4 for n in args.steps):
        parser.error("Require -1<c<1, digits>=25, and each step count>=4.")
    result = exact_geometry(args.wrong_connection_sign)
    if not args.wrong_connection_sign:
        result["ambient_transport"] = {"cos_latitude": str(c), "decimal_digits": args.digits,
            "oriented_angle_mod_2pi": str(2 * sp.pi * (1 - c)),
            "rows": [latitude_transport(c, n, args.digits) for n in args.steps],
            "interpretation": "Finite-step RK4 errors, not uncertainty in the exact curvature."}
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()
