"""Cross-framework tests comparing JAX and PyTorch Riemannian ICA implementations.

Each test runs the same computation in both frameworks, converts to numpy, and
asserts numerical agreement.
"""
import numpy as np
import pytest

import jax
import jax.numpy as jnp

import torch

from local_geometry.rica_jax.riemannian_ica import (
  TaylorData as JaxTaylorData,
  taylor_data as jax_taylor_data,
  christoffel as jax_christoffel,
  christoffel_gradient as jax_christoffel_gradient,
  riemann_tensor as jax_riemann_tensor,
  ricci_tensor as jax_ricci_tensor,
  log_rho_derivatives as jax_log_rho_derivatives,
  entanglement_tensor as jax_entanglement_tensor,
  orthonormal_frame as jax_orthonormal_frame,
  riemannian_ica as jax_riemannian_ica,
)
from local_geometry.rica_torch.riemannian_ica_torch import (
  TaylorData as TorchTaylorData,
  taylor_data as torch_taylor_data,
  christoffel as torch_christoffel,
  christoffel_gradient as torch_christoffel_gradient,
  riemann_tensor as torch_riemann_tensor,
  ricci_tensor as torch_ricci_tensor,
  log_rho_derivatives as torch_log_rho_derivatives,
  entanglement_tensor as torch_entanglement_tensor,
  orthonormal_frame as torch_orthonormal_frame,
  riemannian_ica as torch_riemannian_ica,
)


jax.config.update("jax_enable_x64", True)
torch.set_default_dtype(torch.float64)


# --- Shared test point ---

P_NP = np.array([0.17, -0.23])
P_JAX = jnp.array(P_NP)
P_TORCH = torch.tensor(P_NP)


# --- JAX test functions ---

def _jax_metric(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 _jax_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
  )


# --- PyTorch test functions ---

def _torch_metric(x):
  x0, x1 = x[0], x[1]
  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 torch.stack([torch.stack([g11, g12]), torch.stack([g12, g22])])


def _torch_log_px(x):
  x0, x1 = x[0], x[1]
  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
  )


# --- Helper to convert tensors to numpy ---

def _to_np(x):
  if isinstance(x, torch.Tensor):
    return x.detach().numpy()
  return np.asarray(x)


# =============================================================================
# Cross-framework comparison tests
# =============================================================================


def test_taylor_data_metric_matches():
  jax_td = jax_taylor_data(_jax_metric, P_JAX)
  torch_td = torch_taylor_data(_torch_metric, P_TORCH)

  np.testing.assert_allclose(_to_np(torch_td.value), _to_np(jax_td.value), atol=1e-14)
  np.testing.assert_allclose(_to_np(torch_td.gradient), _to_np(jax_td.gradient), atol=1e-14)
  np.testing.assert_allclose(_to_np(torch_td.hessian), _to_np(jax_td.hessian), atol=1e-14)


def test_taylor_data_scalar_matches():
  jax_td = jax_taylor_data(_jax_log_px, P_JAX)
  torch_td = torch_taylor_data(_torch_log_px, P_TORCH)

  np.testing.assert_allclose(_to_np(torch_td.value), _to_np(jax_td.value), atol=1e-14)
  np.testing.assert_allclose(_to_np(torch_td.gradient), _to_np(jax_td.gradient), atol=1e-14)
  np.testing.assert_allclose(_to_np(torch_td.hessian), _to_np(jax_td.hessian), atol=1e-14)


def test_christoffel_matches():
  jax_mt = jax_taylor_data(_jax_metric, P_JAX)
  torch_mt = torch_taylor_data(_torch_metric, P_TORCH)

  jax_g_inv = jnp.linalg.inv(jax_mt.value)
  torch_g_inv = torch.linalg.inv(torch_mt.value)

  jax_gamma = jax_christoffel(jax_g_inv, jax_mt.gradient)
  torch_gamma = torch_christoffel(torch_g_inv, torch_mt.gradient)

  np.testing.assert_allclose(_to_np(torch_gamma), _to_np(jax_gamma), atol=1e-12)


def test_christoffel_gradient_matches():
  jax_mt = jax_taylor_data(_jax_metric, P_JAX)
  torch_mt = torch_taylor_data(_torch_metric, P_TORCH)

  jax_g_inv = jnp.linalg.inv(jax_mt.value)
  torch_g_inv = torch.linalg.inv(torch_mt.value)

  jax_dg = jax_christoffel_gradient(jax_g_inv, jax_mt.gradient, jax_mt.hessian)
  torch_dg = torch_christoffel_gradient(torch_g_inv, torch_mt.gradient, torch_mt.hessian)

  np.testing.assert_allclose(_to_np(torch_dg), _to_np(jax_dg), atol=1e-12)


def test_ricci_matches():
  jax_mt = jax_taylor_data(_jax_metric, P_JAX)
  torch_mt = torch_taylor_data(_torch_metric, P_TORCH)

  jax_g_inv = jnp.linalg.inv(jax_mt.value)
  torch_g_inv = torch.linalg.inv(torch_mt.value)

  jax_gamma = jax_christoffel(jax_g_inv, jax_mt.gradient)
  jax_dgamma = jax_christoffel_gradient(jax_g_inv, jax_mt.gradient, jax_mt.hessian)
  jax_R = jax_riemann_tensor(jax_gamma, jax_dgamma)
  jax_ric = jax_ricci_tensor(jax_R)

  torch_gamma = torch_christoffel(torch_g_inv, torch_mt.gradient)
  torch_dgamma = torch_christoffel_gradient(torch_g_inv, torch_mt.gradient, torch_mt.hessian)
  torch_R = torch_riemann_tensor(torch_gamma, torch_dgamma)
  torch_ric = torch_ricci_tensor(torch_R)

  np.testing.assert_allclose(_to_np(torch_ric), _to_np(jax_ric), atol=1e-12)


def test_entanglement_tensor_matches():
  jax_mt = jax_taylor_data(_jax_metric, P_JAX)
  jax_lp = jax_taylor_data(_jax_log_px, P_JAX)
  jax_D = jax_entanglement_tensor(jax_mt, jax_lp)

  torch_mt = torch_taylor_data(_torch_metric, P_TORCH)
  torch_lp = torch_taylor_data(_torch_log_px, P_TORCH)
  torch_D = torch_entanglement_tensor(torch_mt, torch_lp)

  np.testing.assert_allclose(_to_np(torch_D), _to_np(jax_D), atol=1e-11)


def test_orthonormal_frame_matches():
  jax_g = _jax_metric(P_JAX)
  torch_g = _torch_metric(P_TORCH)

  jax_E = jax_orthonormal_frame(jax_g)
  torch_E = torch_orthonormal_frame(torch_g)

  np.testing.assert_allclose(_to_np(torch_E), _to_np(jax_E), atol=1e-12)


def test_riemannian_ica_eigenvalues_match():
  jax_result = jax_riemannian_ica(_jax_metric, _jax_log_px, P_JAX)
  torch_result = torch_riemannian_ica(_torch_metric, _torch_log_px, P_TORCH)

  np.testing.assert_allclose(
    _to_np(torch_result.eigenvalues),
    _to_np(jax_result.eigenvalues),
    atol=1e-10,
  )


def test_riemannian_ica_directions_match():
  """Directions should match up to per-column sign flips."""
  jax_result = jax_riemannian_ica(_jax_metric, _jax_log_px, P_JAX)
  torch_result = torch_riemannian_ica(_torch_metric, _torch_log_px, P_TORCH)

  jax_dirs = _to_np(jax_result.directions)
  torch_dirs = _to_np(torch_result.directions)

  d = jax_dirs.shape[0]
  for i in range(d):
    dot = np.dot(jax_dirs[:, i], torch_dirs[:, i])
    norm_product = np.linalg.norm(jax_dirs[:, i]) * np.linalg.norm(torch_dirs[:, i])
    np.testing.assert_allclose(np.abs(dot), norm_product, atol=1e-10)


# =============================================================================
# PyTorch-only correctness tests
# =============================================================================


def test_flat_metric_reduces_to_hessian():
  """When the metric is Euclidean, D reduces to the Hessian of log p."""
  def flat_metric(x):
    d = x.shape[0]
    return torch.eye(d, dtype=x.dtype)

  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 = torch.tensor([0.3, -0.4])
  mt = torch_taylor_data(flat_metric, p)
  lp = torch_taylor_data(log_p_fn, p)
  D = torch_entanglement_tensor(mt, lp)

  expected = torch.autograd.functional.hessian(log_p_fn, p)
  np.testing.assert_allclose(_to_np(D), _to_np(expected), atol=1e-12)


def test_disentangled_hessian_diagonal():
  """The found directions should produce a diagonal log-likelihood Hessian in RNC."""
  mt = torch_taylor_data(_torch_metric, P_TORCH)
  g_inv = torch.linalg.inv(mt.value)
  gamma = torch_christoffel(g_inv, mt.gradient)
  dgamma = torch_christoffel_gradient(g_inv, mt.gradient, mt.hessian)

  result = torch_riemannian_ica(_torch_metric, _torch_log_px, P_TORCH)
  W = result.directions
  z0 = torch.zeros_like(P_TORCH)

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

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

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

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

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

  H = torch.autograd.functional.hessian(log_ps, z0)
  off_diag = H - torch.diag(torch.diag(H))
  np.testing.assert_allclose(_to_np(off_diag), 0.0, atol=1e-8)


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

  def jax_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 jax_log_p_fn(x):
    return -0.5 * jnp.sum(x**2) + 0.03 * jnp.sum(x**3)

  def torch_metric_fn(x):
    diag = 1.0 + 0.1 * x**2
    g = torch.diag(diag)
    for i in range(d - 1):
      off = 0.02 * x[i] * x[i + 1]
      g = g.clone()
      g[i, i + 1] = g[i, i + 1] + off
      g[i + 1, i] = g[i + 1, i] + off
    return g

  def torch_log_p_fn(x):
    return -0.5 * torch.sum(x**2) + 0.03 * torch.sum(x**3)

  p_np = 0.1 * np.arange(d, dtype=np.float64)
  p_jax = jnp.array(p_np)
  p_torch = torch.tensor(p_np)

  jax_mt = jax_taylor_data(jax_metric_fn, p_jax)
  jax_lp = jax_taylor_data(jax_log_p_fn, p_jax)
  jax_D = jax_entanglement_tensor(jax_mt, jax_lp)

  torch_mt = torch_taylor_data(torch_metric_fn, p_torch)
  torch_lp = torch_taylor_data(torch_log_p_fn, p_torch)
  torch_D = torch_entanglement_tensor(torch_mt, torch_lp)

  np.testing.assert_allclose(_to_np(torch_D), _to_np(jax_D), atol=1e-10)

  jax_result = jax_riemannian_ica(jax_metric_fn, jax_log_p_fn, p_jax)
  torch_result = torch_riemannian_ica(torch_metric_fn, torch_log_p_fn, p_torch)

  np.testing.assert_allclose(
    _to_np(torch_result.eigenvalues),
    _to_np(jax_result.eigenvalues),
    atol=1e-10,
  )

  W = torch_result.directions
  g = torch_metric_fn(p_torch)
  np.testing.assert_allclose(_to_np(W.T @ g @ W), np.eye(d), atol=1e-12)
