"""Tests for the standalone Riemannian ICA implementation.

Each test compares the standalone functions in local_geometry.riemannian_ica
against the existing library (local_coordinates) to verify correctness.
"""
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 as riemannian_ica_library,
)
from local_geometry.rica_jax.riemannian_ica import (
  TaylorData,
  taylor_data,
  christoffel,
  christoffel_gradient,
  riemann_tensor,
  ricci_tensor,
  log_rho_derivatives,
  entanglement_tensor,
  orthonormal_frame,
  riemannian_ica,
  riemannian_ica_from_taylor,
)


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


def _metric_components(x):
  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):
  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(p):
  basis = get_standard_basis(p)
  metric_jet = function_to_jet(_metric_components, p)
  return RiemannianMetric(basis=basis, components=metric_jet)


P = jnp.array([0.17, -0.23])


def test_christoffel_matches_library():
  metric = _build_metric(P)
  connection = get_levi_civita_connection(metric)
  expected = connection.christoffel_symbols.value

  g = _metric_components(P)
  g_inv = jnp.linalg.inv(g)
  dg = jax.jacfwd(_metric_components)(P)
  actual = christoffel(g_inv, dg)

  assert jnp.allclose(actual, expected, atol=1e-12)


def test_christoffel_gradient_matches_library():
  metric = _build_metric(P)
  connection = get_levi_civita_connection(metric)
  expected = connection.christoffel_symbols.gradient

  g = _metric_components(P)
  g_inv = jnp.linalg.inv(g)
  dg = jax.jacfwd(_metric_components)(P)
  d2g = jax.jacfwd(jax.jacfwd(_metric_components))(P)
  actual = christoffel_gradient(g_inv, dg, d2g)

  assert jnp.allclose(actual, expected, atol=1e-12)


def test_ricci_matches_library():
  metric = _build_metric(P)
  connection = get_levi_civita_connection(metric)
  expected = get_ricci_tensor(connection).components.value

  g = _metric_components(P)
  g_inv = jnp.linalg.inv(g)
  dg = jax.jacfwd(_metric_components)(P)
  d2g = jax.jacfwd(jax.jacfwd(_metric_components))(P)
  gamma = christoffel(g_inv, dg)
  dgamma = christoffel_gradient(g_inv, dg, d2g)
  R = riemann_tensor(gamma, dgamma)
  actual = ricci_tensor(R)

  assert jnp.allclose(actual, expected, atol=1e-12)


def test_entanglement_tensor_matches_library():
  metric = _build_metric(P)
  expected = get_entanglement_tensor(metric, _log_px).components.value

  mt = taylor_data(_metric_components, P)
  lp = taylor_data(_log_px, P)
  actual = entanglement_tensor(mt, lp)

  assert jnp.allclose(actual, expected, atol=1e-11)


def test_entanglement_tensor_is_symmetric():
  mt = taylor_data(_metric_components, P)
  lp = taylor_data(_log_px, P)
  D = entanglement_tensor(mt, lp)
  assert jnp.allclose(D, D.T, atol=1e-12)


def test_eigenvalues_match_library():
  metric = _build_metric(P)
  j_x_to_s_lib = riemannian_ica_library(metric, _log_px)
  j_s_to_x_lib = j_x_to_s_lib.get_inverse()

  z0 = jnp.zeros_like(P)

  def _x_of_s(s, j_s_to_x):
    linear = jnp.einsum("ia,a->i", j_s_to_x.value, s)
    quadratic = 0.5 * jnp.einsum("iab,a,b->i", j_s_to_x.gradient, s, s)
    cubic = (1.0 / 6.0) * jnp.einsum("iabc,a,b,c->i", j_s_to_x.hessian, s, s, s)
    return P + linear + quadratic + cubic

  def _dx_ds(s, j_s_to_x):
    linear = jnp.einsum("iab,b->ia", j_s_to_x.gradient, s)
    quadratic = 0.5 * jnp.einsum("iabc,b,c->ia", j_s_to_x.hessian, s, s)
    return j_s_to_x.value + linear + quadratic

  def log_ps(s):
    x = _x_of_s(s, j_s_to_x_lib)
    dx_ds = _dx_ds(s, j_s_to_x_lib)
    return _log_px(x) + jnp.linalg.slogdet(dx_ds)[1]

  lib_eigenvalues = jnp.sort(jnp.diag(jax.hessian(log_ps)(z0)))

  result = riemannian_ica(_metric_components, _log_px, P)
  standalone_eigenvalues = jnp.sort(result.eigenvalues)

  assert jnp.allclose(standalone_eigenvalues, lib_eigenvalues, atol=1e-8)


def test_directions_match_library():
  """Standalone directions should match the library frame (up to per-column signs)."""
  metric = _build_metric(P)
  j_x_to_s_lib = riemannian_ica_library(metric, _log_px)
  j_s_to_x_lib = j_x_to_s_lib.get_inverse()
  lib_frame = j_s_to_x_lib.value

  result = riemannian_ica(_metric_components, _log_px, P)
  standalone_frame = result.directions

  for i in range(P.shape[0]):
    dot = jnp.dot(lib_frame[:, i], standalone_frame[:, i])
    assert jnp.allclose(jnp.abs(dot), jnp.linalg.norm(lib_frame[:, i]) * jnp.linalg.norm(standalone_frame[:, i]), atol=1e-10), (
      f"Direction {i} mismatch: dot={dot}"
    )


def test_disentangled_hessian_is_diagonal():
  """The standalone directions should produce a diagonal log-likelihood Hessian.

  Builds the full RNC forward Jacobian (dx/ds) from the standalone Christoffel
  symbols and directions, then verifies diagonality via autodiff.
  """
  mt = taylor_data(_metric_components, P)
  g_inv = jnp.linalg.inv(mt.value)
  gamma = christoffel(g_inv, mt.gradient)
  dgamma = christoffel_gradient(g_inv, mt.gradient, mt.hessian)

  result = riemannian_ica(_metric_components, _log_px, P)
  W = result.directions
  z0 = jnp.zeros_like(P)

  d2xds2 = -jnp.einsum("abi,aj,bk->ijk", gamma, W, W)

  unsym_t1 = -jnp.einsum("pqim,ma,pb,qc->iabc", dgamma, W, W, W)
  unsym_t2 = 2 * jnp.einsum("pqi,mnp,ma,nb,qc->iabc", gamma, gamma, W, W, W)
  unsym = unsym_t1 + unsym_t2
  d3xds3 = (
    unsym
    + jnp.transpose(unsym, (0, 1, 3, 2))
    + jnp.transpose(unsym, (0, 2, 1, 3))
    + jnp.transpose(unsym, (0, 2, 3, 1))
    + jnp.transpose(unsym, (0, 3, 1, 2))
    + jnp.transpose(unsym, (0, 3, 2, 1))
  ) / 6

  def x_of_s(s):
    return (
      P
      + jnp.einsum("ia,a->i", W, s)
      + 0.5 * jnp.einsum("iab,a,b->i", d2xds2, s, s)
      + (1.0 / 6.0) * jnp.einsum("iabc,a,b,c->i", d3xds3, s, s, s)
    )

  def dx_ds(s):
    return (
      W
      + jnp.einsum("iab,b->ia", d2xds2, s)
      + 0.5 * jnp.einsum("iabc,b,c->ia", d3xds3, s, s)
    )

  def log_ps(s):
    x = x_of_s(s)
    J = dx_ds(s)
    return _log_px(x) + jnp.linalg.slogdet(J)[1]

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


def test_higher_dim():
  """Test with d=4 to verify correctness beyond 2D."""
  d = 4

  def metric_fn(x):
    g = jnp.eye(d)
    for i in range(d):
      g = g.at[i, i].add(0.1 * x[i] ** 2)
    for i in range(d - 1):
      off = 0.02 * x[i] * x[i + 1]
      g = g.at[i, i + 1].add(off)
      g = g.at[i + 1, i].add(off)
    return g

  def log_p_fn(x):
    return -0.5 * jnp.sum(x ** 2) + 0.03 * jnp.sum(x ** 3)

  p = 0.1 * jnp.arange(d, dtype=jnp.float64)

  metric = RiemannianMetric(
    basis=get_standard_basis(p),
    components=function_to_jet(metric_fn, p),
  )
  D_lib = get_entanglement_tensor(metric, log_p_fn).components.value

  mt = taylor_data(metric_fn, p)
  lp = taylor_data(log_p_fn, p)
  D_standalone = entanglement_tensor(mt, lp)

  assert jnp.allclose(D_standalone, D_lib, atol=1e-10)

  result = riemannian_ica(metric_fn, log_p_fn, p)
  assert result.directions.shape == (d, d)
  assert result.eigenvalues.shape == (d,)

  E = orthonormal_frame(metric_fn(p))
  assert jnp.allclose(E.T @ metric_fn(p) @ E, jnp.eye(d), atol=1e-12)

  W = result.directions
  assert jnp.allclose(W.T @ metric_fn(p) @ W, jnp.eye(d), atol=1e-12)


def test_flat_metric_reduces_to_ordinary_hessian():
  """When the metric is Euclidean, D reduces to the Hessian of log p."""
  def flat_metric(x):
    return jnp.eye(x.shape[0])

  def log_p_fn(x):
    return -0.5 * (2.0 * x[0] ** 2 + 3.0 * x[1] ** 2 + 1.0 * x[0] * x[1])

  p = jnp.array([0.3, -0.4])
  mt = taylor_data(flat_metric, p)
  lp = taylor_data(log_p_fn, p)
  D = entanglement_tensor(mt, lp)

  expected = jax.hessian(log_p_fn)(p)
  assert jnp.allclose(D, expected, atol=1e-12)


def test_taylor_data_helper():
  """taylor_data should match manual derivatives."""
  td = taylor_data(_metric_components, P)
  assert jnp.allclose(td.value, _metric_components(P), atol=1e-14)

  manual_grad = jax.jacfwd(_metric_components)(P)
  assert jnp.allclose(td.gradient, manual_grad, atol=1e-14)

  td_scalar = taylor_data(_log_px, P)
  assert jnp.allclose(td_scalar.value, _log_px(P), atol=1e-14)
  assert jnp.allclose(td_scalar.gradient, jax.grad(_log_px)(P), atol=1e-14)
  assert jnp.allclose(td_scalar.hessian, jax.hessian(_log_px)(P), atol=1e-14)
