#!/usr/bin/env python3
"""Exact checks for the generated-solution admissibility page.

The main generated benchmark is a normalized Ehlers transformation from a
Schwarzschild seed to Taub--NUT.  A spinning conical quotient then supplies an
adversarial local-versus-global test: it is locally Minkowski and has the
canonical Killing-block determinant, but its quotient data can leave a
conical axis, a nonclosing degenerating direction, and closed timelike
circles.  A Ricci-flat plane wave supplies the independent scalar-invariant
counterexample.
"""

from __future__ import annotations

import sympy as sp


def require_zero(label: str, expression: sp.Expr) -> None:
    simplified = sp.factor(sp.cancel(sp.simplify(expression)))
    if simplified != 0:
        raise AssertionError(f"{label} did not vanish: {simplified}")


def require_zero_matrix(label: str, matrix: sp.Matrix) -> None:
    for row in range(matrix.rows):
        for column in range(matrix.cols):
            require_zero(f"{label}[{row},{column}]", matrix[row, column])


def connection_and_curvature(
    metric: sp.Matrix, coordinates: tuple[sp.Symbol, ...]
) -> tuple[
    list[list[list[sp.Expr]]],
    list[list[list[list[sp.Expr]]]],
    list[list[sp.Expr]],
]:
    dimension = len(coordinates)
    inverse = sp.simplify(metric.inv())
    christoffel = [
        [
            [
                sp.simplify(
                    sum(
                        inverse[a, d]
                        * (
                            sp.diff(metric[d, c], coordinates[b])
                            + sp.diff(metric[d, b], coordinates[c])
                            - sp.diff(metric[b, c], coordinates[d])
                        )
                        / 2
                        for d in range(dimension)
                    )
                )
                for c in range(dimension)
            ]
            for b in range(dimension)
        ]
        for a in range(dimension)
    ]
    riemann = [
        [
            [
                [
                    sp.simplify(
                        sp.diff(christoffel[a][b][d], coordinates[c])
                        - sp.diff(christoffel[a][b][c], coordinates[d])
                        + sum(
                            christoffel[a][e][c] * christoffel[e][b][d]
                            - christoffel[a][e][d] * christoffel[e][b][c]
                            for e in range(dimension)
                        )
                    )
                    for d in range(dimension)
                ]
                for c in range(dimension)
            ]
            for b in range(dimension)
        ]
        for a in range(dimension)
    ]
    ricci = [
        [
            sp.simplify(sum(riemann[a][b][a][d] for a in range(dimension)))
            for d in range(dimension)
        ]
        for b in range(dimension)
    ]
    return christoffel, riemann, ricci


def check_normalized_ehlers_to_taub_nut() -> None:
    seed_mass = sp.symbols("m", positive=True)
    seed_radius = sp.symbols("x", positive=True)
    ehlers = sp.symbols("lambda", real=True)
    seed_lapse = 1 - 2 * seed_mass / seed_radius

    transformed_ernst = (seed_lapse + sp.I * ehlers) / (
        1 + sp.I * ehlers * seed_lapse
    )
    transformed_lapse = (1 + ehlers**2) * seed_lapse / (
        1 + ehlers**2 * seed_lapse**2
    )
    transformed_twist = ehlers * (1 - seed_lapse**2) / (
        1 + ehlers**2 * seed_lapse**2
    )
    require_zero(
        "normalized Ehlers real part",
        sp.re(transformed_ernst).expand(complex=True) - transformed_lapse,
    )
    require_zero(
        "normalized Ehlers imaginary part",
        sp.im(transformed_ernst).expand(complex=True) - transformed_twist,
    )
    require_zero("normalized Ehlers infinity", sp.limit(transformed_ernst, seed_radius, sp.oo) - 1)

    denominator = 1 + ehlers**2
    mass = seed_mass * (1 - ehlers**2) / denominator
    nut = 2 * seed_mass * ehlers / denominator
    radius = seed_radius - 2 * seed_mass * ehlers**2 / denominator
    sigma = sp.expand(radius**2 + nut**2)
    delta = sp.expand(radius**2 - 2 * mass * radius - nut**2)
    require_zero("Ehlers charge circle", mass**2 + nut**2 - seed_mass**2)
    require_zero("Taub--NUT lapse map", transformed_lapse - delta / sigma)
    require_zero(
        "Taub--NUT twist map",
        transformed_twist - 2 * nut * (radius - mass) / sigma,
    )
    jacobian = sp.simplify(
        sp.det(
            sp.Matrix(
                [
                    [sp.diff(mass, seed_mass), sp.diff(mass, ehlers)],
                    [sp.diff(nut, seed_mass), sp.diff(nut, ehlers)],
                ]
            )
        )
    )
    require_zero("Ehlers parameter-map Jacobian", jacobian - 2 * seed_mass / denominator)
    require_zero(
        "Ehlers inverse parameter",
        ehlers - nut / (seed_mass + mass),
    )

    t = sp.symbols("t", real=True)
    r = sp.symbols("r", real=True)
    u = sp.symbols("u", real=True)
    phi = sp.symbols("phi", real=True)
    physical_mass, physical_nut = sp.symbols("M n", real=True)
    coordinates = (t, r, u, phi)
    sigma_r = r**2 + physical_nut**2
    delta_r = r**2 - 2 * physical_mass * r - physical_nut**2
    lapse = delta_r / sigma_r
    metric = sp.Matrix(
        [
            [-lapse, 0, 0, -2 * physical_nut * u * lapse],
            [0, 1 / lapse, 0, 0],
            [0, 0, sigma_r / (1 - u**2), 0],
            [
                -2 * physical_nut * u * lapse,
                0,
                0,
                sigma_r * (1 - u**2) - 4 * physical_nut**2 * u**2 * lapse,
            ],
        ]
    )
    require_zero("Taub--NUT determinant in cosine coordinate", metric.det() + sigma_r**2)
    inverse_metric = sp.simplify(metric.inv())

    _, riemann, ricci = connection_and_curvature(metric, coordinates)
    for a in range(4):
        for b in range(4):
            require_zero(f"Taub--NUT Ricci[{a},{b}]", ricci[a][b])

    # Construct the page's three-dimensional Hodge dual rather than inserting
    # the expected radial derivative by hand.  On 0<theta<pi and Delta>0, the
    # oriented volume density for h=diag(1,Delta,Delta sin^2 theta) is
    # Delta sin theta in the orientation dr wedge dtheta wedge dphi.
    theta = sp.symbols("theta", real=True)
    base_coordinates = (r, theta, phi)
    base_metric = sp.diag(1, delta_r, delta_r * sp.sin(theta) ** 2)
    base_inverse = sp.simplify(base_metric.inv())
    potential = sp.Matrix([0, 0, 2 * physical_nut * sp.cos(theta)])
    field_strength = sp.Matrix(
        3,
        3,
        lambda a, b: sp.diff(potential[b], base_coordinates[a])
        - sp.diff(potential[a], base_coordinates[b]),
    )
    base_volume_density = delta_r * sp.sin(theta)
    hodge_dual = sp.Matrix(
        [
            sp.simplify(
                sum(
                    base_volume_density
                    * sp.LeviCivita(i, j, k)
                    * base_inverse[j, a]
                    * base_inverse[k, b]
                    * field_strength[a, b]
                    / 2
                    for j in range(3)
                    for k in range(3)
                    for a in range(3)
                    for b in range(3)
                )
            )
            for i in range(3)
        ]
    )
    require_zero_matrix(
        "Taub--NUT oriented Hodge dual",
        hodge_dual - sp.Matrix([-2 * physical_nut / delta_r, 0, 0]),
    )

    twist = 2 * physical_nut * (r - physical_mass) / sigma_r
    twist_gradient = sp.Matrix(
        [sp.diff(twist, coordinate) for coordinate in base_coordinates]
    )
    require_zero_matrix(
        "Taub--NUT twist reconstruction",
        twist_gradient - lapse**2 * hodge_dual,
    )

    # At u=0 the metric is diagonal, which makes a direct contraction of the
    # already reconstructed full Riemann tensor compact.  The lifted local
    # SU(2) isometry of Taub--NUT is transitive on each constant-r group orbit,
    # so scalar invariants depend on r alone; the full tensor above is retained
    # at generic u.
    equatorial_metric = sp.simplify(metric.subs(u, 0))
    equatorial_inverse = sp.simplify(equatorial_metric.inv())
    inverse_diagonal = [equatorial_inverse[a, a] for a in range(4)]
    lowered_equatorial = [
        [
            [
                [
                    sp.simplify(
                        sum(
                            equatorial_metric[a, e]
                            * riemann[e][b][c][d].subs(u, 0)
                            for e in range(4)
                        )
                    )
                    for d in range(4)
                ]
                for c in range(4)
            ]
            for b in range(4)
        ]
        for a in range(4)
    ]
    raised_equatorial = [
        [
            [
                [
                    sp.simplify(
                        lowered_equatorial[a][b][c][d]
                        * inverse_diagonal[a]
                        * inverse_diagonal[b]
                        * inverse_diagonal[c]
                        * inverse_diagonal[d]
                    )
                    for d in range(4)
                ]
                for c in range(4)
            ]
            for b in range(4)
        ]
        for a in range(4)
    ]
    kretschmann = sp.factor(
        sum(
            lowered_equatorial[a][b][c][d]
            * raised_equatorial[a][b][c][d]
            for a in range(4)
            for b in range(4)
            for c in range(4)
            for d in range(4)
        )
    )
    pontryagin = sp.factor(
        sum(
            sp.LeviCivita(a, b, e, f)
            * sigma_r
            * inverse_diagonal[e]
            * inverse_diagonal[f]
            * lowered_equatorial[e][f][c][d]
            * raised_equatorial[a][b][c][d]
            / 2
            for a in range(4)
            for b in range(4)
            for e in range(4)
            for f in range(4)
            for c in range(4)
            for d in range(4)
        )
    )
    curvature_magnitude = 48 * (physical_mass**2 + physical_nut**2) / sigma_r**3
    require_zero(
        "Taub--NUT quadratic curvature magnitude",
        kretschmann**2 + pontryagin**2 - curvature_magnitude**2,
    )

    horizon_distance = sp.sqrt(physical_mass**2 + physical_nut**2)
    outer_horizon = physical_mass + horizon_distance
    require_zero("outer Taub--NUT horizon", delta_r.subs(r, outer_horizon))

    # Ingoing v has dv=dt+Sigma/Delta dr.  The transformed metric is finite
    # at Delta=0 and has g_vr=1 rather than a divergent g_rr.
    ingoing_jacobian = sp.Matrix(
        [
            [1, -sigma_r / delta_r, 0, 0],
            [0, 1, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )
    ingoing_metric = sp.simplify(ingoing_jacobian.T * metric * ingoing_jacobian)
    require_zero("horizon-normal norm", inverse_metric[1, 1] - lapse)
    require_zero("ingoing horizon radial square", ingoing_metric[1, 1])
    require_zero("ingoing horizon cross term", ingoing_metric[0, 1] - 1)
    require_zero("ingoing azimuthal radial term", ingoing_metric[1, 3] - 2 * physical_nut * u)
    require_zero("ingoing metric determinant", ingoing_metric.det() + sigma_r**2)
    require_zero(
        "ingoing horizon-generator norm",
        ingoing_metric[0, 0].subs(r, outer_horizon),
    )

    killing_block = metric.extract([0, 3], [0, 3])
    north_generator = sp.Matrix([-2 * physical_nut, 1])
    south_generator = sp.Matrix([2 * physical_nut, 1])
    require_zero_matrix(
        "north-axis kernel",
        killing_block.subs(u, 1) * north_generator,
    )
    require_zero_matrix(
        "south-axis kernel",
        killing_block.subs(u, -1) * south_generator,
    )

    for label, axis_value, generator in (
        ("north", 1, sp.Matrix([-2 * physical_nut, 0, 0, 1])),
        ("south", -1, sp.Matrix([2 * physical_nut, 0, 0, 1])),
    ):
        norm = sp.simplify((generator.T * metric * generator)[0])
        gradient = sp.Matrix([sp.diff(norm, coordinate) for coordinate in coordinates])
        ratio = sp.factor((gradient.T * inverse_metric * gradient)[0] / (4 * norm))
        require_zero(
            f"{label}-axis elementary flatness",
            sp.limit(ratio, u, axis_value, dir="-" if axis_value == 1 else "+") - 1,
        )

    signed_time_shift = 2 * sp.pi * (south_generator[0] - north_generator[0])
    require_zero("two-axis time patching", signed_time_shift - 8 * sp.pi * physical_nut)
    representative_stationary_norm = metric[0, 0].subs(
        {physical_mass: 1, physical_nut: 1, r: 3}
    )
    if representative_stationary_norm >= 0:
        raise AssertionError("periodic time should be timelike in the representative exterior")

    require_zero("asymptotic mass coefficient", sp.limit(r * (1 + metric[0, 0]) / 2, r, sp.oo) - physical_mass)
    require_zero("asymptotic NUT coefficient", sp.limit(r * twist / 2, r, sp.oo) - physical_nut)
    require_zero(
        "nondecaying asymptotic cross term",
        sp.limit(metric[0, 3], r, sp.oo) + 2 * physical_nut * u,
    )

    schwarzschild_cosine = sp.diag(
        -(1 - 2 * seed_mass / r),
        1 / (1 - 2 * seed_mass / r),
        r**2 / (1 - u**2),
        r**2 * (1 - u**2),
    )
    require_zero_matrix(
        "Schwarzschild identity limit",
        metric.subs({physical_mass: seed_mass, physical_nut: 0})
        - schwarzschild_cosine,
    )


def check_spinning_conical_quotient() -> None:
    t, phi, z = sp.symbols("t phi z", real=True)
    rho = sp.symbols("rho", positive=True)
    conicity = sp.symbols("C", positive=True)
    beta = sp.symbols("beta", real=True)
    coordinates = (t, rho, phi, z)

    metric = sp.Matrix(
        [
            [-1, 0, -beta, 0],
            [0, conicity ** -2, 0, 0],
            [-beta, 0, rho**2 - beta**2, 0],
            [0, 0, 0, 1],
        ]
    )

    # Local pullback from cylindrical Minkowski coordinates
    # T=t+beta phi, R=rho/C, Phi=C phi, Z=z.
    jacobian = sp.Matrix(
        [
            [1, 0, beta, 0],
            [0, 1 / conicity, 0, 0],
            [0, 0, conicity, 0],
            [0, 0, 0, 1],
        ]
    )
    minkowski_cylindrical = sp.diag(-1, 1, (rho / conicity) ** 2, 1)
    require_zero_matrix(
        "local Minkowski pullback",
        jacobian.T * minkowski_cylindrical * jacobian - metric,
    )
    require_zero("invertible local map", jacobian.det() - 1)

    killing_block = metric.extract([0, 2], [0, 2])
    require_zero("canonical Killing-block determinant", killing_block.det() + rho**2)
    require_zero("full determinant", metric.det() + rho**2 / conicity**2)

    _, riemann, ricci = connection_and_curvature(metric, coordinates)
    for a in range(4):
        for b in range(4):
            require_zero(f"quotient Ricci[{a},{b}]", ricci[a][b])
            for c in range(4):
                for d in range(4):
                    require_zero(
                        f"punctured-domain Riemann[{a},{b},{c},{d}]",
                        riemann[a][b][c][d],
                    )

    # The degenerating local direction is v=partial_phi-beta partial_t.
    v = sp.Matrix([-beta, 0, 1, 0])
    eta = sp.Matrix([0, 0, 1, 0])
    v_squared = sp.expand((v.T * metric * v)[0])
    eta_squared = sp.expand((eta.T * metric * eta)[0])
    require_zero("degenerating direction norm", v_squared - rho**2)
    require_zero("closed axial-orbit norm", eta_squared - (rho**2 - beta**2))

    gradient = sp.Matrix([sp.diff(v_squared, coordinate) for coordinate in coordinates])
    elementary_flatness = sp.simplify(
        (gradient.T * metric.inv() * gradient)[0] / (4 * v_squared)
    )
    require_zero("elementary-flatness ratio", elementary_flatness - conicity**2)

    if sp.simplify(eta_squared.subs({rho: sp.Rational(1, 5), beta: sp.Rational(2, 5)})) >= 0:
        raise AssertionError("the representative closed phi orbit should be timelike")

    regular_minkowski = sp.diag(-1, 1, rho**2, 1)
    require_zero_matrix(
        "regular-axis identity limit",
        metric.subs({conicity: 1, beta: 0}) - regular_minkowski,
    )

    # Under phi -> phi+2pi, the local flat coordinates shift helically.
    flat_time = t + beta * phi
    flat_angle = conicity * phi
    delta_t_flat = sp.expand(flat_time.subs(phi, phi + 2 * sp.pi) - flat_time)
    delta_phi_flat = sp.expand(flat_angle.subs(phi, phi + 2 * sp.pi) - flat_angle)
    require_zero("helical time shift", delta_t_flat - 2 * sp.pi * beta)
    require_zero("conical angular shift", delta_phi_flat - 2 * sp.pi * conicity)


def check_vsi_plane_wave() -> None:
    u, v, x, y = sp.symbols("u v x y", real=True)
    coordinates = (u, v, x, y)
    profile = x**2 - y**2
    metric = sp.Matrix(
        [
            [profile, -1, 0, 0],
            [-1, 0, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )
    inverse = sp.simplify(metric.inv())
    _, riemann, ricci = connection_and_curvature(metric, coordinates)
    for a in range(4):
        for b in range(4):
            require_zero(f"plane-wave Ricci[{a},{b}]", ricci[a][b])

    lowered = [
        [
            [
                [
                    sp.simplify(sum(metric[a, e] * riemann[e][b][c][d] for e in range(4)))
                    for d in range(4)
                ]
                for c in range(4)
            ]
            for b in range(4)
        ]
        for a in range(4)
    ]
    if all(
        lowered[a][b][c][d] == 0
        for a in range(4)
        for b in range(4)
        for c in range(4)
        for d in range(4)
    ):
        raise AssertionError("the plane-wave curvature should be nonzero")

    raised = [
        [
            [
                [
                    sp.simplify(
                        sum(
                            inverse[a, e]
                            * inverse[b, f]
                            * inverse[c, g]
                            * inverse[d, h]
                            * lowered[e][f][g][h]
                            for e in range(4)
                            for f in range(4)
                            for g in range(4)
                            for h in range(4)
                        )
                    )
                    for d in range(4)
                ]
                for c in range(4)
            ]
            for b in range(4)
        ]
        for a in range(4)
    ]
    kretschmann = sp.simplify(
        sum(
            lowered[a][b][c][d] * raised[a][b][c][d]
            for a in range(4)
            for b in range(4)
            for c in range(4)
            for d in range(4)
        )
    )
    require_zero("nonflat plane-wave Kretschmann scalar", kretschmann)


def main() -> None:
    check_normalized_ehlers_to_taub_nut()
    check_spinning_conical_quotient()
    check_vsi_plane_wave()
    print("Generated-solution audit checks passed:")
    print("- normalized Ehlers map from Schwarzschild to Taub--NUT")
    print("- full Taub--NUT Ricci tensor, twist, curvature magnitude, horizon chart, axes, charges, and limit")
    print("- exact local Minkowski pullback and vanishing punctured-domain curvature")
    print("- canonical determinant, Lorentzian local form, and identity limit")
    print("- nonclosing axis kernel, elementary-flatness ratio, and closed-timelike-orbit screen")
    print("- Ricci-flat nonflat plane wave with zero Kretschmann scalar")


if __name__ == "__main__":
    main()