"""Test all code examples from TUTORIAL.md to ensure they work correctly."""
import jax
import jax.numpy as jnp

# Enable 64-bit precision
jax.config.update("jax_enable_x64", True)


def test_jets_basic():
    """Test basic Jet creation and usage."""
    from local_coordinates.jet import Jet, function_to_jet, jet_decorator

    # Create a Jet from a function
    def f(x):
        return x[0]**2 + jnp.sin(x[1])

    p = jnp.array([1.0, 0.5])
    jet = function_to_jet(f, p)

    print(f"Value: {jet.value}")
    print(f"Gradient: {jet.gradient}")
    print(f"Hessian:\n{jet.hessian}")
    
    # Verify values
    assert jnp.isclose(jet.value, 1.0 + jnp.sin(0.5))
    assert jnp.allclose(jet.gradient, jnp.array([2.0, jnp.cos(0.5)]))


def test_jet_decorator():
    """Test the @jet_decorator."""
    from local_coordinates.jet import function_to_jet, jet_decorator

    @jet_decorator
    def square(x):
        return x**2

    # Create a simple identity jet
    x_jet = function_to_jet(lambda t: t, jnp.array(3.0))
    result = square(x_jet)

    print(f"Value: {result.value}")
    print(f"Gradient: {result.gradient}")
    print(f"Hessian: {result.hessian}")
    
    assert jnp.isclose(result.value, 9.0)


def test_taylor_evaluation():
    """Test Taylor evaluation of Jets."""
    from local_coordinates.jet import function_to_jet

    def f(x):
        return x[0]**2 + jnp.sin(x[1])

    p = jnp.array([1.0, 0.5])
    jet = function_to_jet(f, p)
    dx = jnp.array([0.1, -0.1])
    approx = jet(dx)
    
    # Compare to actual function value
    actual = f(p + dx)
    print(f"Taylor approx: {approx}, Actual: {actual}")
    assert jnp.isclose(approx, actual, rtol=1e-2)


def test_jacobian_basic():
    """Test basic Jacobian creation."""
    from local_coordinates.jacobian import function_to_jacobian

    # Polar to Cartesian transformation
    def polar_to_cartesian(rtheta):
        r, theta = rtheta[0], rtheta[1]
        return jnp.array([r * jnp.cos(theta), r * jnp.sin(theta)])

    p_polar = jnp.array([2.0, jnp.pi/4])
    J = function_to_jacobian(polar_to_cartesian, p_polar)

    print(f"Jacobian matrix:\n{J.value}")
    
    # Check shape
    assert J.value.shape == (2, 2)


def test_jacobian_inverse():
    """Test Jacobian inversion."""
    from local_coordinates.jacobian import function_to_jacobian

    def polar_to_cartesian(rtheta):
        r, theta = rtheta[0], rtheta[1]
        return jnp.array([r * jnp.cos(theta), r * jnp.sin(theta)])

    p_polar = jnp.array([2.0, jnp.pi/4])
    J = function_to_jacobian(polar_to_cartesian, p_polar)
    
    # Get the inverse Jacobian
    J_inv = J.get_inverse()
    
    # J @ J_inv should be identity
    product = J.value @ J_inv.value
    print(f"J @ J_inv:\n{product}")
    assert jnp.allclose(product, jnp.eye(2), atol=1e-10)


def test_basis_vectors():
    """Test BasisVectors creation."""
    from local_coordinates.basis import get_standard_basis

    p = jnp.array([1.0, 2.0])
    standard_basis = get_standard_basis(p)

    print(f"Basis components:\n{standard_basis.components.value}")
    
    assert jnp.allclose(standard_basis.components.value, jnp.eye(2))


def test_frame_and_lie_brackets():
    """Test Frame and Lie bracket computation."""
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.frame import basis_to_frame, get_lie_bracket_between_frame_pairs

    p = jnp.array([1.0, 2.0])
    basis = get_standard_basis(p)
    
    # Convert BasisVectors to Frame
    frame = basis_to_frame(basis)
    
    # Compute Lie brackets [E_i, E_j] for all pairs
    lie_brackets = get_lie_bracket_between_frame_pairs(frame)
    
    # For standard basis, Lie brackets should vanish
    print(f"Lie brackets (should be zero):\n{lie_brackets.components.value}")
    assert jnp.allclose(lie_brackets.components.value, 0.0, atol=1e-10)


def test_tensor_creation():
    """Test Tensor creation."""
    from local_coordinates.jet import Jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.tensor import Tensor, TensorType

    p = jnp.array([1.0, 2.0])
    basis = get_standard_basis(p)
    
    # Create a (1, 1) mixed tensor
    tensor_type = TensorType(k=1, l=1)
    components_jet = Jet(
        value=jnp.eye(2),
        gradient=jnp.zeros((2, 2, 2)),
        hessian=jnp.zeros((2, 2, 2, 2))
    )
    tensor = Tensor(
        tensor_type=tensor_type,
        basis=basis,
        components=components_jet
    )
    
    print(f"Tensor type: ({tensor.tensor_type.k}, {tensor.tensor_type.l})")
    assert tensor.tensor_type.k == 1
    assert tensor.tensor_type.l == 1


def test_euclidean_metric():
    """Test Euclidean metric creation."""
    from local_coordinates.metric import get_euclidean_metric

    p = jnp.array([1.0, 2.0])
    euclidean_metric = get_euclidean_metric(p)
    
    print(f"Euclidean metric:\n{euclidean_metric.components.value}")
    assert jnp.allclose(euclidean_metric.components.value, jnp.eye(2))


def test_custom_metric():
    """Test custom metric creation."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.0],
            [0.0, 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    print(f"Custom metric:\n{metric.components.value}")
    expected = jnp.array([[1.1, 0.0], [0.0, 1.1]])
    assert jnp.allclose(metric.components.value, expected)


def test_levi_civita_connection():
    """Test Levi-Civita connection computation."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.connection import get_levi_civita_connection

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.0],
            [0.0, 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    connection = get_levi_civita_connection(metric)
    Gamma = connection.christoffel_symbols.value
    
    print(f"Christoffel symbols shape: {Gamma.shape}")
    print(f"Christoffel symbols:\n{Gamma}")
    assert Gamma.shape == (2, 2, 2)


def test_riemann_curvature():
    """Test Riemann curvature tensor computation."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.connection import get_levi_civita_connection
    from local_coordinates.riemann import get_riemann_curvature_tensor, get_ricci_tensor

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.0],
            [0.0, 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    connection = get_levi_civita_connection(metric)
    riemann = get_riemann_curvature_tensor(connection)
    R = riemann.components.value
    
    print(f"Riemann tensor shape: {R.shape}")
    assert R.shape == (2, 2, 2, 2)
    
    ricci = get_ricci_tensor(connection, R=riemann)
    print(f"Ricci tensor:\n{ricci.components.value}")
    assert ricci.components.value.shape == (2, 2)


def test_rnc_jacobians():
    """Test RNC Jacobian computation."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.normal_coords import get_rnc_jacobians

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.0],
            [0.0, 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    J_x_to_v, J_v_to_x = get_rnc_jacobians(metric)
    
    print(f"J_x_to_v value shape: {J_x_to_v.value.shape}")
    print(f"J_v_to_x value shape: {J_v_to_x.value.shape}")
    
    # They should be inverses
    product = J_x_to_v.value @ J_v_to_x.value
    print(f"J_x_to_v @ J_v_to_x:\n{product}")
    assert jnp.allclose(product, jnp.eye(2), atol=1e-10)


def test_to_rnc_metric():
    """Test transforming metric to RNC."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.normal_coords import to_riemann_normal_coordinates

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.05*x[0]*x[1]],
            [0.05*x[0]*x[1], 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)

    print("Original metric at p:")
    print(metric.components.value)

    # Transform to RNC
    metric_rnc = to_riemann_normal_coordinates(metric)

    print("\nMetric in RNC (should be identity at origin):")
    print(metric_rnc.components.value)
    
    # Should be identity
    assert jnp.allclose(metric_rnc.components.value, jnp.eye(2), atol=1e-10)


def test_christoffel_vanish_in_rnc():
    """Test that Christoffel symbols vanish in RNC."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.connection import get_levi_civita_connection
    from local_coordinates.normal_coords import to_riemann_normal_coordinates

    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.05*x[0]*x[1]],
            [0.05*x[0]*x[1], 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    connection = get_levi_civita_connection(metric)
    connection_rnc = to_riemann_normal_coordinates(connection, metric)

    print("\nChristoffel symbols in RNC (should vanish at origin):")
    print(f"Max abs value: {jnp.max(jnp.abs(connection_rnc.christoffel_symbols.value))}")
    
    assert jnp.allclose(connection_rnc.christoffel_symbols.value, 0.0, atol=1e-10)


def test_exponential_map_taylor():
    """Test exponential map with Taylor approximation."""
    from local_coordinates.jet import Jet, function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.tangent import TangentVector
    from local_coordinates.exponential_map import exponential_map_taylor

    # Use Euclidean metric for simplicity
    p = jnp.array([0.0, 0.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(lambda x: jnp.eye(2), p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    
    # Create tangent vector
    v_components = jnp.array([1.0, 0.5])
    v_jet = Jet(
        value=v_components,
        gradient=jnp.zeros((2, 2)),
        hessian=jnp.zeros((2, 2, 2))
    )
    v = TangentVector(p=p, components=v_jet, basis=basis)
    
    # For Euclidean metric, exp_p(v) = p + v
    q = exponential_map_taylor(metric, v)
    
    print(f"exp_p(v) = {q}")
    print(f"Expected (p + v) = {p + v_components}")
    
    # Should be close to p + v for flat metric
    assert jnp.allclose(q, p + v_components, atol=1e-6)


def test_pullback_metric():
    """Test pullback metric computation with same-dimension mapping."""
    from local_coordinates.jet import Jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric, pullback_metric

    # Polar to Cartesian transformation (R^2 -> R^2)
    def polar_to_cartesian(q):
        r, phi = q[0], q[1]
        return jnp.array([r * jnp.cos(phi), r * jnp.sin(phi)])

    r_val, phi_val = 2.0, jnp.pi / 4.0
    p_polar = jnp.array([r_val, phi_val])
    p_cart = polar_to_cartesian(p_polar)

    # Euclidean metric on R^2 (Cartesian coordinates)
    basis_cart = get_standard_basis(p_cart)
    euclidean_2d = RiemannianMetric(
        basis=basis_cart,
        components=Jet(
            value=jnp.eye(2),
            gradient=jnp.zeros((2, 2, 2)),
            hessian=jnp.zeros((2, 2, 2, 2))
        )
    )

    # Pullback to polar coordinates
    polar_metric = pullback_metric(p_polar, polar_to_cartesian, euclidean_2d)
    
    print(f"Polar metric at r={r_val}, phi={phi_val}:")
    print(polar_metric.components.value)
    
    # Expected: diag(1, r^2)
    expected = jnp.diag(jnp.array([1.0, r_val**2]))
    print(f"Expected:\n{expected}")
    assert jnp.allclose(polar_metric.components.value, expected, atol=1e-10)


def test_full_curvature_pipeline():
    """Test the full pipeline: metric -> connection -> curvature."""
    from local_coordinates.jet import function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.connection import get_levi_civita_connection
    from local_coordinates.riemann import get_riemann_curvature_tensor, get_ricci_tensor

    # Position-dependent metric
    def metric_components(x):
        return jnp.array([
            [1 + 0.1*x[0]**2, 0.0],
            [0.0, 1 + 0.1*x[1]**2]
        ])

    p = jnp.array([1.0, 1.0])
    basis = get_standard_basis(p)
    metric_jet = function_to_jet(metric_components, p)
    metric = RiemannianMetric(basis=basis, components=metric_jet)

    print("Metric components:")
    print(metric.components.value)

    # Compute the Levi-Civita connection
    connection = get_levi_civita_connection(metric)
    print("\nChristoffel symbols:")
    print(connection.christoffel_symbols.value)

    # Compute the Riemann curvature tensor
    riemann = get_riemann_curvature_tensor(connection)
    R = riemann.components.value
    print(f"\nRiemann tensor shape: {R.shape}")
    assert R.shape == (2, 2, 2, 2)

    # Compute the Ricci tensor
    ricci = get_ricci_tensor(connection, R=riemann)
    print("\nRicci tensor:")
    print(ricci.components.value)
    assert ricci.components.value.shape == (2, 2)

    # Compute the Ricci scalar
    g_inv = jnp.linalg.inv(metric.components.value)
    ricci_scalar = jnp.einsum("ij,ij->", g_inv, ricci.components.value)
    print(f"\nRicci scalar: {ricci_scalar}")


def test_geodesic_trajectory():
    """Test geodesic computation via exponential map ODE."""
    from local_coordinates.jet import Jet, function_to_jet
    from local_coordinates.basis import get_standard_basis
    from local_coordinates.metric import RiemannianMetric
    from local_coordinates.tangent import TangentVector
    from local_coordinates.exponential_map import exponential_map_ode

    # Define a metric that varies with position
    def make_metric(x):
        basis = get_standard_basis(x)
        g_jet = function_to_jet(
            lambda y: (1.0 / (1.0 + 0.1 * jnp.sum(y**2))) * jnp.eye(2),
            x
        )
        return RiemannianMetric(basis=basis, components=g_jet)

    # Starting point and initial velocity
    p = jnp.array([0.0, 0.0])
    v_components = jnp.array([1.0, 0.5])

    # Create tangent vector
    basis = get_standard_basis(p)
    v_jet = Jet(
        value=v_components,
        gradient=jnp.zeros((2, 2)),
        hessian=jnp.zeros((2, 2, 2))
    )
    v = TangentVector(p=p, components=v_jet, basis=basis)

    # Compute geodesic trajectory
    ts, trajectory = exponential_map_ode(
        v,
        metric_fn=make_metric,
        num_steps=50,
        return_trajectory=True
    )

    print("Geodesic trajectory (first and last points):")
    print(f"  t={ts[0]:.2f}: ({trajectory[0, 0]:.4f}, {trajectory[0, 1]:.4f})")
    print(f"  t={ts[-1]:.2f}: ({trajectory[-1, 0]:.4f}, {trajectory[-1, 1]:.4f})")
    
    # Check that trajectory starts at origin
    assert jnp.allclose(trajectory[0], p, atol=1e-6)
    # Check that we moved somewhere
    assert jnp.linalg.norm(trajectory[-1] - p) > 0.5


if __name__ == "__main__":
    print("=" * 60)
    print("Testing Jets basic")
    print("=" * 60)
    test_jets_basic()
    
    print("\n" + "=" * 60)
    print("Testing jet_decorator")
    print("=" * 60)
    test_jet_decorator()
    
    print("\n" + "=" * 60)
    print("Testing Taylor evaluation")
    print("=" * 60)
    test_taylor_evaluation()
    
    print("\n" + "=" * 60)
    print("Testing Jacobian basic")
    print("=" * 60)
    test_jacobian_basic()
    
    print("\n" + "=" * 60)
    print("Testing Jacobian inverse")
    print("=" * 60)
    test_jacobian_inverse()
    
    print("\n" + "=" * 60)
    print("Testing BasisVectors")
    print("=" * 60)
    test_basis_vectors()
    
    print("\n" + "=" * 60)
    print("Testing Frame and Lie brackets")
    print("=" * 60)
    test_frame_and_lie_brackets()
    
    print("\n" + "=" * 60)
    print("Testing Tensor creation")
    print("=" * 60)
    test_tensor_creation()
    
    print("\n" + "=" * 60)
    print("Testing Euclidean metric")
    print("=" * 60)
    test_euclidean_metric()
    
    print("\n" + "=" * 60)
    print("Testing custom metric")
    print("=" * 60)
    test_custom_metric()
    
    print("\n" + "=" * 60)
    print("Testing Levi-Civita connection")
    print("=" * 60)
    test_levi_civita_connection()
    
    print("\n" + "=" * 60)
    print("Testing Riemann curvature")
    print("=" * 60)
    test_riemann_curvature()
    
    print("\n" + "=" * 60)
    print("Testing RNC Jacobians")
    print("=" * 60)
    test_rnc_jacobians()
    
    print("\n" + "=" * 60)
    print("Testing metric in RNC")
    print("=" * 60)
    test_to_rnc_metric()
    
    print("\n" + "=" * 60)
    print("Testing Christoffel symbols vanish in RNC")
    print("=" * 60)
    test_christoffel_vanish_in_rnc()
    
    print("\n" + "=" * 60)
    print("Testing exponential map (Taylor)")
    print("=" * 60)
    test_exponential_map_taylor()
    
    print("\n" + "=" * 60)
    print("Testing pullback metric")
    print("=" * 60)
    test_pullback_metric()
    
    print("\n" + "=" * 60)
    print("Testing full curvature pipeline")
    print("=" * 60)
    test_full_curvature_pipeline()
    
    print("\n" + "=" * 60)
    print("Testing geodesic trajectory")
    print("=" * 60)
    test_geodesic_trajectory()
    
    print("\n" + "=" * 60)
    print("ALL TESTS PASSED!")
    print("=" * 60)
