import jax
import jax.numpy as jnp
import pytest
from local_coordinates.metric import RiemannianMetric, pullback_metric
from local_coordinates.basis import BasisVectors, get_standard_basis
from local_coordinates.jet import Jet

def test_pullback_metric_basis_handling():
  """
  Test that pullback_metric correctly handles non-standard bases on N.

  pullback_metric internally converts the N-metric to the standard basis,
  computes the pullback G^T h_std G, and returns the result in the standard
  basis of M.
  """
  dim_m = 2
  dim_n = 3

  def f(x):
    return jnp.array([x[0]**2, x[1]**2, x[0]*x[1]])

  p_m = jnp.array([1.0, 2.0])
  p_n = f(p_m)

  E_N_val = jnp.array([
    [1.0, 0.5, 0.0],
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0]
  ])
  basis_n = BasisVectors(
    p=p_n,
    components=Jet(value=E_N_val, gradient=jnp.zeros((3, 3, 3)), hessian=jnp.zeros((3, 3, 3, 3)))
  )

  h_components = Jet(value=jnp.eye(dim_n), gradient=jnp.zeros((dim_n, dim_n, dim_n)), hessian=jnp.zeros((dim_n, dim_n, dim_n, dim_n)))
  metric_n = RiemannianMetric(basis=basis_n, components=h_components)

  metric_m = pullback_metric(p_m, f, metric_n)

  G = jax.jacobian(f)(p_m)

  inv_E_N = jnp.linalg.inv(E_N_val)
  h_std = inv_E_N.T @ jnp.eye(dim_n) @ inv_E_N

  expected_g_std = G.T @ h_std @ G

  assert jnp.allclose(metric_m.components.value, expected_g_std, atol=1e-6)
  assert jnp.allclose(metric_m.basis.components.value, jnp.eye(dim_m))

def test_pullback_metric_standard_basis_parity():
  """
  Test that pullback_metric with standard bases matches the manual formula.
  """
  dim_m = 2
  dim_n = 2

  def f(x):
    return jnp.array([jnp.sin(x[0]), jnp.cos(x[1])])

  p_m = jnp.array([0.5, 0.5])
  p_n = f(p_m)

  basis_n = get_standard_basis(p_n)

  h_components = Jet(value=jnp.eye(dim_n), gradient=jnp.zeros((dim_n, dim_n, dim_n)), hessian=jnp.zeros((dim_n, dim_n, dim_n, dim_n)))
  metric_n = RiemannianMetric(basis=basis_n, components=h_components)

  metric_m = pullback_metric(p_m, f, metric_n)

  G = jax.jacobian(f)(p_m)
  expected_g = G.T @ jnp.eye(dim_n) @ G

  assert jnp.allclose(metric_m.components.value, expected_g)

if __name__ == "__main__":
  pytest.main([__file__])
