#!/usr/bin/env python3
"""Exact and pointwise checks for the Volume 5 Bäcklund-to-Kerr benchmark."""

from __future__ import annotations

import sympy as sp


def require_zero(label: str, expression: sp.Expr) -> None:
    reduced = sp.factor(sp.cancel(sp.trigsimp(expression)))
    if reduced != 0:
        raise AssertionError(f"{label}: expected zero, got {reduced}")


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


def check_parameter_sheet() -> None:
    m = sp.symbols("m", positive=True)
    sheet_coordinate = sp.symbols("sheet_coordinate", real=True)
    p = 1 / sp.sqrt(1 + sheet_coordinate**2)
    q = sheet_coordinate / sp.sqrt(1 + sheet_coordinate**2)
    sigma = m * p
    a = m * q

    require_zero("positive-sheet identity p^2+q^2=1", p**2 + q**2 - 1)
    require_zero("subextremal endpoint relation", sigma**2 - (m**2 - a**2))
    require_zero("nonrotating point q=0", q.subs(sheet_coordinate, 0))
    require_zero("nonrotating point p=1", p.subs(sheet_coordinate, 0) - 1)
    if p.is_positive is not True or sigma.is_positive is not True:
        raise AssertionError("the selected subextremal sheet must have p>0 and sigma>0")


def check_pure_gauge_zero_curvature() -> None:
    u, v, spectral = sp.symbols("u v spectral", real=True)
    polynomial_gauge = sp.Matrix(
        [
            [1 + u * v * spectral**2, u * spectral],
            [v * spectral, 1],
        ]
    )
    require_zero("polynomial gauge determinant", polynomial_gauge.det() - 1)

    # For dPhi=A Phi and Phi=P, A_u=P_u P^{-1}, A_v=P_v P^{-1}.
    inverse_gauge = polynomial_gauge.inv()
    connection_u = sp.simplify(sp.diff(polynomial_gauge, u) * inverse_gauge)
    connection_v = sp.simplify(sp.diff(polynomial_gauge, v) * inverse_gauge)
    curvature_uv = sp.simplify(
        sp.diff(connection_u, v)
        - sp.diff(connection_v, u)
        + connection_u * connection_v
        - connection_v * connection_u
    )
    require_zero_matrix("pure-gauge zero curvature", curvature_uv)


def check_spectral_branch() -> None:
    k, z = sp.symbols("K z", real=True)
    rho = sp.symbols("rho", positive=True)
    i = sp.I
    zeta = rho + i * z
    bar_zeta = sp.conjugate(zeta)
    radius = sp.sqrt((k - z) ** 2 + rho**2)
    lambda_squared = (k - i * bar_zeta) / (k + i * zeta)

    # This continuous real-K sheet approaches +1 as K tends to +infinity.
    lambda_branch = (k - z - i * rho) / radius
    require_zero(
        "real-K radial identity",
        (k - i * bar_zeta) * (k + i * zeta) - radius**2,
    )
    require_zero("spectral square-root branch", lambda_branch**2 - lambda_squared)
    require_zero(
        "real-K spectral reality",
        sp.conjugate(lambda_branch) - 1 / lambda_branch,
    )
    require_zero(
        "Lambda tends to one on the selected sheet",
        sp.limit(lambda_branch, k, sp.oo) - 1,
    )


def check_projective_riccati_system() -> None:
    """Derive the Riccati coefficients directly from the displayed Lax pair."""

    alpha = sp.symbols("alpha")
    spectral = sp.symbols("Lambda", nonzero=True)
    l_component, m_component = sp.symbols("l_component m_component", nonzero=True)
    m3, m4, mstar3, mstar4 = sp.symbols("M3 M4 Mstar3 Mstar4")

    lax_zeta = sp.Matrix(
        [
            [m4, spectral * m4],
            [spectral * m3, m3],
        ]
    )
    lax_bar_zeta = sp.Matrix(
        [
            [mstar4, mstar4 / spectral],
            [mstar3 / spectral, mstar3],
        ]
    )
    projective_vector = sp.Matrix([l_component, m_component])

    def quotient_derivative(lax_matrix: sp.MatrixBase) -> sp.Expr:
        vector_derivative = lax_matrix * projective_vector
        derivative = (
            vector_derivative[0] * m_component
            - l_component * vector_derivative[1]
        ) / m_component**2
        return sp.cancel(derivative.subs(l_component, alpha * m_component))

    derived_zeta = quotient_derivative(lax_zeta)
    derived_bar_zeta = quotient_derivative(lax_bar_zeta)
    corrected_zeta = m4 * (alpha + spectral) - alpha * m3 * (
        1 + spectral * alpha
    )
    corrected_bar_zeta = mstar4 * (alpha + 1 / spectral) - alpha * mstar3 * (
        1 + alpha / spectral
    )
    require_zero(
        "projective Riccati zeta coefficient from alpha=l/m",
        derived_zeta - corrected_zeta,
    )
    require_zero(
        "projective Riccati bar-zeta coefficient from alpha=l/m",
        derived_bar_zeta - corrected_bar_zeta,
    )

    # Stephani et al. Eq. (34.93) prints the two coefficients in the
    # factorized form below. In the displayed left-acting Lax normalization
    # and the alpha=l/m projective chart, that form is not generically equal
    # to the direct quotient-rule result.
    stephani_factorized_zeta = (alpha + spectral) * (m4 - alpha * m3)
    stephani_factorized_bar_zeta = (alpha + 1 / spectral) * (
        mstar4 - alpha * mstar3
    )
    zeta_discrepancy = sp.cancel(corrected_zeta - stephani_factorized_zeta)
    bar_zeta_discrepancy = sp.cancel(
        corrected_bar_zeta - stephani_factorized_bar_zeta
    )
    require_zero(
        "Stephani 34.93 zeta discrepancy factor",
        zeta_discrepancy - alpha * m3 * (alpha - 1) * (1 - spectral),
    )
    require_zero(
        "Stephani 34.93 bar-zeta discrepancy factor",
        bar_zeta_discrepancy
        - alpha * mstar3 * (alpha - 1) * (1 - 1 / spectral),
    )

    generic_point = {
        alpha: sp.Integer(2),
        spectral: sp.Integer(3),
        m3: sp.Integer(5),
        m4: sp.Integer(7),
        mstar3: sp.Integer(11),
        mstar4: sp.Integer(13),
    }
    require_zero(
        "generic nonzero zeta discrepancy witness",
        zeta_discrepancy.subs(generic_point) + 20,
    )
    require_zero(
        "generic nonzero bar-zeta discrepancy witness",
        bar_zeta_discrepancy.subs(generic_point) - sp.Rational(44, 3),
    )
    if zeta_discrepancy.subs(generic_point) == 0:
        raise AssertionError("the factorized zeta coefficient must differ generically")
    if bar_zeta_discrepancy.subs(generic_point) == 0:
        raise AssertionError("the factorized bar-zeta coefficient must differ generically")


def check_rogers_shadwick_permutability_maps() -> None:
    """Regress Rogers--Shadwick Eqs. (1.562) and (1.564) exactly."""

    rs_alpha, rs_gamma = sp.symbols("RS_alpha RS_gamma")

    # Rogers--Shadwick Eq. (1.562): the three transformed parameter pairs.
    rs_alpha_prime = (rs_alpha - rs_gamma) / (
        rs_gamma * (rs_alpha - 1)
    )
    rs_gamma_prime = 1 / rs_gamma
    rs_alpha_double_prime = rs_gamma / rs_alpha
    rs_gamma_double_prime = rs_gamma
    rs_alpha_triple_prime = (rs_alpha - 1) / (rs_alpha - rs_gamma)
    rs_gamma_triple_prime = 1 / rs_gamma

    eq_1562_regressions = (
        (
            "RS 1.562 alpha-prime map",
            rs_gamma * (rs_alpha - 1) * rs_alpha_prime
            - (rs_alpha - rs_gamma),
        ),
        ("RS 1.562 gamma-prime map", rs_gamma * rs_gamma_prime - 1),
        (
            "RS 1.562 alpha-double-prime map",
            rs_alpha * rs_alpha_double_prime - rs_gamma,
        ),
        (
            "RS 1.562 gamma-double-prime map",
            rs_gamma_double_prime - rs_gamma,
        ),
        (
            "RS 1.562 alpha-triple-prime map",
            (rs_alpha - rs_gamma) * rs_alpha_triple_prime - (rs_alpha - 1),
        ),
        ("RS 1.562 gamma-triple-prime map", rs_gamma * rs_gamma_triple_prime - 1),
    )
    for label, expression in eq_1562_regressions:
        require_zero(label, expression)

    # Equations (1.559)--(1.562) imply that the combined same-type parameter
    # products on opposite sides of the loop are mutual inverses.
    require_zero(
        "RS 1.562 alpha loop-product closure",
        (rs_alpha * rs_alpha_double_prime)
        * (rs_alpha_prime * rs_alpha_triple_prime)
        - 1,
    )
    require_zero(
        "RS 1.562 gamma loop-product closure",
        (rs_gamma * rs_gamma_double_prime)
        * (rs_gamma_prime * rs_gamma_triple_prime)
        - 1,
    )

    # The rational maps are used only on their stated nonsingular chart.
    # alpha=0, gamma=0, alpha=1, and alpha=gamma are excluded here instead
    # of being passed through cancellation as if the maps extended there.
    rs_loop_denominators = (
        rs_alpha,
        rs_gamma,
        rs_alpha - 1,
        rs_alpha - rs_gamma,
    )
    regular_loop_point = {rs_alpha: sp.Integer(2), rs_gamma: sp.Integer(3)}
    if any(
        sp.simplify(denominator.subs(regular_loop_point)) == 0
        for denominator in rs_loop_denominators
    ):
        raise AssertionError("the regular RS Eq. 1.562 witness crossed an excluded branch")
    excluded_loop_points = (
        ("RS alpha=0", {rs_alpha: 0, rs_gamma: 2}),
        ("RS gamma=0", {rs_alpha: 2, rs_gamma: 0}),
        ("RS alpha=1", {rs_alpha: 1, rs_gamma: 2}),
        ("RS alpha=gamma", {rs_alpha: 2, rs_gamma: 2}),
    )
    for label, point in excluded_loop_points:
        if not any(
            sp.simplify(denominator.subs(point)) == 0
            for denominator in rs_loop_denominators
        ):
            raise AssertionError(f"{label} must be recorded as an excluded branch")

    rs_alpha1, rs_alpha2 = sp.symbols("RS_alpha1 RS_alpha2")
    rs_gamma1, rs_gamma2 = sp.symbols("RS_gamma1 RS_gamma2")
    rs_alpha_ratio = rs_alpha1 / rs_alpha2
    rs_gamma_ratio = rs_gamma1 / rs_gamma2

    # Substitute the ratios of two first-generation data sets into the prime
    # map of Eq. (1.562). This derives Eq. (1.564) rather than re-entering it as
    # an unrelated formula.
    rs_alpha_second_generation_derived = sp.cancel(
        rs_alpha_prime.subs(
            {rs_alpha: rs_alpha_ratio, rs_gamma: rs_gamma_ratio},
            simultaneous=True,
        )
    )
    rs_gamma_second_generation_derived = sp.cancel(
        rs_gamma2
        * rs_gamma_prime.subs(
            {rs_alpha: rs_alpha_ratio, rs_gamma: rs_gamma_ratio},
            simultaneous=True,
        )
    )
    rs_alpha_second_generation_numerator = (
        rs_gamma1 * rs_alpha2 - rs_alpha1 * rs_gamma2
    )
    rs_alpha_second_generation = rs_alpha_second_generation_numerator / (
        rs_gamma1 * (rs_alpha2 - rs_alpha1)
    )
    rs_gamma_second_generation = rs_gamma2**2 / rs_gamma1
    require_zero(
        "RS 1.564 alpha second-generation map from 1.562",
        rs_alpha_second_generation_derived - rs_alpha_second_generation,
    )
    require_zero(
        "RS 1.564 gamma second-generation map from 1.562",
        rs_gamma_second_generation_derived - rs_gamma_second_generation,
    )

    # Both first-generation parameter pairs must be nonzero. The ratio
    # substitution additionally exposes alpha2 and gamma2 as denominators,
    # while Eq. (1.564) requires gamma1 and alpha2-alpha1 to be nonzero. Its
    # alpha numerator must also be nonzero so the second-generation pair stays
    # inside the nonzero transformation chart. Every excluded factor receives
    # a witness below; no rational expression is evaluated on those branches.
    rs_second_generation_domain_factors = (
        rs_alpha1,
        rs_alpha2,
        rs_gamma1,
        rs_gamma2,
        rs_alpha2 - rs_alpha1,
        rs_alpha_second_generation_numerator,
    )
    regular_second_generation_point = {
        rs_alpha1: sp.Integer(2),
        rs_alpha2: sp.Integer(5),
        rs_gamma1: sp.Integer(3),
        rs_gamma2: sp.Integer(7),
    }
    if any(
        sp.simplify(factor.subs(regular_second_generation_point)) == 0
        for factor in rs_second_generation_domain_factors
    ):
        raise AssertionError("the regular RS Eq. 1.564 witness crossed an excluded branch")
    excluded_second_generation_points = (
        (
            "RS alpha1=0 first-generation data",
            0,
            {rs_alpha1: 0, rs_alpha2: 5, rs_gamma1: 3, rs_gamma2: 7},
        ),
        (
            "RS alpha2=0 ratio chart",
            1,
            {rs_alpha1: 2, rs_alpha2: 0, rs_gamma1: 3, rs_gamma2: 7},
        ),
        (
            "RS gamma1=0 second-generation map",
            2,
            {rs_alpha1: 2, rs_alpha2: 5, rs_gamma1: 0, rs_gamma2: 7},
        ),
        (
            "RS gamma2=0 ratio chart",
            3,
            {rs_alpha1: 2, rs_alpha2: 5, rs_gamma1: 3, rs_gamma2: 0},
        ),
        (
            "RS coincident alpha data",
            4,
            {rs_alpha1: 5, rs_alpha2: 5, rs_gamma1: 3, rs_gamma2: 7},
        ),
        (
            "RS alpha-second-generation=0 output",
            5,
            {
                rs_alpha1: 2,
                rs_alpha2: 3,
                rs_gamma1: 5,
                rs_gamma2: sp.Rational(15, 2),
            },
        ),
    )
    for label, excluded_index, point in excluded_second_generation_points:
        evaluated_factors = tuple(
            sp.simplify(factor.subs(point))
            for factor in rs_second_generation_domain_factors
        )
        if evaluated_factors[excluded_index] != 0:
            raise AssertionError(f"{label} must make its domain factor vanish")
        if any(
            value == 0
            for index, value in enumerate(evaluated_factors)
            if index != excluded_index
        ):
            raise AssertionError(f"{label} must isolate exactly one excluded branch")


def check_backlund_output() -> None:
    m, p = sp.symbols("m p", positive=True)
    q, x, y = sp.symbols("q x y", real=True)
    i = sp.I
    sigma = m * p
    r1 = sigma * (x + y)
    r2 = sigma * (x - y)
    phase_minus = p - i * q
    phase_plus = p + i * q
    numerator = phase_minus * r1 + phase_plus * r2 - 2 * m * p
    denominator = phase_minus * r1 + phase_plus * r2 + 2 * m * p
    ernst_bt = sp.cancel(numerator / denominator)
    ernst_prolate = (p * x - i * q * y - 1) / (p * x - i * q * y + 1)
    require_zero("paired Bäcklund output in prolate variables", ernst_bt - ernst_prolate)
    require_zero(
        "q=0 Schwarzschild Ernst branch",
        ernst_prolate.subs({p: 1, q: 0}) - (x - 1) / (x + 1),
    )

    r, a = sp.symbols("r a", real=True)
    ernst_bl = sp.cancel(
        ernst_prolate.subs({x: (r - m) / (m * p), q: a / m})
    )
    ernst_kerr = 1 - 2 * m / (r - i * a * y)
    require_zero("prolate-to-Boyer-Lindquist Ernst map", ernst_bl - ernst_kerr)

    sigma_free = sp.symbols("sigma", real=True, nonzero=True)
    ernst_free = (sigma_free * x - i * a * y - m) / (
        sigma_free * x - i * a * y + m
    )
    real_ernst = (ernst_free + sp.conjugate(ernst_free)) / 2
    laplace_part = sp.diff((x**2 - 1) * sp.diff(ernst_free, x), x) + sp.diff(
        (1 - y**2) * sp.diff(ernst_free, y), y
    )
    gradient_part = (x**2 - 1) * sp.diff(ernst_free, x) ** 2 + (
        1 - y**2
    ) * sp.diff(ernst_free, y) ** 2
    residual = sp.factor(sp.together(real_ernst * laplace_part - gradient_part))
    expected_factor = sp.factor(
        4
        * m
        * (i * a * y + sigma_free * x)
        * (a**2 - m**2 + sigma_free**2)
        / (
            (m + sigma_free * x - i * a * y) ** 2
            * (a**2 * y**2 + (m + sigma_free * x) ** 2)
        )
    )
    require_zero("unspecialized Ernst residual factor", residual - expected_factor)
    require_zero(
        "Kerr endpoint relation solves the Ernst equation",
        residual.subs(sigma_free**2, m**2 - a**2),
    )


def kerr_metric(m: sp.Symbol, a: sp.Symbol, r: sp.Symbol, theta: sp.Symbol) -> sp.Matrix:
    sin2 = sp.sin(theta) ** 2
    sigma2 = r**2 + a**2 * sp.cos(theta) ** 2
    delta = r**2 - 2 * m * r + a**2
    metric = sp.zeros(4)
    metric[0, 0] = -(1 - 2 * m * r / sigma2)
    metric[0, 3] = metric[3, 0] = -2 * m * a * r * sin2 / sigma2
    metric[1, 1] = sigma2 / delta
    metric[2, 2] = sigma2
    metric[3, 3] = (
        sin2 * ((r**2 + a**2) ** 2 - a**2 * delta * sin2) / sigma2
    )
    return metric


def check_reconstruction() -> None:
    m, a, r, theta = sp.symbols("m a r theta", real=True)
    sigma2 = r**2 + a**2 * sp.cos(theta) ** 2
    delta = r**2 - 2 * m * r + a**2
    qbase = (r - m) ** 2 - (m**2 - a**2) * sp.cos(theta) ** 2
    f = (sigma2 - 2 * m * r) / sigma2
    a_phi = 2 * m * a * r * sp.sin(theta) ** 2 / (sigma2 - 2 * m * r)
    rho2 = delta * sp.sin(theta) ** 2
    e2k = (sigma2 - 2 * m * r) / qbase

    metric = kerr_metric(m, a, r, theta)
    require_zero("stationary norm reconstruction", metric[0, 0] + f)
    require_zero("rotational shift reconstruction", metric[0, 3] + f * a_phi)
    require_zero(
        "azimuthal coefficient reconstruction",
        metric[3, 3] - (-f * a_phi**2 + rho2 / f),
    )
    require_zero(
        "canonical Killing-block determinant",
        metric[0, 0] * metric[3, 3] - metric[0, 3] ** 2 + rho2,
    )
    require_zero("Weyl base radial reconstruction", e2k * qbase / f - sigma2)
    require_zero(
        "Weyl base angular reconstruction", e2k * qbase / (f * delta) - sigma2 / delta
    )

    ernst = 1 - 2 * m / (r - sp.I * a * sp.cos(theta))
    f_from_ernst = sp.simplify((ernst + sp.conjugate(ernst)) / 2)
    chi_from_ernst = sp.simplify((ernst - sp.conjugate(ernst)) / (2 * sp.I))
    require_zero("Ernst real part", f_from_ernst - f)
    require_zero("Ernst twist part", chi_from_ernst + 2 * m * a * sp.cos(theta) / sigma2)

    mass_coefficient = sp.limit(r * (metric[0, 0] + 1) / 2, r, sp.oo)
    spin_coefficient = sp.limit(
        -r * metric[0, 3] / (2 * sp.sin(theta) ** 2), r, sp.oo
    )
    require_zero("ADM mass coefficient", mass_coefficient - m)
    require_zero("angular-momentum coefficient", spin_coefficient - m * a)

    schwarzschild = sp.diag(
        -(1 - 2 * m / r),
        1 / (1 - 2 * m / r),
        r**2,
        r**2 * sp.sin(theta) ** 2,
    )
    flat_spherical = sp.diag(
        -1,
        1,
        r**2,
        r**2 * sp.sin(theta) ** 2,
    )
    require_zero_matrix("full Schwarzschild metric limit", metric.subs(a, 0) - schwarzschild)
    require_zero_matrix(
        "full flat metric limit",
        metric.subs({m: 0, a: 0}, simultaneous=True) - flat_spherical,
    )


def check_rods_axes_and_horizon() -> None:
    m, sigma = sp.symbols("m sigma", positive=True)
    a, r, theta = sp.symbols("a r theta", real=True)
    metric = kerr_metric(m, a, r, theta)
    delta = r**2 - 2 * m * r + a**2
    r_plus = m + sigma
    subextremal_relation = {a**2: m**2 - sigma**2}

    require_zero(
        "outer-horizon root",
        delta.subs(r, r_plus).subs(subextremal_relation),
    )
    rho_squared = delta * sp.sin(theta) ** 2
    z_weyl = (r - m) * sp.cos(theta)
    require_zero(
        "finite horizon rod has rho=0",
        rho_squared.subs(r, r_plus).subs(subextremal_relation),
    )
    horizon_z = z_weyl.subs(r, r_plus)
    require_zero("horizon rod north endpoint", horizon_z.subs(theta, 0) - sigma)
    require_zero("horizon rod south endpoint", horizon_z.subs(theta, sp.pi) + sigma)

    omega_horizon = a / (2 * m * r_plus)
    killing_block = metric.extract([0, 3], [0, 3]).subs(r, r_plus)
    horizon_generator = sp.Matrix([1, omega_horizon])
    horizon_kernel = (killing_block * horizon_generator).applyfunc(
        lambda expression: expression.subs(subextremal_relation)
    )
    require_zero_matrix("null horizon generator", horizon_kernel)

    axis_offset = sp.symbols("axis_offset", positive=True)
    exterior_axis_radius = r_plus + axis_offset
    north_axis_z = z_weyl.subs({r: exterior_axis_radius, theta: 0})
    south_axis_z = z_weyl.subs({r: exterior_axis_radius, theta: sp.pi})
    require_zero("north axis interval", north_axis_z - (sigma + axis_offset))
    require_zero("south axis interval", south_axis_z + sigma + axis_offset)
    if (north_axis_z - sigma).is_positive is not True:
        raise AssertionError("the north exterior axis must have z>sigma")
    if (-sigma - south_axis_z).is_positive is not True:
        raise AssertionError("the south exterior axis must have z<-sigma")

    for label, axis_angle in (("north", 0), ("south", sp.pi)):
        require_zero(f"{label} axis g_tphi", metric[0, 3].subs(theta, axis_angle))
        require_zero(f"{label} axis g_phiphi", metric[3, 3].subs(theta, axis_angle))

    elementary_flatness_ratio = sp.cancel(
        metric[3, 3] / (sp.sin(theta) ** 2 * metric[2, 2])
    )
    require_zero(
        "north-axis elementary flatness for period 2pi",
        sp.limit(elementary_flatness_ratio, theta, 0) - 1,
    )
    require_zero(
        "south-axis elementary flatness for period 2pi",
        sp.limit(elementary_flatness_ratio, theta, sp.pi) - 1,
    )


def tensors_at_point(
    metric: sp.Matrix,
    coordinates: tuple[sp.Symbol, ...],
    substitutions: dict[sp.Symbol, sp.Expr],
) -> tuple[list[list[list[sp.Expr]]], list[list[sp.Expr]], sp.Matrix]:
    dimension = len(coordinates)
    metric_at = metric.subs(substitutions)
    inverse_at = sp.simplify(metric_at.inv())
    first = [
        [
            [sp.diff(metric[i, j], coordinates[c]).subs(substitutions) for j in range(dimension)]
            for i in range(dimension)
        ]
        for c in range(dimension)
    ]
    second = [
        [
            [
                [
                    sp.diff(metric[i, j], coordinates[c], coordinates[e]).subs(substitutions)
                    for j in range(dimension)
                ]
                for i in range(dimension)
            ]
            for c in range(dimension)
        ]
        for e in range(dimension)
    ]

    inverse_derivative = [sp.zeros(dimension) for _ in range(dimension)]
    for e in range(dimension):
        dg = sp.Matrix(first[e])
        inverse_derivative[e] = sp.simplify(-inverse_at * dg * inverse_at)

    gamma = [[[sp.S.Zero for _ in range(dimension)] for _ in range(dimension)] for _ in range(dimension)]
    derivative_gamma = [
        [[[sp.S.Zero for _ in range(dimension)] for _ in range(dimension)] for _ in range(dimension)]
        for _ in range(dimension)
    ]
    for upper in range(dimension):
        for lower1 in range(dimension):
            for lower2 in range(dimension):
                connection_terms = [
                    first[lower1][lower2][d]
                    + first[lower2][lower1][d]
                    - first[d][lower1][lower2]
                    for d in range(dimension)
                ]
                gamma[upper][lower1][lower2] = sp.simplify(
                    sum(inverse_at[upper, d] * connection_terms[d] for d in range(dimension)) / 2
                )
                for e in range(dimension):
                    second_terms = [
                        second[e][lower1][lower2][d]
                        + second[e][lower2][lower1][d]
                        - second[e][d][lower1][lower2]
                        for d in range(dimension)
                    ]
                    derivative_gamma[e][upper][lower1][lower2] = sp.simplify(
                        (
                            sum(
                                inverse_derivative[e][upper, d] * connection_terms[d]
                                for d in range(dimension)
                            )
                            + sum(
                                inverse_at[upper, d] * second_terms[d]
                                for d in range(dimension)
                            )
                        )
                        / 2
                    )
    return derivative_gamma, gamma, inverse_at


def check_pp_wave_counterexample() -> None:
    u, v, x, y = sp.symbols("u v x y", real=True)
    coordinates = (u, v, x, y)
    dimension = len(coordinates)
    profile = x**2 - y**2
    metric = sp.Matrix(
        [
            [profile, -1, 0, 0],
            [-1, 0, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )
    derivative_gamma, gamma, inverse_metric = tensors_at_point(
        metric, coordinates, {}
    )

    riemann_up = [
        [
            [
                [sp.S.Zero for _ in range(dimension)] for _ in range(dimension)
            ]
            for _ in range(dimension)
        ]
        for _ in range(dimension)
    ]
    for upper in range(dimension):
        for lower in range(dimension):
            for first in range(dimension):
                for second in range(dimension):
                    value = (
                        derivative_gamma[first][upper][second][lower]
                        - derivative_gamma[second][upper][first][lower]
                    )
                    for contracted in range(dimension):
                        value += (
                            gamma[upper][first][contracted]
                            * gamma[contracted][second][lower]
                            - gamma[upper][second][contracted]
                            * gamma[contracted][first][lower]
                        )
                    riemann_up[upper][lower][first][second] = sp.simplify(value)

    ricci = sp.Matrix(
        dimension,
        dimension,
        lambda lower, second: sp.simplify(
            sum(
                riemann_up[upper][lower][upper][second]
                for upper in range(dimension)
            )
        ),
    )
    require_zero_matrix("pp-wave Ricci tensor", ricci)

    riemann_lower = [
        [
            [
                [
                    sp.simplify(
                        sum(
                            metric[first, upper]
                            * riemann_up[upper][second][third][fourth]
                            for upper in range(dimension)
                        )
                    )
                    for fourth in range(dimension)
                ]
                for third in range(dimension)
            ]
            for second in range(dimension)
        ]
        for first in range(dimension)
    ]
    r_uxux = riemann_lower[0][2][0][2]
    r_uyuy = riemann_lower[0][3][0][3]
    require_zero("pp-wave R_uxux", r_uxux + 1)
    require_zero("pp-wave R_uyuy", r_uyuy - 1)
    require_zero("opposite pp-wave transverse signs", r_uxux + r_uyuy)

    riemann_matrix = sp.Matrix(
        dimension**2,
        dimension**2,
        lambda first_pair, second_pair: riemann_lower[
            first_pair // dimension
        ][first_pair % dimension][second_pair // dimension][second_pair % dimension],
    )
    pair_raising = sp.kronecker_product(inverse_metric, inverse_metric)
    riemann_raised = sp.simplify(
        pair_raising * riemann_matrix * pair_raising.transpose()
    )
    kretschmann = sp.simplify(
        sum(
            riemann_matrix[row, column] * riemann_raised[row, column]
            for row in range(dimension**2)
            for column in range(dimension**2)
        )
    )
    require_zero("pp-wave Kretschmann scalar", kretschmann)


def check_ricci_and_twist() -> None:
    t, r, theta, phi = sp.symbols("t r theta phi", real=True)
    m, a = sp.symbols("m a", real=True)
    coordinates = (t, r, theta, phi)
    metric = kerr_metric(m, a, r, theta)
    test_points = (
        {m: sp.Rational(3, 2), a: sp.Rational(2, 5), r: sp.Rational(17, 4), theta: sp.pi / 3},
        {m: sp.Rational(5, 4), a: -sp.Rational(1, 3), r: sp.Rational(19, 5), theta: sp.pi / 4},
    )

    for point_number, substitutions in enumerate(test_points, start=1):
        derivative_gamma, gamma, inverse_at = tensors_at_point(metric, coordinates, substitutions)
        ricci = sp.zeros(4)
        for alpha in range(4):
            for beta in range(4):
                value = sp.S.Zero
                for c in range(4):
                    value += derivative_gamma[c][c][alpha][beta]
                    value -= derivative_gamma[beta][c][alpha][c]
                    for d in range(4):
                        value += gamma[c][alpha][beta] * gamma[d][c][d]
                        value -= gamma[d][alpha][c] * gamma[c][beta][d]
                ricci[alpha, beta] = sp.simplify(value)
                require_zero(f"Ricci component {alpha}{beta} at point {point_number}", ricci[alpha, beta])

        first_metric = [
            sp.Matrix(
                [
                    [sp.diff(metric[i, j], coordinates[c]).subs(substitutions) for j in range(4)]
                    for i in range(4)
                ]
            )
            for c in range(4)
        ]
        nabla_lower = sp.zeros(4)
        for alpha in range(4):
            for beta in range(4):
                nabla_lower[alpha, beta] = (
                    first_metric[alpha][beta, 0] - first_metric[beta][alpha, 0]
                ) / 2
        nabla_upper = sp.simplify(inverse_at * nabla_lower * inverse_at)
        volume = sp.sqrt(-sp.det(metric.subs(substitutions)))
        omega = []
        for mu in range(4):
            component = sp.S.Zero
            for rho in range(4):
                for sigma in range(4):
                    component += (
                        sp.LeviCivita(mu, 0, rho, sigma)
                        * volume
                        * nabla_upper[rho, sigma]
                    )
            omega.append(sp.simplify(component))
        sigma2 = r**2 + a**2 * sp.cos(theta) ** 2
        chi = -2 * m * a * sp.cos(theta) / sigma2
        for mu, coordinate in enumerate(coordinates):
            require_zero(
                f"twist sign component {mu} at point {point_number}",
                omega[mu] - sp.diff(chi, coordinate).subs(substitutions),
            )


def main() -> None:
    check_parameter_sheet()
    check_pure_gauge_zero_curvature()
    check_spectral_branch()
    check_projective_riccati_system()
    check_rogers_shadwick_permutability_maps()
    check_backlund_output()
    check_reconstruction()
    check_rods_axes_and_horizon()
    check_pp_wave_counterexample()
    check_ricci_and_twist()
    print("Bäcklund-to-Kerr checks passed:")
    print("- positive unit-circle sheet, q=0 endpoint, and spectral branch reality")
    print("- explicit 2x2 polynomial pure-gauge zero curvature")
    print("- corrected projective Riccati system derived from the displayed Lax pair")
    print("- generic discrepancy from the factorized Stephani Eq. 34.93 coefficients")
    print("- exact Rogers-Shadwick 1.562/1.564 maps on their nonsingular charts")
    print("- paired Minkowski output and prolate/Boyer-Lindquist maps")
    print("- exact Ernst residual and endpoint relation")
    print("- full Kerr reconstruction, determinant, asymptotics, and full limits")
    print("- finite horizon rod, null generator, and both elementary-flat axes")
    print("- Ricci-flat pp-wave with nonzero curvature and zero Kretschmann scalar")
    print("- vacuum Ricci tensor and site-oriented twist sign at two exact points")


if __name__ == "__main__":
    main()