"""
Sanity check for the Taylor series coefficient symmetries in normal_coords.py.

The Taylor series coefficients for x(v) should have the following symmetries
(because partial derivatives commute):
- d2xdv2[i,j,k] should be symmetric in j,k
- d3xdv3[i,j,k,l] should be symmetric in j,k,l
"""
import jax
import jax.numpy as jnp
from jax import random
from local_coordinates.metric import RiemannianMetric
from local_coordinates.basis import get_standard_basis
from local_coordinates.tensor import change_basis
from local_coordinates.connection import get_levi_civita_connection
from local_coordinates.riemann import get_riemann_curvature_tensor
from local_coordinates.metric import lower_index
from local_coordinates.jet import Jet
import pdb


def compute_taylor_coefficients_original(metric: RiemannianMetric):
  """
  Compute the Taylor series coefficients for x(v) at v=0 using the ORIGINAL
  formula from the notes (which may have errors).
  Returns d2xdv2 and d3xdv3.
  """
  standard_basis = get_standard_basis(metric.basis.p)
  metric = change_basis(metric, standard_basis)

  connection = get_levi_civita_connection(metric)
  curvature_endomorphism = get_riemann_curvature_tensor(connection)
  curvature_tensor = lower_index(curvature_endomorphism, metric, 4)
  gamma_bar = connection.christoffel_symbols
  R_bar = curvature_tensor.components

  # Construct orthonormal basis
  gij = metric.components.value
  eigenvalues, eigenvectors = jnp.linalg.eigh(gij)
  dxdv = jnp.einsum("ij,j->ij", eigenvectors, jax.lax.rsqrt(eigenvalues))

  # Second order: d2xdv2[i,j,k] = -Gamma^i_{ab} * dxdv[a,j] * dxdv[b,k]
  d2xdv2 = -jnp.einsum("abi,aj,bk->ijk", gamma_bar.value, dxdv, dxdv)

  # Third order terms (ORIGINAL - with minus signs on Gamma*Gamma terms)
  term1 = -jnp.einsum("abic,aj,bk,cl->ijkl", gamma_bar.gradient, dxdv, dxdv, dxdv)
  term2 = -jnp.einsum("mbi,cam,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)
  term3 = -jnp.einsum("ami,cbm,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)
  term4 = 1/3*jnp.einsum("mbac,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)
  term5 = 1/3*jnp.einsum("mcab,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)
  d3xdv3 = term1 + term2 + term3 + term4 + term5

  return dxdv, d2xdv2, d3xdv3, {
    "term1": term1,
    "term2": term2,
    "term3": term3,
    "term4": term4,
    "term5": term5,
    "gamma_bar": gamma_bar,
    "R_bar": R_bar,
  }


def compute_taylor_coefficients_corrected(metric: RiemannianMetric):
  """
  Compute the Taylor series coefficients with CORRECTED signs.

  The derivation shows:
  -Gamma * (-Gamma * ...) = +Gamma * Gamma * ...

  So terms 2 and 3 should have PLUS signs, not minus.
  """
  standard_basis = get_standard_basis(metric.basis.p)
  metric = change_basis(metric, standard_basis)

  connection = get_levi_civita_connection(metric)
  curvature_endomorphism = get_riemann_curvature_tensor(connection)
  curvature_tensor = lower_index(curvature_endomorphism, metric, 4)
  gamma_bar = connection.christoffel_symbols
  R_bar = curvature_tensor.components

  # Construct orthonormal basis
  gij = metric.components.value
  eigenvalues, eigenvectors = jnp.linalg.eigh(gij)
  dxdv = jnp.einsum("ij,j->ij", eigenvectors, jax.lax.rsqrt(eigenvalues))

  # Second order
  d2xdv2 = -jnp.einsum("abi,aj,bk->ijk", gamma_bar.value, dxdv, dxdv)

  # Third order terms (CORRECTED - with plus signs on Gamma*Gamma terms)
  term1 = -jnp.einsum("abic,aj,bk,cl->ijkl", gamma_bar.gradient, dxdv, dxdv, dxdv)
  term2 = +jnp.einsum("mbi,cam,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)
  term3 = +jnp.einsum("ami,cbm,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)
  term4 = 1/3*jnp.einsum("mbac,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)
  term5 = 1/3*jnp.einsum("mcab,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)
  d3xdv3 = term1 + term2 + term3 + term4 + term5

  return dxdv, d2xdv2, d3xdv3, {
    "term1": term1,
    "term2": term2,
    "term3": term3,
    "term4": term4,
    "term5": term5,
    "gamma_bar": gamma_bar,
    "R_bar": R_bar,
  }


def compute_from_derivation(metric: RiemannianMetric):
  """
  Compute the Taylor series coefficients by directly implementing
  each step of the derivation to check for errors.

  From the derivation:
  d3x/dv^3 = d/dv_l [ -Gamma^i_ab * J^a_j * J^b_k + Gamma^a_jk * J^i_a ]

  At v=0, Gamma_v = 0, so:
  = -dGamma^i_ab/dx^c * J^c_l * J^a_j * J^b_k
    + Gamma^i_ab * Gamma^a_mn * J^m_l * J^n_j * J^b_k
    + Gamma^i_ab * Gamma^b_mn * J^a_j * J^m_l * J^n_k
    + dGamma^a_jk/dv^l * J^i_a

  where dGamma^a_jk/dv^l = (1/3)(R^a_jkl + R^a_kjl) in v-coordinates.
  """
  standard_basis = get_standard_basis(metric.basis.p)
  metric = change_basis(metric, standard_basis)

  connection = get_levi_civita_connection(metric)
  curvature_endo = get_riemann_curvature_tensor(connection)  # R^i_jkl
  curvature_tensor = lower_index(curvature_endo, metric, 4)  # R_ijkl
  gamma_bar = connection.christoffel_symbols
  R_endo = curvature_endo.components  # R^i_jkl (endomorphism form)
  R_bar = curvature_tensor.components  # R_ijkl (fully lowered)

  gij = metric.components.value
  eigenvalues, eigenvectors = jnp.linalg.eigh(gij)
  dxdv = jnp.einsum("ij,j->ij", eigenvectors, jax.lax.rsqrt(eigenvalues))
  # dxdv[i,j] = dx^i/dv^j (Jacobian from v to x)

  # Term 1: -dGamma^i_ab/dx^c * J^a_j * J^b_k * J^c_l
  # gamma_bar.gradient[a,b,i,c] = dGamma^i_ab/dx^c
  term1 = -jnp.einsum("abic,aj,bk,cl->ijkl", gamma_bar.gradient, dxdv, dxdv, dxdv)

  # Term 2: +Gamma^i_ab * Gamma^a_mn * J^m_l * J^n_j * J^b_k
  # Relabeling: a->m, m->c, n->a gives Gamma^i_mb * Gamma^m_ca * J^c_l * J^a_j * J^b_k
  term2 = +jnp.einsum("mbi,cam,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)

  # Term 3: +Gamma^i_ab * Gamma^b_mn * J^a_j * J^m_l * J^n_k
  # Relabeling: b->m, m->c, n->b gives Gamma^i_am * Gamma^m_cb * J^a_j * J^c_l * J^b_k
  term3 = +jnp.einsum("ami,cbm,aj,bk,cl->ijkl", gamma_bar.value, gamma_bar.value, dxdv, dxdv, dxdv)

  # Term 4: dGamma^a_jk/dv^l * J^i_a = (1/3)(R^a_jkl + R^a_kjl) * J^i_a
  #
  # The Riemann curvature endomorphism R^a_bcd in v-coords relates to x-coords by:
  #   R^a_jkl (v-coords) = R^m_npq (x-coords) * J^a_m * (J^{-1})^n_j * (J^{-1})^p_k * (J^{-1})^q_l
  #
  # But since J is orthogonal (J^T J = I when g_v = delta), J^{-1} = J^T, so:
  #   R^a_jkl * J^i_a = R^m_npq * J^a_m * J^i_a * (J^T)^n_j * (J^T)^p_k * (J^T)^q_l
  #                   = R^m_npq * delta^{mi} * J^j_n * J^k_p * J^l_q  (orthogonality)
  #                   = R^i_npq * J^j_n * J^k_p * J^l_q
  #
  # Wait, that's not right either. Let me be more careful about the index structure.
  # Actually for orthonormal frames: J^a_m * J^i_a = (J J^T)^{im} which is the inverse
  # metric in x-coords evaluated at the point.

  # Let's just compute this term using the endomorphism form directly.
  # R_endo.value[i,j,k,l] = R^i_jkl (upper first index)
  # The term is: (1/3)(R^a_jkl + R^a_kjl) * J^i_a
  # In Einstein notation: R^a_bcd * J^i_a -> need to contract a with first index

  # R_endo.value has shape (n,n,n,n) with R_endo.value[a,b,c,d] = R^a_bcd
  # Wait, let me check the convention in this codebase...
  # Actually, looking at riemann.py, the endomorphism form is R^a_bcd

  # The curvature term contribution:
  # (1/3)(R^a_jkl + R^a_kjl) J^i_a
  # = (1/3) R^a_jkl J^i_a + (1/3) R^a_kjl J^i_a

  # In v-coordinates, R^a_jkl transforms from x-coordinates as:
  # R^a_jkl = R^m_npq * J^a_m * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l

  # So R^a_jkl * J^i_a = R^m_npq * (J^a_m J^i_a) * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l

  # Now, J^a_m J^i_a = sum_a J[a,m] * J[i,a] = (J^T J)[m,i]
  # For orthonormal frame, g_v = I, so J^T g_x J = I, meaning J^T J = g_x^{-1}

  # Actually this is getting complicated. Let me just try a simpler approach:
  # use the identity J^i_a (J^{-1})^a_m = delta^i_m

  # Actually, the simplest thing is to work in x-coordinates throughout.
  # The term dGamma^a_jk/dv^l * J^i_a comes from differentiating Gamma^a_jk (in v-coords).
  # In v-coords at p: Gamma=0, dGamma = curvature terms.

  # From the formula for Christoffel derivatives in normal coords:
  # d Gamma^k_ij / dv^l = (1/3)(R^k_ijl + R^k_jil) at the origin
  #
  # So the contribution to d3x/dv3 is:
  # (1/3)(R^a_jkl + R^a_kjl) * J^i_a
  #
  # Now, R^a_jkl in v-coords can be expressed using the Jacobian transformation.
  # But since we're at the origin where the metric is Euclidean in v-coords,
  # R^a_jkl (v) is related to R^m_npq (x) by tensor transformation rules.

  # For the curvature endomorphism: R^a_bcd transforms as
  # R^a_bcd (v) = J^a_m * (J^-1)^n_b * (J^-1)^p_c * (J^-1)^q_d * R^m_npq (x)

  # For orthonormal frame, J^-1 = J^T (since J^T g_x J = I and g_v = I at origin)
  # Wait, that's not right. g_v = I means J^T g_x J = I, so J^T = J^{-1} g_x^{-1}

  # Hmm, let me just try computing the term using the fully-lowered R and raising index.
  # Or better yet, use the endomorphism directly.

  # R_endo = R^i_jkl, so:
  # term4_part1 = (1/3) R^a_jkl J^i_a = (1/3) R_endo[a,j,k,l] * dxdv[i,a]
  # But wait, j,k,l here are v-coordinate indices, and R_endo is in x-coordinates.
  # We need to transform.

  # Actually I think the issue is the notes' formula conflates coordinate systems.
  # The cleanest approach is to compute dGamma^a_jk/dv^l directly from the
  # Christoffel transformation formula.

  # For now, let me just try the interpretation that the curvature term should
  # give something symmetric. The formula (R^a_jkl + R^a_kjl) is symmetric in j,k
  # but what about l?

  # Using the symmetry R^a_bjkl = -R^a_bkjl:
  # R^a_jkl + R^a_kjl... these have indices (j,k,l) and (k,j,l), both contracted
  # with Jacobians. For j<->k symmetry, we get R^a_kjl + R^a_jkl = same. Good.
  # For j<->l symmetry: R^a_lkj + R^a_klj vs R^a_jkl + R^a_kjl... not the same.

  # So the curvature term as written is NOT symmetric in (j,k,l)!
  # This means either:
  # 1. The other terms (Gamma derivatives, Gamma*Gamma) must cancel this asymmetry
  # 2. The formula is wrong

  # Let me compute all terms separately and check if they sum to something symmetric.

  ginv = jnp.linalg.inv(gij)  # g^{ij} in x-coordinates
  dvdx = jnp.linalg.inv(dxdv)  # J^{-1}: dv/dx

  # The curvature term from differentiating Gamma_v:
  # d/dv^l (Gamma^a_jk) = (1/3)(R^a_jkl + R^a_kjl) evaluated in v-coords at origin
  #
  # Using tensor transformation:
  # R^a_jkl (v) = J^a_m * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l * R^m_npq (x)
  #
  # Then: (1/3)(R^a_jkl + R^a_kjl) * J^i_a
  #     = (1/3) J^i_a J^a_m (J^-1)^n_j (J^-1)^p_k (J^-1)^q_l (R^m_npq + R^m_pnq)
  #     = (1/3) delta^i_m (J^-1)^n_j (J^-1)^p_k (J^-1)^q_l (R^i_npq + R^i_pnq)
  #
  # Wait, J^i_a J^a_m = (J J^T)[i,m], which equals delta^im only if J is orthogonal.
  # For an orthonormal frame, J^T g J = I, so J J^T = g^{-1}.
  # Thus J^i_a J^a_m = g^{im}.

  # So the curvature term becomes:
  # (1/3) g^{im} (J^-1)^n_j (J^-1)^p_k (J^-1)^q_l (R^m_npq + R^m_pnq)
  #
  # Hmm, this still involves the inverse Jacobian which contracts with (j,k,l).
  # The formula in the notes has J contracted with (j,k,l), not J^{-1}!

  # I think the notes have the wrong Jacobian direction in the curvature term.

  # Let me try both versions:

  # Version A: Notes' formula (J for curvature term, with R_abcd fully lowered)
  term4_notes = 1/3*jnp.einsum("mbac,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)
  term5_notes = 1/3*jnp.einsum("mcab,mi,aj,bk,cl->ijkl", R_bar.value, dxdv, dxdv, dxdv, dxdv)

  # Version B: Correct transformation of the curvature term
  #
  # Starting from: dGamma^a_jk/dv^l * J^i_a = (1/3)(R^a_jkl + R^a_kjl) * J^i_a
  # where R^a_jkl is the curvature endomorphism in v-coords.
  #
  # Transform R^a_jkl from v to x coords:
  # R^a_jkl (v) = J^a_m * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l * R^m_npq (x)
  #
  # Then: R^a_jkl * J^i_a
  #     = J^i_a * J^a_m * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l * R^m_npq (x)
  #     = delta^i_m * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l * R^m_npq (x)
  #     = R^i_npq (x) * (J^-1)^n_j * (J^-1)^p_k * (J^-1)^q_l
  #
  # Wait, J^i_a * J^a_m = (J J^T)^{im}, NOT delta^{im}!
  # For orthonormal frame: J^T g_x J = I (in v-coords g_v=I), so J J^T = g_x^{-1}.
  # Thus J^i_a * J^a_m = g^{im}_x.
  #
  # Hmm, but actually I think the issue is different. Let me reconsider...
  #
  # The key insight: dv/dx * dx/dv = I, so:
  # J^i_a * (J^-1)^a_m = delta^i_m
  #
  # The curvature term is:
  # (1/3)(R^a_jkl + R^a_kjl) (v-coords) * J^i_a (converts a from v to x, gives i)
  #
  # For R^a_jkl in v-coords, we have:
  # R^a_jkl (v) = (J^-1)^a_m * J^n_j * J^p_k * J^q_l * R^m_npq (x)
  #
  # Note the inverse on the upper index! So:
  # R^a_jkl * J^i_a = J^i_a * (J^-1)^a_m * J^n_j * J^p_k * J^q_l * R^m_npq
  #                 = delta^i_m * J^n_j * J^p_k * J^q_l * R^m_npq
  #                 = R^i_npq * J^n_j * J^p_k * J^q_l
  #
  # Relabeling n->a, p->b, q->c:
  # = R^i_abc * J^a_j * J^b_k * J^c_l
  #
  # So the curvature term should be:
  # (1/3)(R^i_abc + R^i_bac) * J^a_j * J^b_k * J^c_l
  #
  # where R^i_abc is the curvature ENDOMORPHISM in x-coords (upper first index).

  # R_endo.value[i,a,b,c] = R^i_abc (endomorphism form)
  term4_correct = 1/3*jnp.einsum("iabc,aj,bk,cl->ijkl", R_endo.value, dxdv, dxdv, dxdv)
  term5_correct = 1/3*jnp.einsum("ibac,aj,bk,cl->ijkl", R_endo.value, dxdv, dxdv, dxdv)

  d3xdv3_notes = term1 + term2 + term3 + term4_notes + term5_notes
  d3xdv3_correct = term1 + term2 + term3 + term4_correct + term5_correct

  return dxdv, d3xdv3_notes, d3xdv3_correct, {
    "term1": term1,
    "term2": term2,
    "term3": term3,
    "term4_notes": term4_notes,
    "term5_notes": term5_notes,
    "term4_correct": term4_correct,
    "term5_correct": term5_correct,
    "gamma_bar": gamma_bar,
    "R_bar": R_bar,
    "R_endo": R_endo,
    "ginv": ginv,
    "dvdx": dvdx,
  }


def check_symmetry_2nd_order(d2xdv2):
  """Check that d2xdv2[i,j,k] is symmetric in j,k."""
  diff = d2xdv2 - jnp.transpose(d2xdv2, (0, 2, 1))  # Swap j,k
  max_diff = jnp.max(jnp.abs(diff))
  return max_diff


def symmetrize_last_3_indices(tensor):
  """Symmetrize a tensor[i,j,k,l] in the last 3 indices (j,k,l)."""
  perms = [
    (0, 1, 2, 3),  # i,j,k,l
    (0, 1, 3, 2),  # i,j,l,k
    (0, 2, 1, 3),  # i,k,j,l
    (0, 2, 3, 1),  # i,k,l,j
    (0, 3, 1, 2),  # i,l,j,k
    (0, 3, 2, 1),  # i,l,k,j
  ]
  result = jnp.zeros_like(tensor)
  for perm in perms:
    result = result + jnp.transpose(tensor, perm)
  return result / 6


def check_symmetry_3rd_order(d3xdv3, verbose=False):
  """Check that d3xdv3[i,j,k,l] is symmetric in j,k,l."""
  # Check all 6 permutations of (j,k,l)
  perms = [
    ((0, 1, 2, 3), "i,j,k,l"),  # original
    ((0, 1, 3, 2), "i,j,l,k"),  # swap k,l
    ((0, 2, 1, 3), "i,k,j,l"),  # swap j,k
    ((0, 2, 3, 1), "i,k,l,j"),  # cycle
    ((0, 3, 1, 2), "i,l,j,k"),  # cycle
    ((0, 3, 2, 1), "i,l,k,j"),  # reverse
  ]

  max_diff = 0.0
  original = d3xdv3
  if verbose:
    print("    Pairwise asymmetries:")
  for perm, name in perms[1:]:
    permuted = jnp.transpose(d3xdv3, perm)
    diff = jnp.max(jnp.abs(original - permuted))
    max_diff = max(max_diff, float(diff))
    if verbose:
      print(f"      {name}: {diff:.2e}")

  return max_diff


def create_random_metric(key, dim=3):
  """Create a random Riemannian metric with non-trivial curvature."""
  k1, k2, k3, k4 = random.split(key, 4)

  # Random point
  p = random.normal(k1, (dim,))

  # Random positive definite metric
  A = random.normal(k2, (dim, dim))
  g_value = A @ A.T + jnp.eye(dim)

  # Random gradient of metric (symmetric in first two indices)
  g_grad = random.normal(k3, (dim, dim, dim))
  g_grad = (g_grad + jnp.transpose(g_grad, (1, 0, 2))) / 2

  # Random hessian of metric (symmetric in first two indices and last two indices)
  g_hess = random.normal(k4, (dim, dim, dim, dim))
  g_hess = (g_hess + jnp.transpose(g_hess, (1, 0, 2, 3))) / 2  # sym in i,j
  g_hess = (g_hess + jnp.transpose(g_hess, (0, 1, 3, 2))) / 2  # sym in k,l

  basis = get_standard_basis(p)
  components = Jet(value=g_value, gradient=g_grad, hessian=g_hess)

  return RiemannianMetric(basis=basis, components=components)


def verify_christoffel_derivative_formula(metric: RiemannianMetric):
  """
  Verify the formula for dGamma^k_ij/dv^l at the origin of RNC.

  The notes claim (eq 201):
    dGamma^k_ij/dv^l = (1/3)(R_{kijl} + R_{kjil})

  But we can also derive it from the definition:
    Gamma^k_ij = (1/2) g^{km} (d_i g_{mj} + d_j g_{mi} - d_m g_{ij})

  At origin: g = I, dg = 0, so:
    d_l Gamma^k_ij = (1/2)(d_l d_i g_{kj} + d_l d_j g_{ki} - d_l d_k g_{ij})

  Using d_a d_b g_{ij} = (1/3)(R_{aibj} + R_{biaj}):
    d_l Gamma^k_ij = (1/6)(R_{lkij} + R_{iklj} + R_{lkji} + R_{jkli} - R_{likj} - R_{kilj})
  """
  from local_coordinates.tensor import change_basis
  from local_coordinates.normal_coords import to_riemann_normal_coordinates

  # Transform metric to RNC
  rnc_metric = to_riemann_normal_coordinates(metric)

  # Get the Riemann tensor in RNC (should be same as in x-coords at origin)
  connection = get_levi_civita_connection(rnc_metric)
  curv_endo = get_riemann_curvature_tensor(connection)
  curv_tensor = lower_index(curv_endo, rnc_metric, 4)
  R = curv_tensor.components.value  # R_{ijkl}

  # Get Christoffel derivative in RNC
  gamma = connection.christoffel_symbols
  dGamma = gamma.gradient  # dGamma[i,j,k,l] = d_l Gamma^k_{ij}

  n = R.shape[0]

  # Compute dGamma from formula (notes eq 201): (1/3)(R_{kijl} + R_{kjil})
  # Note: R_{kijl} means first index is k (lowered from upper), then i,j,l
  dGamma_formula1 = (1/3) * (R + jnp.transpose(R, (0, 2, 1, 3)))  # R_{kijl} + R_{kjil}
  # But wait, our R has indices R[a,b,c,d] = R_{abcd}
  # So R_{kijl} = R[k,i,j,l] and R_{kjil} = R[k,j,i,l]
  # dGamma^k_ij/dv^l -> need to match index ordering

  # Actually, let me compute from the definition instead
  # d_l Gamma^k_ij = (1/2)(d_l d_i g_{kj} + d_l d_j g_{ki} - d_l d_k g_{ij})
  # Using d_a d_b g_{ij} = (1/3)(R_{aibj} + R_{biaj})

  # d_l d_i g_{kj} has derivative indices (l,i) and metric indices (k,j)
  # So it's (1/3)(R_{lkij} + R_{iklj})
  d2g_li_kj = (1/3) * (jnp.transpose(R, (3, 0, 1, 2)) + jnp.transpose(R, (1, 0, 3, 2)))
  # R[l,k,i,j] + R[i,k,l,j] -> wait this indexing is confusing

  # Let me be more careful. R_{abcd} is stored as R[a,b,c,d]
  # R_{lkij} = R[l,k,i,j]
  # R_{iklj} = R[i,k,l,j]

  # d2g/dv^l dv^i evaluated at g_{kj} = (1/3)(R_{lkij} + R_{iklj})
  # This should give a tensor with shape (n,n,n,n) where indices are (l,i,k,j)

  # Let's construct it properly
  # d_l d_i g_{kj} = (1/3)(R_{lkij} + R_{iklj})
  # We need R[l,k,i,j] + R[i,k,l,j]
  term_li_kj = (1/3) * (R + jnp.einsum("abcd->cbad", R))  # This gives R_{lkij} + R_{iklj} for l,k,i,j
  # Wait, R[a,b,c,d] has indices (a,b,c,d) = (l,k,i,j) for the first term
  # For second term R_{iklj} = R[i,k,l,j], we need to rearrange from R[l,k,i,j] to R[i,k,l,j]
  # That's (l,k,i,j) -> (i,k,l,j), or (0,1,2,3) -> (2,1,0,3)
  term_li_kj = (1/3) * (R + jnp.transpose(R, (2, 1, 0, 3)))

  # Similarly:
  # d_l d_j g_{ki} = (1/3)(R_{lkji} + R_{jkli})
  # R_{lkji} = R[l,k,j,i] and R_{jkli} = R[j,k,l,i]
  # From base R[l,k,i,j]: R[l,k,j,i] needs (0,1,3,2), R[j,k,l,i] needs (3,1,0,2)
  term_lj_ki = (1/3) * (jnp.transpose(R, (0, 1, 3, 2)) + jnp.transpose(R, (3, 1, 0, 2)))

  # d_l d_k g_{ij} = (1/3)(R_{likj} + R_{kilj})
  # R_{likj} = R[l,i,k,j] and R_{kilj} = R[k,i,l,j]
  # From base R[l,k,i,j]: R[l,i,k,j] needs (0,2,1,3), R[k,i,l,j] needs (1,2,0,3)
  term_lk_ij = (1/3) * (jnp.transpose(R, (0, 2, 1, 3)) + jnp.transpose(R, (1, 2, 0, 3)))

  # Now: d_l Gamma^k_ij = (1/2)(d_l d_i g_{kj} + d_l d_j g_{ki} - d_l d_k g_{ij})
  # term_li_kj has indices (l,k,i,j) for d_l d_i g_{kj}
  # term_lj_ki has indices (l,k,j,i) for d_l d_j g_{ki} -> need to swap i,j -> (l,k,i,j) via (0,1,3,2)
  # term_lk_ij has indices (l,i,k,j) for d_l d_k g_{ij} -> need (l,k,i,j) via (0,2,1,3)

  # Hmm this is getting confusing. Let me just compute it element by element
  dGamma_from_def = jnp.zeros((n, n, n, n))

  def compute_dGamma_elem(l, k, i, j):
    # d_l Gamma^k_ij = (1/2)(d_l d_i g_{kj} + d_l d_j g_{ki} - d_l d_k g_{ij})
    # d_a d_b g_{cd} = (1/3)(R_{acbd} + R_{bcad})
    d2g_li_kj = (1/3) * (R[l, k, i, j] + R[i, k, l, j])
    d2g_lj_ki = (1/3) * (R[l, k, j, i] + R[j, k, l, i])
    d2g_lk_ij = (1/3) * (R[l, i, k, j] + R[k, i, l, j])
    return 0.5 * (d2g_li_kj + d2g_lj_ki - d2g_lk_ij)

  # Vectorize computation
  ls, ks, iis, js = jnp.meshgrid(jnp.arange(n), jnp.arange(n), jnp.arange(n), jnp.arange(n), indexing='ij')

  # d_l d_i g_{kj} = (1/3)(R_{lkij} + R_{iklj})
  d2g_li_kj = (1/3) * (R[ls, ks, iis, js] + R[iis, ks, ls, js])
  # d_l d_j g_{ki} = (1/3)(R_{lkji} + R_{jkli})
  d2g_lj_ki = (1/3) * (R[ls, ks, js, iis] + R[js, ks, ls, iis])
  # d_l d_k g_{ij} = (1/3)(R_{likj} + R_{kilj})
  d2g_lk_ij = (1/3) * (R[ls, iis, ks, js] + R[ks, iis, ls, js])

  dGamma_from_def = 0.5 * (d2g_li_kj + d2g_lj_ki - d2g_lk_ij)
  # dGamma_from_def[l,k,i,j] = d_l Gamma^k_ij

  # Compare to notes formula: (1/3)(R_{kijl} + R_{kjil})
  # R_{kijl} = R[k,i,j,l] and R_{kjil} = R[k,j,i,l]
  # For dGamma^k_ij/dv^l, this should give tensor with indices (k,i,j,l)
  dGamma_notes = (1/3) * (jnp.transpose(R, (0, 1, 2, 3)) + jnp.transpose(R, (0, 2, 1, 3)))
  # dGamma_notes[k,i,j,l] = (1/3)(R_{kijl} + R_{kjil})

  # Rearrange to (l,k,i,j) for comparison
  dGamma_notes_rearranged = jnp.transpose(dGamma_notes, (3, 0, 1, 2))

  # Compare
  diff = jnp.max(jnp.abs(dGamma_from_def - dGamma_notes_rearranged))

  return {
    "dGamma_from_def": dGamma_from_def,
    "dGamma_notes": dGamma_notes_rearranged,
    "diff": diff,
    "R": R,
  }


def compute_third_derivative_autodiff(metric: RiemannianMetric):
  """
  Compute d3x/dv3 using autodiff as ground truth.

  The key is to differentiate the second derivative formula with respect to v,
  where the second derivative depends on x(v) through the Christoffel symbols.
  """
  from local_coordinates.tensor import change_basis

  standard_basis = get_standard_basis(metric.basis.p)
  metric_std = change_basis(metric, standard_basis)

  connection = get_levi_civita_connection(metric_std)
  gamma = connection.christoffel_symbols

  gij = metric_std.components.value
  eigenvalues, eigenvectors = jnp.linalg.eigh(gij)
  J = jnp.einsum("ij,j->ij", eigenvectors, jax.lax.rsqrt(eigenvalues))  # dx/dv at v=0

  # The second derivative at v=0 is:
  # d2x^i/dv^j dv^k = -Gamma^i_ab(p) * J^a_j * J^b_k
  #
  # For the third derivative, we need to differentiate this w.r.t. v^l.
  # The key is that Gamma^i_ab depends on x, and x depends on v.
  #
  # Using chain rule:
  # d/dv^l [d2x^i/dv^j dv^k] = -d(Gamma^i_ab)/dx^c * dx^c/dv^l * J^a_j * J^b_k
  #                            -Gamma^i_ab * d(J^a_j)/dv^l * J^b_k
  #                            -Gamma^i_ab * J^a_j * d(J^b_k)/dv^l
  #
  # At v=0:
  # - dx^c/dv^l = J^c_l
  # - d(J^a_j)/dv^l = d2x^a/dv^l dv^j = -Gamma^a_mn(p) * J^m_l * J^n_j
  #
  # But we also need to account for Gamma_v (Christoffel in v-coords).
  # The full formula for d2x/dv2 is:
  # d2x^i/dv^j dv^k = -Gamma^i_ab(x) * J^a_j * J^b_k + Gamma^m_jk(v) * J^i_m
  #
  # Differentiating and evaluating at v=0 where Gamma(v)=0:
  # d3x^i/dv^j dv^k dv^l =
  #   -dGamma^i_ab/dx^c * J^c_l * J^a_j * J^b_k       (term 1)
  #   +Gamma^i_ab * Gamma^a_mn * J^m_l * J^n_j * J^b_k  (term 2: -Gamma*(-Gamma))
  #   +Gamma^i_ab * Gamma^b_mn * J^a_j * J^m_l * J^n_k  (term 3: -Gamma*(-Gamma))
  #   +dGamma^m_jk(v)/dv^l * J^i_m                     (term 4: curvature)

  # Let me verify using JAX autodiff by computing derivatives of x(v)
  # defined implicitly through the geodesic equation.

  # Define Gamma as a function of position (Taylor expansion around p)
  p = metric_std.basis.p
  Gamma_val = gamma.value  # Gamma^i_ab at p
  Gamma_grad = gamma.gradient  # dGamma^i_ab/dx^c at p

  def gamma_at_x(x):
    """Christoffel symbols at position x, using Taylor expansion around p."""
    dx = x - p
    # Gamma(x) ≈ Gamma(p) + dGamma/dx * dx + O(dx^2)
    return Gamma_val + jnp.einsum("abic,c->abi", Gamma_grad, dx)

  def x_of_v(v):
    """
    Compute x(v) = exp_p(v * J) iteratively using the geodesic equation.
    For small v, we can use Taylor expansion.
    """
    # x(v) ≈ p + J.v + (1/2) * d2x/dv2 * v * v + ...
    # This is circular, so instead we use:
    # x(v) ≈ p + J.v (first order)
    # Then d2x/dv2 ≈ -Gamma(p) . J . J (since x ≈ p)
    x1 = p + jnp.einsum("ij,j->i", J, v)
    return x1

  def d2x_dv2_func(v):
    """Second derivative d2x/dv^j dv^k as a function of v."""
    x = x_of_v(v)
    Gamma_x = gamma_at_x(x)
    # d2x/dv^j dv^k = -Gamma^i_ab(x) * J^a_j * J^b_k
    # But we also need to account for how J changes with v...
    # This is getting complicated. Let's just verify at v=0.
    return -jnp.einsum("abi,aj,bk->ijk", Gamma_x, J, J)

  # Compute d3x/dv3 at v=0 using JAX jacobian
  v0 = jnp.zeros_like(p)

  # d(d2x[i,j,k])/dv[l] gives d3x[i,j,k,l]
  d3x_autodiff = jax.jacfwd(d2x_dv2_func)(v0)  # shape (n,n,n,n)

  # Check symmetry
  sym_jk = jnp.max(jnp.abs(d3x_autodiff - jnp.transpose(d3x_autodiff, (0, 2, 1, 3))))
  sym_kl = jnp.max(jnp.abs(d3x_autodiff - jnp.transpose(d3x_autodiff, (0, 1, 3, 2))))
  sym_jl = jnp.max(jnp.abs(d3x_autodiff - jnp.transpose(d3x_autodiff, (0, 3, 2, 1))))

  return {
    "d3x_autodiff": d3x_autodiff,
    "sym_jk": sym_jk,
    "sym_kl": sym_kl,
    "sym_jl": sym_jl,
    "J": J,
    "gamma": gamma,
  }


def main():
  print("=" * 60)
  print("Symmetry Check for Taylor Series Coefficients")
  print("=" * 60)

  key = random.PRNGKey(42)

  for dim in [2, 3]:
    print(f"\n{'='*60}")
    print(f"Dimension: {dim}")
    print("=" * 60)

    key, subkey = random.split(key)
    metric = create_random_metric(subkey, dim=dim)

    # Test ORIGINAL formula (from notes)
    print("\n--- ORIGINAL formula (from notes) ---")
    dxdv, d2xdv2, d3xdv3_orig, terms_orig = compute_taylor_coefficients_original(metric)

    sym_2nd = check_symmetry_2nd_order(d2xdv2)
    sym_3rd_orig = check_symmetry_3rd_order(d3xdv3_orig)

    print(f"  d2xdv2 symmetry in (j,k): {sym_2nd:.2e}")
    print(f"  d3xdv3 symmetry in (j,k,l): {sym_3rd_orig:.2e}")

    # Test CORRECTED formula (with + signs on Gamma*Gamma)
    print("\n--- CORRECTED signs (+ on Gamma*Gamma) ---")
    _, _, d3xdv3_corr, terms_corr = compute_taylor_coefficients_corrected(metric)

    sym_3rd_corr = check_symmetry_3rd_order(d3xdv3_corr)
    print(f"  d3xdv3 symmetry in (j,k,l): {sym_3rd_corr:.2e}")

    # Test with derivation-based computation
    print("\n--- DERIVATION-BASED (corrected curvature term) ---")
    _, d3_notes, d3_correct, terms_deriv = compute_from_derivation(metric)

    sym_notes = check_symmetry_3rd_order(d3_notes)
    sym_correct = check_symmetry_3rd_order(d3_correct)

    print(f"  d3xdv3 (notes curvature term): {sym_notes:.2e}")
    print(f"  d3xdv3 (corrected curvature term): {sym_correct:.2e}")

    # Individual term analysis
    print("\n  Individual term symmetries:")
    for name in ["term1", "term2", "term3"]:
      term = terms_deriv[name]
      sym = check_symmetry_3rd_order(term)
      print(f"    {name}: {sym:.2e}")

    curv_notes = terms_deriv["term4_notes"] + terms_deriv["term5_notes"]
    curv_correct = terms_deriv["term4_correct"] + terms_deriv["term5_correct"]
    print(f"    curv (notes):    {check_symmetry_3rd_order(curv_notes):.2e}")
    print(f"    curv (corrected): {check_symmetry_3rd_order(curv_correct):.2e}")

    # Check if term1+term2+term3 is symmetric (it should be if Gamma terms cancel curvature)
    non_curv = terms_deriv["term1"] + terms_deriv["term2"] + terms_deriv["term3"]
    print(f"    term1+2+3:       {check_symmetry_3rd_order(non_curv):.2e}")

    # Summary
    print("\n  SUMMARY:")
    print(f"    Original (notes):     {sym_3rd_orig:.2e}")
    print(f"    Sign correction only: {sym_3rd_corr:.2e}")
    print(f"    Full correction:      {sym_correct:.2e}")

    # Check: does symmetrizing give something reasonable?
    d3_sym = symmetrize_last_3_indices(d3_correct)
    diff_from_sym = jnp.max(jnp.abs(d3_correct - d3_sym))
    print(f"\n    |d3 - symmetrized(d3)|: {diff_from_sym:.2e}")

    # Verify the Christoffel derivative formula
    print("\n  Christoffel derivative formula verification:")
    result = verify_christoffel_derivative_formula(metric)
    print(f"    |dGamma(def) - dGamma(notes)| = {result['diff']:.2e}")

    # Check if the internal tensor T^i_abc is symmetric in (a,b,c)
    print("\n  Internal tensor T^i_abc symmetry check:")
    gamma = terms_deriv["gamma_bar"]
    R_endo = terms_deriv["R_endo"]

    # T^i_abc = -dGamma^i_ab/dx_c + Gamma^i_mb*Gamma^m_ca + Gamma^i_am*Gamma^m_cb
    #           + (1/3)(R^i_abc + R^i_bac)
    # gamma.gradient[a,b,i,c] = dGamma^i_ab/dx_c
    # gamma.value[a,b,i] = Gamma^i_ab
    # R_endo.value[i,a,b,c] = R^i_abc

    n = dim
    # Term 1: -dGamma^i_ab/dx_c  (indices: i,a,b,c)
    T1 = -jnp.transpose(gamma.gradient, (2, 0, 1, 3))  # [a,b,i,c] -> [i,a,b,c]

    # Term 2: Gamma^i_mb * Gamma^m_ca (sum over m)
    # gamma.value[m,b,i] * gamma.value[c,a,m] -> [i,a,b,c]
    T2 = jnp.einsum("mbi,cam->iabc", gamma.value, gamma.value)

    # Term 3: Gamma^i_am * Gamma^m_cb (sum over m)
    T3 = jnp.einsum("ami,cbm->iabc", gamma.value, gamma.value)

    # Term 4+5: (1/3)(R^i_abc + R^i_bac)
    T4 = (1/3) * R_endo.value
    T5 = (1/3) * jnp.transpose(R_endo.value, (0, 2, 1, 3))  # swap a,b

    T = T1 + T2 + T3 + T4 + T5

    # Check symmetry in (a,b,c) - the last 3 indices
    T_sym_ab = jnp.max(jnp.abs(T - jnp.transpose(T, (0, 2, 1, 3))))
    T_sym_bc = jnp.max(jnp.abs(T - jnp.transpose(T, (0, 1, 3, 2))))
    T_sym_ac = jnp.max(jnp.abs(T - jnp.transpose(T, (0, 3, 2, 1))))

    print(f"    T^i_abc sym in (a,b): {T_sym_ab:.2e}")
    print(f"    T^i_abc sym in (b,c): {T_sym_bc:.2e}")
    print(f"    T^i_abc sym in (a,c): {T_sym_ac:.2e}")

    # Individually check each term's symmetry
    print("    Individual terms:")
    for name, term in [("T1 (grad)", T1), ("T2 (G*G)", T2), ("T3 (G*G)", T3),
                        ("T4+T5 (curv)", T4+T5), ("T (total)", T)]:
      sym_ab = jnp.max(jnp.abs(term - jnp.transpose(term, (0, 2, 1, 3))))
      sym_bc = jnp.max(jnp.abs(term - jnp.transpose(term, (0, 1, 3, 2))))
      print(f"      {name}: (a,b)={sym_ab:.2e}, (b,c)={sym_bc:.2e}")

    # The gradient term T1 = -dGamma^i_ab/dx_c should relate to curvature via:
    # dGamma^i_ab/dx_c - dGamma^i_ac/dx_b = R^i_acb + Gamma*Gamma terms
    # At the origin (Gamma=0): dGamma^i_ab/dx_c - dGamma^i_ac/dx_b = R^i_acb
    #
    # This means T1 is NOT symmetric in (b,c), with asymmetry = -R^i_abc
    # Let's verify this:
    T1_asym_bc = T1 - jnp.transpose(T1, (0, 1, 3, 2))  # T1[i,a,b,c] - T1[i,a,c,b]
    # This should equal -R^i_acb + R^i_abc = R^i_abc - R^i_acb = R^i_abc + R^i_abc = 2*R^i_abc...
    # wait that's not right. Let me think again.
    #
    # dGamma^i_ab/dx_c - dGamma^i_ac/dx_b = R^i_abc (at origin, Gamma=0)
    # So -dGamma^i_ab/dx_c + dGamma^i_ac/dx_b = -R^i_abc
    # i.e. T1[i,a,b,c] - T1[i,a,c,b] = -R^i_abc

    # Check: T1_asym_bc should equal -R^i_abc
    diff_T1_R = jnp.max(jnp.abs(T1_asym_bc + R_endo.value))
    print(f"\n    Checking T1 asymmetry = -R^i_abc:")
    print(f"      |T1[i,a,b,c] - T1[i,a,c,b] + R^i_abc| = {diff_T1_R:.2e}")

    # The curvature term is (1/3)(R^i_abc + R^i_bac)
    # Its (b,c) asymmetry is:
    # (1/3)(R^i_abc + R^i_bac) - (1/3)(R^i_acb + R^i_cab)
    # = (1/3)(R^i_abc - R^i_acb + R^i_bac - R^i_cab)
    # Using R^i_abc = -R^i_acb: R^i_abc - R^i_acb = 2*R^i_abc
    # Using R^i_bac = -R^i_bca: ...this is getting complicated

    # For the total to be symmetric, we need:
    # T1_asym + T2_asym + T3_asym + curv_asym = 0
    # where _asym means the (b,c) antisymmetric part

    T2_asym_bc = T2 - jnp.transpose(T2, (0, 1, 3, 2))
    T3_asym_bc = T3 - jnp.transpose(T3, (0, 1, 3, 2))
    curv_asym_bc = (T4+T5) - jnp.transpose(T4+T5, (0, 1, 3, 2))
    total_asym = T1_asym_bc + T2_asym_bc + T3_asym_bc + curv_asym_bc

    print(f"    Antisymmetric parts in (b,c):")
    print(f"      |T1_asym|    = {jnp.max(jnp.abs(T1_asym_bc)):.2e}")
    print(f"      |T2_asym|    = {jnp.max(jnp.abs(T2_asym_bc)):.2e}")
    print(f"      |T3_asym|    = {jnp.max(jnp.abs(T3_asym_bc)):.2e}")
    print(f"      |curv_asym|  = {jnp.max(jnp.abs(curv_asym_bc)):.2e}")
    print(f"      |total_asym| = {jnp.max(jnp.abs(total_asym)):.2e}")

    # Compute third derivative using autodiff and compare
    print("\n  Autodiff verification:")
    autodiff_result = compute_third_derivative_autodiff(metric)
    print(f"    d3x_autodiff symmetry: (j,k)={autodiff_result['sym_jk']:.2e}, (k,l)={autodiff_result['sym_kl']:.2e}")

    # Compare autodiff to our formula
    d3_formula = d3_correct  # our corrected formula
    d3_auto = autodiff_result['d3x_autodiff']
    diff_formula_auto = jnp.max(jnp.abs(d3_formula - d3_auto))
    print(f"    |d3(formula) - d3(autodiff)| = {diff_formula_auto:.2e}")

  pdb.set_trace()


if __name__ == "__main__":
  main()

