import jax
import jax.numpy as jnp
import numpy as np
from jax import random
from jaxtyping import Array, PRNGKeyArray
from scipy import stats
from typing import Callable

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

from local_coordinates.basis import get_standard_basis
from local_coordinates.jet import function_to_jet
from local_coordinates.metric import RiemannianMetric
from local_coordinates.connection import get_levi_civita_connection, get_covariant_hessian
from local_coordinates.riemann import get_ricci_tensor

from local_geometry.rica_jax.hessian_inner_product import (
  estimate_metric_log_det_grad,
  estimate_metric_log_det_covariant_hessian_inner,
  christoffel_reduction,
  christoffel_term_1,
  christoffel_term_2,
  log_likelihood_covariant_hessian_inner,
  estimate_rnc_hessian_inner_product,
  make_dense_matrix_operations,
  make_metric_application_operators,
)

DIM = 5
N_SAMPLES = 2048


def make_random_metric_tensor_fn(
  key: PRNGKeyArray,
  dim: int,
) -> Callable[[Array], Array]:
  k1, k2 = random.split(key)
  base = random.normal(k1, (dim, dim))
  tangent = 0.2 * random.normal(k2, (dim, dim, dim))

  def g(x: Array) -> Array:
    matrix = base + jnp.tensordot(x, tangent, axes=1)
    return matrix @ matrix.T + jnp.eye(dim)

  return g


def sample_rademacher(
  key: PRNGKeyArray,
  shape: tuple[int, ...],
) -> Array:
  return 2.0*random.bernoulli(key, 0.5, shape).astype(jnp.float64) - 1.0


def make_test_fixtures(key: PRNGKeyArray):
  """Return (metric_fn, metric_ops, connection, x, w) for use across tests."""
  metric_key, x_key, w_key = random.split(key, 3)
  metric_fn = make_random_metric_tensor_fn(metric_key, DIM)
  metric_ops = make_dense_matrix_operations(metric_fn)
  x = random.normal(x_key, (DIM,))
  w_raw = random.normal(w_key, (DIM,))
  w = w_raw / jnp.linalg.norm(w_raw)

  basis = get_standard_basis(x)
  metric_jet = function_to_jet(metric_fn, x)
  metric = RiemannianMetric(basis=basis, components=metric_jet)
  connection = get_levi_civita_connection(metric)

  return metric_fn, metric_ops, connection, x, w


def run_hutchinson_ttest(
  exact_scalar: float,
  estimator_fn: Callable[[Array], Array],
  sample_key: PRNGKeyArray,
  alpha: float = 0.01,
):
  """Run a one-sample t-test comparing the mean of Hutchinson estimates to an exact value."""
  v_keys = random.split(sample_key, N_SAMPLES)

  def single_estimate(v_key: PRNGKeyArray) -> Array:
    v = sample_rademacher(v_key, (DIM,))
    return estimator_fn(v)

  estimates = jax.vmap(single_estimate)(v_keys)
  result = stats.ttest_1samp(np.asarray(estimates), float(exact_scalar))
  assert result.pvalue > alpha, (
    f"Hutchinson estimate mean {np.mean(estimates):.6f} != exact {exact_scalar:.6f} "
    f"(p={result.pvalue:.4f})"
  )


def test_estimate_metric_log_det_grad():
  dim = 5
  n_samples = 2048
  alpha = 0.01 / dim

  key = random.PRNGKey(0)
  metric_key, x_key, sample_key = random.split(key, 3)
  metric_fn = make_random_metric_tensor_fn(metric_key, dim)
  metric_ops = make_dense_matrix_operations(metric_fn)
  x = random.normal(x_key, (dim,))

  exact_grad = jax.grad(lambda x: jnp.linalg.slogdet(metric_fn(x))[1])(x)

  v_keys = random.split(sample_key, n_samples)

  def single_estimate(v_key: PRNGKeyArray) -> Array:
    v_hutch = sample_rademacher(v_key, (dim,))
    return estimate_metric_log_det_grad(metric_ops, x, v_hutch)

  estimates = jax.vmap(single_estimate)(v_keys)
  estimates_np = np.asarray(estimates)
  exact_grad_np = np.asarray(exact_grad)

  for idx in range(dim):
    test_result = stats.ttest_1samp(
      estimates_np[:, idx],
      float(exact_grad_np[idx]),
    )
    assert test_result.pvalue > alpha


def test_christoffel_reduction():
  """
  Ground truth: v_k Gamma^k_ij w^i w^j = einsum("k,ijk,i,j->", v, gamma, w, w).
  This is an exact (non-stochastic) test.
  """
  key = random.PRNGKey(1)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(key)

  v_key, _ = random.split(key)
  v = random.normal(v_key, (DIM,))

  # Gamma^k_ij = gamma[i, j, k]
  gamma = connection.christoffel_symbols.value
  exact = jnp.einsum("k,ijk,i,j->", v, gamma, w, w)

  result = christoffel_reduction(metric_ops, x, w, v)

  np.testing.assert_allclose(float(result), float(exact), rtol=1e-5)


def test_log_likelihood_covariant_hessian_inner():
  """
  Ground truth: nabla^2 log_p_ij w^i w^j via get_covariant_hessian from local_coordinates.
  This is an exact test.
  """
  key = random.PRNGKey(2)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(key)

  # Simple Gaussian log-likelihood
  def log_p(x: Array) -> Array:
    return -0.5 * jnp.dot(x, x)

  cov_hess = get_covariant_hessian(connection, log_p)
  # components.value has shape (D, D), indexed (i, j)
  exact = jnp.einsum("ij,i,j->", cov_hess.components.value, w, w)

  result = log_likelihood_covariant_hessian_inner(metric_ops, log_p, x, w)

  np.testing.assert_allclose(float(result), float(exact), rtol=1e-5)


def test_estimate_metric_log_det_covariant_hessian_inner():
  """
  Ground truth: nabla^2 log det g_ij w^i w^j via get_covariant_hessian.
  Stochastic: use Hutchinson t-test with two independent probe vectors.
  """
  key = random.PRNGKey(3)
  sample_key, fixture_key = random.split(key)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(fixture_key)

  def log_det_g(x: Array) -> Array:
    return jnp.linalg.slogdet(metric_fn(x))[1]

  cov_hess = get_covariant_hessian(connection, log_det_g)
  exact = float(jnp.einsum("ij,i,j->", cov_hess.components.value, w, w))

  v_keys_a = random.split(random.fold_in(sample_key, 0), N_SAMPLES)
  v_keys_b = random.split(random.fold_in(sample_key, 1), N_SAMPLES)

  def single_estimate(v_key_a: PRNGKeyArray, v_key_b: PRNGKeyArray) -> Array:
    v_hess = sample_rademacher(v_key_a, (DIM,))
    v_grad = sample_rademacher(v_key_b, (DIM,))
    return estimate_metric_log_det_covariant_hessian_inner(metric_ops, x, w, v_hess, v_grad)

  estimates = jax.vmap(single_estimate)(v_keys_a, v_keys_b)
  result = stats.ttest_1samp(np.asarray(estimates), exact)
  assert result.pvalue > 0.01, (
    f"Hutchinson estimate mean {np.mean(estimates):.6f} != exact {exact:.6f} "
    f"(p={result.pvalue:.4f})"
  )


def test_christoffel_term_1():
  """
  Ground truth: (d_a Gamma^a_ij) w^i w^j.

  Using the Jacobian of the Christoffel-symbol function:
    J[i,j,k,a] = d/dx^a Gamma^k_ij
  The contraction  d_a Gamma^a_ij w^i w^j = einsum("ijaa,i,j->", J, w, w)
  where the trace is over k=a (upper index equals derivative index).
  Stochastic: use Hutchinson t-test.
  """
  key = random.PRNGKey(4)
  sample_key, fixture_key = random.split(key)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(fixture_key)

  def gamma_fn(x: Array) -> Array:
    basis = get_standard_basis(x)
    metric_jet = function_to_jet(metric_fn, x)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    conn = get_levi_civita_connection(metric)
    return conn.christoffel_symbols.value  # shape (D, D, D), Gamma^k_ij = out[i,j,k]

  # J[i,j,k,a] = d/dx^a Gamma^k_ij
  J = jax.jacobian(gamma_fn)(x)  # shape (D, D, D, D)
  # Trace: d_a Gamma^a_ij = J[i,j,a,a] (upper index k equals derivative index a)
  exact = float(jnp.einsum("ijaa,i,j->", J, w, w))

  def estimator(v_hutch: Array) -> Array:
    return christoffel_term_1(metric_ops, x, w, v_hutch)

  run_hutchinson_ttest(exact, estimator, sample_key)


def test_christoffel_term_2():
  """
  Ground truth: Gamma^a_ib Gamma^b_aj w^i w^j = einsum("iab,baj,i,j->", gamma, gamma, w, w).
  Stochastic: use Hutchinson t-test.
  """
  key = random.PRNGKey(5)
  sample_key, fixture_key = random.split(key)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(fixture_key)

  # gamma[i,j,k] = Gamma^k_ij
  gamma = connection.christoffel_symbols.value
  # Gamma^a_ib Gamma^b_aj w^i w^j
  # = gamma[i,b,a] * gamma[a,j,b] * w^i * w^j
  exact = float(jnp.einsum("iba,ajb,i,j->", gamma, gamma, w, w))

  def estimator(v_hutch: Array) -> Array:
    return christoffel_term_2(metric_ops, x, w, v_hutch)

  run_hutchinson_ttest(exact, estimator, sample_key)


def test_rnc_hessian_via_ricci():
  """
  Ground truth from the top-level formula in Proposition 1:
    D(x) = nabla^2 log rho(x) - 1/3 Ric(x)
  where log rho = log p - 1/2 log det g is the Riemannian log-density.

  This is an independent check of estimate_rnc_hessian_inner_product using only
  get_covariant_hessian and get_ricci_tensor (no Christoffel Jacobians).
  Stochastic: use Hutchinson t-test with two independent probe vectors.
  """
  key = random.PRNGKey(7)
  sample_key, fixture_key = random.split(key)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(fixture_key)

  def log_p(x: Array) -> Array:
    return -0.5 * jnp.dot(x, x)

  def log_rho(x: Array) -> Array:
    return log_p(x) - 0.5 * jnp.linalg.slogdet(metric_fn(x))[1]

  cov_hess_log_rho = get_covariant_hessian(connection, log_rho)
  ricci = get_ricci_tensor(connection)
  exact = float(
    jnp.einsum("ij,i,j->", cov_hess_log_rho.components.value, w, w)
    - (1/3) * jnp.einsum("ij,i,j->", ricci.components.value, w, w)
  )

  v_keys_a = random.split(random.fold_in(sample_key, 0), N_SAMPLES)
  v_keys_b = random.split(random.fold_in(sample_key, 1), N_SAMPLES)

  def single_estimate(v_key_a: PRNGKeyArray, v_key_b: PRNGKeyArray) -> Array:
    v1 = sample_rademacher(v_key_a, (DIM,))
    v2 = sample_rademacher(v_key_b, (DIM,))
    return estimate_rnc_hessian_inner_product(metric_ops, log_p, x, w, v1, v2)

  estimates = jax.vmap(single_estimate)(v_keys_a, v_keys_b)
  result = stats.ttest_1samp(np.asarray(estimates), exact)
  assert result.pvalue > 0.01, (
    f"Hutchinson estimate mean {np.mean(estimates):.6f} != exact {exact:.6f} "
    f"(p={result.pvalue:.4f})"
  )


def make_pullback_fixtures(
  key: PRNGKeyArray,
  n: int = 4,
  m: int = 6,
  _lambda: float = 1.5,
):
  """Return (metric_ops, G_exact, x, w, u, v) for a linear encoder f(x) = A @ x.

  For f(x) = A @ x the Jacobian is constant J = A, so the exact metric matrix is
  G = lambda*I + A^T A, which can be formed explicitly for comparison.
  """
  k_A, k_x, k_w, k_u, k_v = random.split(key, 5)
  A = random.normal(k_A, (m, n))

  def f(x: Array) -> Array:
    return A @ x

  metric_ops = make_metric_application_operators(f, _lambda)
  G_exact = _lambda * jnp.eye(n) + A.T @ A

  x = random.normal(k_x, (n,))
  w_raw = random.normal(k_w, (n,))
  w = w_raw / jnp.linalg.norm(w_raw)
  u = random.normal(k_u, (n,))
  v = random.normal(k_v, (n,))

  return metric_ops, G_exact, x, w, u, v


def test_metric_application_g_inner():
  """g_inner(x, u, v) == u^T (lambda*I + A^T A) v for a linear encoder."""
  key = random.PRNGKey(10)
  metric_ops, G_exact, x, _, u, v = make_pullback_fixtures(key)

  result = metric_ops.g_inner(x, u, v)
  expected = jnp.vdot(u, G_exact @ v)

  np.testing.assert_allclose(float(result), float(expected), rtol=1e-5)


def test_metric_application_g_mvp():
  """g_mvp(x, v) == (lambda*I + A^T A) v for a linear encoder."""
  key = random.PRNGKey(11)
  metric_ops, G_exact, x, _, _, v = make_pullback_fixtures(key)

  result = metric_ops.g_mvp(x, v)
  expected = G_exact @ v

  np.testing.assert_allclose(np.asarray(result), np.asarray(expected), rtol=1e-5)


def test_metric_application_g_inv_mvp():
  """g_inv_mvp(x, v) == (lambda*I + A^T A)^{-1} v for a linear encoder.

  Tolerance is relaxed to 1e-4 to match the CG solver tolerance of rtol=1e-6,
  which maps to roughly that level of residual error in the solution.
  """
  key = random.PRNGKey(12)
  metric_ops, G_exact, x, _, _, v = make_pullback_fixtures(key)

  result = metric_ops.g_inv_mvp(x, v)
  expected = jnp.linalg.solve(G_exact, v)

  np.testing.assert_allclose(np.asarray(result), np.asarray(expected), rtol=1e-4)


def test_estimate_rnc_hessian_inner_product():
  """
  Ground truth: w^T D(x) w where
    D(x) = nabla^2 log p - 1/3 nabla^2 log det g - 1/3 d_a Gamma^a_ij + 1/3 Gamma^a_ib Gamma^b_aj.
  Stochastic: use Hutchinson t-test with two independent probe vectors.
  """
  key = random.PRNGKey(6)
  sample_key, fixture_key = random.split(key)
  metric_fn, metric_ops, connection, x, w = make_test_fixtures(fixture_key)

  def log_p(x: Array) -> Array:
    return -0.5 * jnp.dot(x, x)

  def log_det_g(x: Array) -> Array:
    return jnp.linalg.slogdet(metric_fn(x))[1]

  def gamma_fn(x: Array) -> Array:
    basis = get_standard_basis(x)
    metric_jet = function_to_jet(metric_fn, x)
    metric = RiemannianMetric(basis=basis, components=metric_jet)
    conn = get_levi_civita_connection(metric)
    return conn.christoffel_symbols.value

  gamma = connection.christoffel_symbols.value

  cov_hess_logp = get_covariant_hessian(connection, log_p)
  t1 = float(jnp.einsum("ij,i,j->", cov_hess_logp.components.value, w, w))

  cov_hess_logdet = get_covariant_hessian(connection, log_det_g)
  t2 = float(jnp.einsum("ij,i,j->", cov_hess_logdet.components.value, w, w))

  J = jax.jacobian(gamma_fn)(x)
  t3 = float(jnp.einsum("ijaa,i,j->", J, w, w))

  t4 = float(jnp.einsum("iba,ajb,i,j->", gamma, gamma, w, w))

  exact = t1 - (1/3)*t2 - (1/3)*t3 + (1/3)*t4

  v_keys_a = random.split(random.fold_in(sample_key, 0), N_SAMPLES)
  v_keys_b = random.split(random.fold_in(sample_key, 1), N_SAMPLES)

  def single_estimate(v_key_a: PRNGKeyArray, v_key_b: PRNGKeyArray) -> Array:
    v1 = sample_rademacher(v_key_a, (DIM,))
    v2 = sample_rademacher(v_key_b, (DIM,))
    return estimate_rnc_hessian_inner_product(metric_ops, log_p, x, w, v1, v2)

  estimates = jax.vmap(single_estimate)(v_keys_a, v_keys_b)
  result = stats.ttest_1samp(np.asarray(estimates), exact)
  assert result.pvalue > 0.01, (
    f"Hutchinson estimate mean {np.mean(estimates):.6f} != exact {exact:.6f} "
    f"(p={result.pvalue:.4f})"
  )
