import jax
import jax.numpy as jnp
import pytest
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

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

def test_monge_math_verification():
  # Setup
  n = 3
  key = jax.random.PRNGKey(0)
  x = jax.random.normal(key, (n,))
  
  # Sample energy function f(x) = 0.5 * x^T A x
  A = jax.random.normal(jax.random.PRNGKey(1), (n, n))
  A = A @ A.T  # Make it positive definite
  def f(x):
    return 0.5 * jnp.dot(x, jnp.dot(A, x))

  # Monge patch phi(x) = (x, f(x))
  def phi(x):
    return jnp.concatenate([x, jnp.atleast_1d(f(x))])

  # Jacobian J = d phi / dx
  J_num = jax.jacobian(phi)(x)
  grad_f = jax.grad(f)(x)
  J_analytical = jnp.vstack([jnp.eye(n), grad_f[None, :]])
  assert jnp.allclose(J_num, J_analytical)

  # 1. J^T J
  jtj_num = J_analytical.T @ J_analytical
  jtj_analytical = jnp.eye(n) + jnp.outer(grad_f, grad_f)
  assert jnp.allclose(jtj_num, jtj_analytical)

  # 2. (J^T J)^{-1}
  inv_jtj_num = jnp.linalg.inv(jtj_num)
  grad_f_norm_sq = jnp.sum(grad_f**2)
  inv_jtj_analytical = jnp.eye(n) - (1.0 / (1.0 + grad_f_norm_sq)) * jnp.outer(grad_f, grad_f)
  assert jnp.allclose(inv_jtj_num, inv_jtj_analytical)

  # 3. J^+ = (J^T J)^{-1} J^T
  pinv_num = jnp.linalg.pinv(J_analytical)
  term1 = jnp.eye(n) - (1.0 / (1.0 + grad_f_norm_sq)) * jnp.outer(grad_f, grad_f)
  term2 = (1.0 / (1.0 + grad_f_norm_sq)) * grad_f[:, None]
  pinv_analytical = jnp.block([term1, term2])
  assert jnp.allclose(pinv_num, pinv_analytical, atol=1e-6)

  # 4. J_parallel = J J^+
  j_para_num = J_analytical @ pinv_analytical
  j_para_analytical = jnp.block([
    [term1, term2],
    [term2.T, jnp.array([[grad_f_norm_sq / (1.0 + grad_f_norm_sq)]])]
  ])
  assert jnp.allclose(j_para_num, j_para_analytical, atol=1e-6)

  # 5. J_perp = I - J_parallel
  j_perp_num = jnp.eye(n + 1) - j_para_num
  j_perp_analytical = (1.0 / (1.0 + grad_f_norm_sq)) * jnp.block([
    [jnp.outer(grad_f, grad_f), -grad_f[:, None]],
    [-grad_f[None, :], jnp.array([[1.0]])]
  ])
  assert jnp.allclose(j_perp_num, j_perp_analytical)

  # Invariants
  I_n = jnp.eye(n)
  I_np1 = jnp.eye(n + 1)
  assert jnp.allclose(pinv_analytical @ J_analytical, I_n, atol=1e-6)
  assert jnp.allclose(j_para_analytical @ J_analytical, J_analytical, atol=1e-6)
  assert jnp.allclose(j_perp_analytical @ J_analytical, jnp.zeros((n + 1, n)), atol=1e-6)
  assert jnp.allclose(j_para_analytical + j_perp_analytical, I_np1, atol=1e-6)
  assert jnp.allclose(j_para_analytical @ j_para_analytical, j_para_analytical, atol=1e-6)
  assert jnp.allclose(j_perp_analytical @ j_perp_analytical, j_perp_analytical, atol=1e-6)
  assert jnp.allclose(j_para_analytical @ j_perp_analytical, jnp.zeros((n + 1, n + 1)), atol=1e-6)

  # --- Intrinsic Geometry Verification ---

  # 6. Induced Metric g(x)
  # Analytical: g(x) = I + grad_f grad_f^T
  g_analytical = jnp.eye(n) + jnp.outer(grad_f, grad_f)
  
  # Numerical using local_coordinates
  def metric_func(x):
    grad_f_x = jax.grad(f)(x)
    return jnp.eye(n) + jnp.outer(grad_f_x, grad_f_x)
  
  basis = get_standard_basis(x)
  metric_jet = function_to_jet(metric_func, x)
  metric = RiemannianMetric(basis=basis, components=metric_jet)
  
  assert jnp.allclose(metric.components.value, g_analytical)

  # 7. Christoffel Symbols Gamma_{ij}^k
  # Analytical: Gamma_{ij}^k = (1 / (1 + |grad_f|^2)) * (partial^2_{ij} f) * (partial_k f)
  hess_f = jax.hessian(f)(x)
  gamma_analytical = (1.0 / (1.0 + grad_f_norm_sq)) * jnp.einsum('ij,k->ijk', hess_f, grad_f)
  
  # Numerical using local_coordinates
  connection = get_levi_civita_connection(metric)
  gamma_num = connection.christoffel_symbols.value # Gamma[i, j, k]
  
  if not jnp.allclose(gamma_num, gamma_analytical, atol=1e-6):
    print("Christoffel symbols mismatch!")
    print("Numerical Gamma (first slice):")
    print(gamma_num[0])
    print("Analytical Gamma (first slice):")
    print(gamma_analytical[0])
    print("Difference (first slice):")
    print(gamma_num[0] - gamma_analytical[0])
  
  assert jnp.allclose(gamma_num, gamma_analytical, atol=1e-6)

  # --- Extrinsic Geometry Verification (Normal Vector) ---

  # 8. Unit Normal Vector N
  # Analytical: N = (1 / sqrt(1 + |grad_f|^2)) * [-grad_f, 1]
  n_analytical = (1.0 / jnp.sqrt(1.0 + grad_f_norm_sq)) * jnp.concatenate([-grad_f, jnp.array([1.0])])
  
  # Check unit length: |N|^2 = 1
  assert jnp.allclose(jnp.sum(n_analytical**2), 1.0)
  
  # Check orthogonality to tangent space: J^T N = 0
  assert jnp.allclose(J_analytical.T @ n_analytical, jnp.zeros(n), atol=1e-6)
  
  # Check N N^T = J_perp
  assert jnp.allclose(jnp.outer(n_analytical, n_analytical), j_perp_analytical, atol=1e-6)


def test_monge_ambient_covariant_derivative_expansion():
  """
  Numerically verify the derivation in notes/monge.md (lines 225-232):
  (bar nabla_V W)^j = V^i ( partial^2 y^j / partial x^i partial x^a W^a
                         + partial y^j / partial x^a partial W^a / partial x^i ).
  LHS: directional derivative of tilde_W in direction V (via JAX).
  RHS: the derived formula using J, H, and dW/dx.
  """
  n = 3
  key = jax.random.PRNGKey(0)
  x = jax.random.normal(key, (n,))

  A = jax.random.normal(jax.random.PRNGKey(1), (n, n))
  A = A @ A.T

  def f(x):
    return 0.5 * jnp.dot(x, jnp.dot(A, x))

  def phi(x):
    return jnp.concatenate([x, jnp.atleast_1d(f(x))])

  def W(x):
    return jnp.sin(x)

  V = jax.random.normal(jax.random.PRNGKey(2), (n,))

  def tilde_W(x):
    J_at_x = jax.jacrev(phi)(x)
    return J_at_x @ W(x)

  grad_tilde_W = jax.jacrev(tilde_W)(x)
  lhs = grad_tilde_W @ V

  J = jax.jacrev(phi)(x)
  H = jax.jacfwd(jax.jacrev(phi))(x)
  W_val = W(x)
  dW_dx = jax.jacrev(W)(x)

  rhs = jnp.array([
    jnp.dot(V, H[j] @ W_val + J[j] @ dW_dx)
    for j in range(n + 1)
  ])

  assert jnp.allclose(lhs, rhs, atol=1e-5, rtol=1e-5)


def _monge_setup(n=3):
  """Shared setup for Monge patch tests."""
  key = jax.random.PRNGKey(0)
  x = jax.random.normal(key, (n,))

  A = jax.random.normal(jax.random.PRNGKey(1), (n, n))
  A = A @ A.T

  def f(x):
    return 0.5 * jnp.dot(x, jnp.dot(A, x))

  def phi(x):
    return jnp.concatenate([x, jnp.atleast_1d(f(x))])

  grad_f = jax.grad(f)(x)
  hess_f = jax.hessian(f)(x)
  grad_f_norm_sq = jnp.sum(grad_f**2)
  J = jax.jacrev(phi)(x)
  N = (1.0 / jnp.sqrt(1.0 + grad_f_norm_sq)) * jnp.concatenate([-grad_f, jnp.array([1.0])])

  return x, f, phi, grad_f, hess_f, grad_f_norm_sq, J, N, n


def test_decomposition_of_ambient_basis_vector():
  """
  The coordinate vector d/dy^{n+1} is NOT the surface normal. It has both
  tangential and normal components relative to the surface. Verify that:

    d/dy^{n+1} = sum_k (df/dx^k)/(1+|grad f|^2) * phi_*(d/dx^k)
                 + (1/sqrt(1+|grad f|^2)) * N

  This is the key fact that the original derivation in monge.md missed.
  """
  x, f, phi, grad_f, hess_f, grad_f_norm_sq, J, N, n = _monge_setup()

  # The ambient basis vector e_{n+1}
  e_np1 = jnp.zeros(n + 1).at[n].set(1.0)

  # Tangential projection via J_parallel
  j_para = J @ jnp.linalg.pinv(J)
  j_perp = jnp.eye(n + 1) - j_para

  tangent_part_num = j_para @ e_np1
  normal_part_num = j_perp @ e_np1

  # Analytical tangential part: sum_k (df/dx^k)/(1+|grad f|^2) * phi_*(d/dx^k)
  # phi_*(d/dx^k) is the k-th column of J, so this is J @ (grad_f / (1+|grad_f|^2))
  tangent_part_analytical = J @ (grad_f / (1.0 + grad_f_norm_sq))

  # Analytical normal part: (1/sqrt(1+|grad f|^2)) * N
  normal_part_analytical = (1.0 / jnp.sqrt(1.0 + grad_f_norm_sq)) * N

  assert jnp.allclose(tangent_part_num, tangent_part_analytical, atol=1e-10)
  assert jnp.allclose(normal_part_num, normal_part_analytical, atol=1e-10)

  # Verify the decomposition reconstructs e_{n+1}
  assert jnp.allclose(tangent_part_analytical + normal_part_analytical, e_np1, atol=1e-10)


def test_gauss_formula():
  """
  Verify the Gauss formula: bar{nabla}_V W = phi_*(nabla_V W) + II(V, W).

  The tangential projection of bar{nabla}_V W should equal
  phi_*(nabla_V W), where nabla_V W uses the Levi-Civita connection
  of the induced metric (with Christoffel symbols).

  The normal projection should equal (H_f(V,W) / sqrt(1+|grad f|^2)) * N.
  """
  x, f, phi, grad_f, hess_f, grad_f_norm_sq, J, N, n = _monge_setup()

  # Vector fields V and W on the base manifold
  V_val = jax.random.normal(jax.random.PRNGKey(2), (n,))
  def W_field(x):
    return jnp.sin(x)

  W_val = W_field(x)
  dW_dx = jax.jacrev(W_field)(x)

  # --- LHS: bar{nabla}_V W computed directly in the ambient space ---
  def tilde_W(x):
    return jax.jacrev(phi)(x) @ W_field(x)

  bar_nabla_V_W = jax.jacrev(tilde_W)(x) @ V_val

  # Projection matrices
  j_para = J @ jnp.linalg.pinv(J)
  j_perp = jnp.eye(n + 1) - j_para

  tangent_proj = j_para @ bar_nabla_V_W
  normal_proj = j_perp @ bar_nabla_V_W

  # --- RHS tangential part: phi_*(nabla_V W) using Christoffel symbols ---
  # nabla_V W has components: (nabla_V W)^k = V^i dW^k/dx^i + Gamma^k_{ij} V^i W^j
  gamma = (1.0 / (1.0 + grad_f_norm_sq)) * jnp.einsum('ij,k->ijk', hess_f, grad_f)
  nabla_V_W_components = dW_dx @ V_val + jnp.einsum('ijk,i,j->k', gamma, V_val, W_val)

  # Push forward to ambient space: phi_*(nabla_V W) = J @ (nabla_V W components)
  tangent_proj_analytical = J @ nabla_V_W_components

  # --- RHS normal part: II(V, W) = (H_f(V,W) / sqrt(1+|grad f|^2)) * N ---
  H_f_VW = V_val @ hess_f @ W_val
  normal_proj_analytical = (H_f_VW / jnp.sqrt(1.0 + grad_f_norm_sq)) * N

  assert jnp.allclose(tangent_proj, tangent_proj_analytical, atol=1e-6), (
    f"Tangential part mismatch.\n"
    f"  max diff: {jnp.max(jnp.abs(tangent_proj - tangent_proj_analytical))}"
  )
  assert jnp.allclose(normal_proj, normal_proj_analytical, atol=1e-6), (
    f"Normal part mismatch.\n"
    f"  max diff: {jnp.max(jnp.abs(normal_proj - normal_proj_analytical))}"
  )

  # Also verify the full Gauss formula reconstructs bar_nabla_V_W
  full_rhs = tangent_proj_analytical + normal_proj_analytical
  assert jnp.allclose(bar_nabla_V_W, full_rhs, atol=1e-6)


def test_incorrect_identification_fails():
  """
  The original derivation in monge.md incorrectly identifies:
    - Hessian term H_f(V,W) * d/dy^{n+1} as the normal part
    - V^i dW^a/dx^i * phi_*(d/dx^a) as the tangential part

  This test shows that identification is WRONG by checking that the
  tangential projection of bar{nabla}_V W does NOT equal the naive
  tangential term V^i dW^a/dx^i * phi_*(d/dx^a) (it is missing the
  Christoffel symbol contribution).
  """
  x, f, phi, grad_f, hess_f, grad_f_norm_sq, J, N, n = _monge_setup()

  V_val = jax.random.normal(jax.random.PRNGKey(2), (n,))
  def W_field(x):
    return jnp.sin(x)

  W_val = W_field(x)
  dW_dx = jax.jacrev(W_field)(x)

  # Compute bar{nabla}_V W
  def tilde_W(x):
    return jax.jacrev(phi)(x) @ W_field(x)
  bar_nabla_V_W = jax.jacrev(tilde_W)(x) @ V_val

  # Project to get true tangential part
  j_para = J @ jnp.linalg.pinv(J)
  tangent_proj_true = j_para @ bar_nabla_V_W

  # The INCORRECT tangential identification from the old derivation:
  # just V^i dW^k/dx^i pushed forward, with NO Christoffel symbols
  naive_tangent_components = dW_dx @ V_val  # missing Gamma^k_{ij} V^i W^j
  naive_tangent = J @ naive_tangent_components

  # These should NOT be equal (the naive version is missing the Christoffel term)
  assert not jnp.allclose(tangent_proj_true, naive_tangent, atol=1e-6), (
    "The naive tangential identification should NOT match the true tangential "
    "projection. If they match, the Christoffel symbols are zero (degenerate case)."
  )

  # The difference should be exactly the missing Christoffel term
  gamma = (1.0 / (1.0 + grad_f_norm_sq)) * jnp.einsum('ij,k->ijk', hess_f, grad_f)
  christoffel_term = J @ jnp.einsum('ijk,i,j->k', gamma, V_val, W_val)
  difference = tangent_proj_true - naive_tangent
  assert jnp.allclose(difference, christoffel_term, atol=1e-6), (
    f"Difference should be the Christoffel term.\n"
    f"  max diff: {jnp.max(jnp.abs(difference - christoffel_term))}"
  )

  # Similarly, the Hessian term is NOT purely normal
  H_f_VW = V_val @ hess_f @ W_val
  e_np1 = jnp.zeros(n + 1).at[n].set(1.0)
  hessian_term = H_f_VW * e_np1  # the old "normal" identification

  j_perp = jnp.eye(n + 1) - j_para
  normal_proj_true = j_perp @ bar_nabla_V_W

  assert not jnp.allclose(normal_proj_true, hessian_term, atol=1e-6), (
    "The Hessian term H_f(V,W) * e_{n+1} should NOT be purely normal."
  )


if __name__ == "__main__":
  test_monge_math_verification()
  test_monge_ambient_covariant_derivative_expansion()
  test_decomposition_of_ambient_basis_vector()
  test_gauss_formula()
  test_incorrect_identification_fails()
