import jax
import jax.numpy as jnp

from local_coordinates.basis import get_standard_basis
from local_coordinates.connection import get_levi_civita_connection
from local_coordinates.jet import function_to_jet, jet_decorator
from local_coordinates.metric import RiemannianMetric
from local_coordinates.normal_coords import get_rnc_jacobians
from local_coordinates.riemann import get_ricci_tensor
from local_geometry.independent_components import get_entanglement_tensor, riemannian_ica


jax.config.update("jax_enable_x64", True)


def _metric_components(x: jax.Array) -> jax.Array:
  x0, x1 = x
  g11 = 1.5 + 0.2 * x0**2 + 0.1 * x1**2
  g22 = 1.2 + 0.15 * x0**2 + 0.25 * x1**2
  g12 = 0.05 * x0 * x1
  return jnp.array([[g11, g12], [g12, g22]])


def _log_px(x: jax.Array) -> jax.Array:
  x0, x1 = x
  return (
    -0.3 * x0**2
    - 0.7 * x1**2
    - 0.2 * x0 * x1
    + 0.05 * x0**3
    - 0.04 * x1**3
    + 0.03 * x0**2 * x1
  )


def _build_metric(x0: jax.Array) -> RiemannianMetric:
  basis = get_standard_basis(x0)
  metric_jet = function_to_jet(_metric_components, x0)
  return RiemannianMetric(basis=basis, components=metric_jet)


def _x_of_z(z: jax.Array, x0: jax.Array, j_z_to_x) -> jax.Array:
  linear = jnp.einsum("ia,a->i", j_z_to_x.value, z)
  quadratic = 0.5 * jnp.einsum("iab,a,b->i", j_z_to_x.gradient, z, z)
  cubic = (1.0 / 6.0) * jnp.einsum("iabc,a,b,c->i", j_z_to_x.hessian, z, z, z)
  return x0 + linear + quadratic + cubic


def _dx_dz_of_z(z: jax.Array, j_z_to_x) -> jax.Array:
  linear = jnp.einsum("iab,b->ia", j_z_to_x.gradient, z)
  quadratic = 0.5 * jnp.einsum("iabc,b,c->ia", j_z_to_x.hessian, z, z)
  return j_z_to_x.value + linear + quadratic


def _covariant_hessian(scalar_jet, gamma_value: jax.Array) -> jax.Array:
  correction = jnp.einsum("abc,c->ab", gamma_value, scalar_jet.gradient)
  return scalar_jet.hessian - correction


def test_log_pz_hessian_matches_riemannian_ica_formula() -> None:
  x0 = jnp.array([0.17, -0.23])
  z0 = jnp.zeros_like(x0)

  metric = _build_metric(x0)
  connection = get_levi_civita_connection(metric)
  ricci = get_ricci_tensor(connection).components.value
  _, j_z_to_x = get_rnc_jacobians(metric)

  def log_pz_from_change_of_variables(z: jax.Array) -> jax.Array:
    x = _x_of_z(z, x0, j_z_to_x)
    dx_dz = _dx_dz_of_z(z, j_z_to_x)
    log_abs_det_jacobian = jnp.linalg.slogdet(dx_dz)[1]
    return _log_px(x) + log_abs_det_jacobian

  hessian_direct = jax.hessian(log_pz_from_change_of_variables)(z0)

  log_px_jet = function_to_jet(_log_px, x0)

  @jet_decorator
  def log_det_metric(gij: jax.Array) -> jax.Array:
    return jnp.linalg.slogdet(gij)[1]

  log_det_g_jet = log_det_metric(metric.components.get_value_jet())

  gamma_value = connection.christoffel_symbols.value
  cov_hess_log_px = _covariant_hessian(log_px_jet, gamma_value)
  cov_hess_log_det_g = _covariant_hessian(log_det_g_jet, gamma_value)
  cov_hess_log_rho = cov_hess_log_px - 0.5 * cov_hess_log_det_g

  core_term = cov_hess_log_rho - (1.0 / 3.0) * ricci
  hessian_formula = jnp.einsum("ab,ai,bj->ij", core_term, j_z_to_x.value, j_z_to_x.value)

  assert jnp.allclose(hessian_direct, hessian_formula, atol=1e-8)


def test_entanglement_tensor_matches_manual_computation() -> None:
  """get_entanglement_tensor should match the manual cov-Hessian-minus-Ricci computation."""
  x0 = jnp.array([0.17, -0.23])
  metric = _build_metric(x0)

  D = get_entanglement_tensor(metric, _log_px)
  D_value = D.components.value

  connection = get_levi_civita_connection(metric)
  ricci = get_ricci_tensor(connection).components.value

  log_px_jet = function_to_jet(_log_px, x0)

  @jet_decorator
  def log_det_metric(gij: jax.Array) -> jax.Array:
    return jnp.linalg.slogdet(gij)[1]

  log_det_g_jet = log_det_metric(metric.components.get_value_jet())

  gamma_value = connection.christoffel_symbols.value
  cov_hess_log_px = _covariant_hessian(log_px_jet, gamma_value)
  cov_hess_log_det_g = _covariant_hessian(log_det_g_jet, gamma_value)
  cov_hess_log_rho = cov_hess_log_px - 0.5 * cov_hess_log_det_g

  expected = cov_hess_log_rho - (1.0 / 3.0) * ricci
  assert jnp.allclose(D_value, expected, atol=1e-12)


def test_entanglement_tensor_is_symmetric() -> None:
  """The entanglement tensor must be symmetric."""
  x0 = jnp.array([0.17, -0.23])
  metric = _build_metric(x0)
  D = get_entanglement_tensor(metric, _log_px)
  D_value = D.components.value
  assert jnp.allclose(D_value, D_value.T, atol=1e-12)


def test_riemannian_ica_diagonalizes_hessian() -> None:
  """The Jacobian from riemannian_ica should produce a diagonal log-likelihood Hessian."""
  x0 = jnp.array([0.17, -0.23])
  z0 = jnp.zeros_like(x0)
  metric = _build_metric(x0)

  j_x_to_s = riemannian_ica(metric, _log_px)
  j_s_to_x = j_x_to_s.get_inverse()

  def log_ps_from_change_of_variables(s: jax.Array) -> jax.Array:
    x = _x_of_z(s, x0, j_s_to_x)
    dx_ds = _dx_dz_of_z(s, j_s_to_x)
    log_abs_det_jacobian = jnp.linalg.slogdet(dx_ds)[1]
    return _log_px(x) + log_abs_det_jacobian

  hessian = jax.hessian(log_ps_from_change_of_variables)(z0)

  off_diag = hessian - jnp.diag(jnp.diag(hessian))
  assert jnp.allclose(off_diag, 0.0, atol=1e-8)


def test_riemannian_ica_eigenvalues_match_diagonal() -> None:
  """Diagonal entries of the disentangled Hessian should be eigenvalues of D."""
  x0 = jnp.array([0.17, -0.23])
  z0 = jnp.zeros_like(x0)
  metric = _build_metric(x0)

  D = get_entanglement_tensor(metric, _log_px)
  _, j_z_to_x = get_rnc_jacobians(metric)

  D_rnc_value = jnp.einsum(
    "ab,ai,bj->ij", D.components.value, j_z_to_x.value, j_z_to_x.value
  )
  eigenvalues_expected = jnp.sort(jnp.linalg.eigh(D_rnc_value)[0])

  j_x_to_s = riemannian_ica(metric, _log_px)
  j_s_to_x = j_x_to_s.get_inverse()

  def log_ps(s: jax.Array) -> jax.Array:
    x = _x_of_z(s, x0, j_s_to_x)
    dx_ds = _dx_dz_of_z(s, j_s_to_x)
    return _log_px(x) + jnp.linalg.slogdet(dx_ds)[1]

  hessian_diag = jnp.sort(jnp.diag(jax.hessian(log_ps)(z0)))
  assert jnp.allclose(hessian_diag, eigenvalues_expected, atol=1e-8)
