"""
Test the change of basis formula for Christoffel symbols from orthogonal_coordinates.md

Formula 1 (line 191):
  Γ̃^k_{ij} = T^l_j T^a_i T̂^k_m Γ^m_{al} + Ẽ_i(T^m_j) T̂^k_m

Formula 2 (line 197):
  E_i(T^k_j) = Γ̃^m_{aj} T^k_m T̂^a_i - T^m_j Γ^k_{im}
"""
import jax
import jax.numpy as jnp
from jax import random
from local_coordinates.basis import BasisVectors, get_basis_transform, get_standard_basis
from local_coordinates.jet import Jet, jet_decorator
from local_coordinates.metric import RiemannianMetric
from local_coordinates.connection import Connection, get_levi_civita_connection, change_basis


def test_christoffel_change_of_basis_formula1():
  """
  Test Formula 1: Γ̃^k_{ij} = T^l_j T^a_i T̂^k_m Γ^m_{al} + Ẽ_i(T^m_j) T̂^k_m

  where:
  - (E_1, ..., E_n) is the old frame with Christoffel symbols Γ^m_{al}
  - (Ẽ_1, ..., Ẽ_n) is the new frame with Christoffel symbols Γ̃^k_{ij}
  - Ẽ_i = T^j_i E_j
  - T̂ = T^{-1}
  """
  key = random.PRNGKey(0)
  dim = 3

  # Create a point
  p = jnp.array([0.5, 0.3, 0.2])

  # Create a random positive definite metric (as a function of position)
  key, subkey = random.split(key)
  A = random.normal(subkey, (dim, dim))
  metric_matrix = A @ A.T + 0.1 * jnp.eye(dim)

  # Metric gradient (random but consistent)
  key, subkey = random.split(key)
  metric_grad = random.normal(subkey, (dim, dim, dim)) * 0.1
  # Make it symmetric in first two indices
  metric_grad = 0.5 * (metric_grad + jnp.swapaxes(metric_grad, 0, 1))

  # Create standard basis at point p
  std_basis = get_standard_basis(p)

  # Create metric
  metric_jet = Jet(
    value=metric_matrix,
    gradient=metric_grad,
    hessian=jnp.zeros((dim, dim, dim, dim))
  )
  metric = RiemannianMetric(basis=std_basis, components=metric_jet)

  # Get Levi-Civita connection in standard basis
  connection_old = get_levi_civita_connection(metric)

  # Create a new basis via a random invertible transformation T
  key, subkey = random.split(key)
  T_matrix = random.normal(subkey, (dim, dim)) + 0.5 * jnp.eye(dim)

  # Gradient of T (how T changes with position)
  key, subkey = random.split(key)
  T_grad = random.normal(subkey, (dim, dim, dim)) * 0.1

  # New basis: Ẽ_i = T^j_i E_j
  # In standard basis coords: new_E[a, i] = T[j, i] * old_E[a, j]
  # Since old_E = I, we have new_E[a, i] = T[a, i] (columns are basis vectors)
  new_basis_value = T_matrix
  new_basis_gradient = T_grad
  new_basis = BasisVectors(
    p=p,
    components=Jet(
      value=new_basis_value,
      gradient=new_basis_gradient,
      hessian=jnp.zeros((dim, dim, dim, dim))
    )
  )

  # Use the library's change_basis to get the new Christoffel symbols
  connection_new = change_basis(connection_old, new_basis)
  Gamma_tilde_lib = connection_new.christoffel_symbols.value

  # Now verify Formula 1 manually:
  # Γ̃^k_{ij} = T^l_j T^a_i T̂^k_m Γ^m_{al} + Ẽ_i(T^m_j) T̂^k_m
  Gamma_old = connection_old.christoffel_symbols.value  # Γ^m_{al} indexed as [a, l, m]
  E_old = connection_old.basis.components.value  # E[a, j] = E^a_j (identity here)
  T_hat = jnp.linalg.inv(T_matrix)  # T̂ = T^{-1}

  # Term 1: T^l_j T^a_i T̂^k_m Γ^m_{al}
  # Using einsum: T[l,j] T[a,i] T_hat[k,m] Gamma[a,l,m]
  term1 = jnp.einsum("lj,ai,km,alm->ijk", T_matrix, T_matrix, T_hat, Gamma_old)

  # Term 2: Ẽ_i(T^m_j) T̂^k_m
  # Ẽ_i(f) = Ẽ_i^α ∂_α f = T^a_i E_a^α ∂_α f = T^a_i ∂_a f (since E = I)
  # So Ẽ_i(T^m_j) = T^a_i ∂_a(T^m_j) = T[a,i] T_grad[m,j,a]
  # Then term2 = T[a,i] T_grad[m,j,a] T_hat[k,m]
  term2 = jnp.einsum("ai,mja,km->ijk", T_matrix, T_grad, T_hat)

  Gamma_tilde_manual = term1 + term2

  print("Formula 1 Test:")
  print(f"  Max difference: {jnp.max(jnp.abs(Gamma_tilde_lib - Gamma_tilde_manual)):.2e}")
  assert jnp.allclose(Gamma_tilde_lib, Gamma_tilde_manual, atol=1e-5), \
    f"Formula 1 failed! Max diff: {jnp.max(jnp.abs(Gamma_tilde_lib - Gamma_tilde_manual))}"
  print("  PASSED!")


def test_christoffel_change_of_basis_formula2():
  """
  Test Formula 2: E_i(T^k_j) = Γ̃^m_{aj} T^k_m T̂^a_i - T^m_j Γ^k_{im}

  This is a rearrangement of Formula 1, solving for the derivative of T.
  """
  key = random.PRNGKey(42)
  dim = 3

  # Create a point
  p = jnp.array([0.5, 0.3, 0.2])

  # Create a random positive definite metric
  key, subkey = random.split(key)
  A = random.normal(subkey, (dim, dim))
  metric_matrix = A @ A.T + 0.1 * jnp.eye(dim)

  key, subkey = random.split(key)
  metric_grad = random.normal(subkey, (dim, dim, dim)) * 0.1
  metric_grad = 0.5 * (metric_grad + jnp.swapaxes(metric_grad, 0, 1))

  std_basis = get_standard_basis(p)
  metric_jet = Jet(
    value=metric_matrix,
    gradient=metric_grad,
    hessian=jnp.zeros((dim, dim, dim, dim))
  )
  metric = RiemannianMetric(basis=std_basis, components=metric_jet)
  connection_old = get_levi_civita_connection(metric)

  # Create transformation T and its gradient
  key, subkey = random.split(key)
  T_matrix = random.normal(subkey, (dim, dim)) + 0.5 * jnp.eye(dim)

  key, subkey = random.split(key)
  T_grad = random.normal(subkey, (dim, dim, dim)) * 0.1

  # New basis
  new_basis = BasisVectors(
    p=p,
    components=Jet(
      value=T_matrix,
      gradient=T_grad,
      hessian=jnp.zeros((dim, dim, dim, dim))
    )
  )

  connection_new = change_basis(connection_old, new_basis)

  # Verify Formula 2: E_i(T^k_j) = Γ̃^m_{aj} T^k_m T̂^a_i - T^m_j Γ^k_{im}
  Gamma_old = connection_old.christoffel_symbols.value  # [i, m, k]
  Gamma_new = connection_new.christoffel_symbols.value  # [a, j, m]
  T_hat = jnp.linalg.inv(T_matrix)

  # LHS: E_i(T^k_j) = ∂_i(T^k_j) since E = I (standard basis)
  # T_grad has shape [k, j, i] where last index is derivative direction
  # So E_i(T^k_j) = T_grad[k, j, i]
  LHS = T_grad  # [k, j, i]

  # RHS Term 1: Γ̃^m_{aj} T^k_m T̂^a_i
  # Gamma_new indexed as [a, j, m]
  rhs_term1 = jnp.einsum("ajm,km,ai->kji", Gamma_new, T_matrix, T_hat)

  # RHS Term 2: -T^m_j Γ^k_{im}
  # Gamma_old indexed as [i, m, k]
  rhs_term2 = -jnp.einsum("mj,imk->kji", T_matrix, Gamma_old)

  RHS = rhs_term1 + rhs_term2

  print("\nFormula 2 Test:")
  print(f"  Max difference: {jnp.max(jnp.abs(LHS - RHS)):.2e}")
  assert jnp.allclose(LHS, RHS, atol=1e-5), \
    f"Formula 2 failed! Max diff: {jnp.max(jnp.abs(LHS - RHS))}"
  print("  PASSED!")


if __name__ == "__main__":
  test_christoffel_change_of_basis_formula1()
  test_christoffel_change_of_basis_formula2()
  print("\nAll tests passed!")

