"""Tests for the dHSIC independence test."""

import jax
import jax.numpy as jnp
import pytest

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

from local_geometry.independence_test_paper.dhsic import (
  dhsic,
  dhsic_weighted,
  dhsic_permutation_test,
  gaussian_kernel_matrix,
  median_heuristic,
)


def test_median_heuristic_positive():
  """Bandwidth should be strictly positive for non-degenerate data."""
  key = jax.random.PRNGKey(0)
  x = jax.random.normal(key, (50,))
  sigma = median_heuristic(x)
  assert sigma > 0.0


def test_gaussian_kernel_matrix_shape_and_symmetry():
  """Gram matrix should be n x n and symmetric."""
  key = jax.random.PRNGKey(1)
  x = jax.random.normal(key, (20,))
  sigma = median_heuristic(x)
  K = gaussian_kernel_matrix(x, sigma)
  assert K.shape == (20, 20)
  assert jnp.allclose(K, K.T, atol=1e-12)


def test_gaussian_kernel_matrix_diagonal_ones():
  """Diagonal of the Gaussian kernel Gram matrix should be 1."""
  key = jax.random.PRNGKey(2)
  x = jax.random.normal(key, (15,))
  sigma = median_heuristic(x)
  K = gaussian_kernel_matrix(x, sigma)
  assert jnp.allclose(jnp.diag(K), 1.0, atol=1e-12)


def test_gaussian_kernel_matrix_values_in_unit_interval():
  """All entries should be in [0, 1]."""
  key = jax.random.PRNGKey(3)
  x = jax.random.normal(key, (30,))
  sigma = median_heuristic(x)
  K = gaussian_kernel_matrix(x, sigma)
  assert jnp.all(K >= 0.0)
  assert jnp.all(K <= 1.0 + 1e-12)


def test_dhsic_nonnegative():
  """dHSIC should always be non-negative."""
  key = jax.random.PRNGKey(10)
  for _ in range(5):
    key, subkey = jax.random.split(key)
    X = jax.random.normal(subkey, (50, 3))
    assert dhsic(X) >= -1e-10


def test_dhsic_known_value_d2():
  """Verify dHSIC against a hand-computed value for a tiny d=2 dataset.

  With n=3, d=2, and data X = [[0, 1], [1, 0], [2, 1]], we manually
  trace through Algorithm 1 to get a reference value.
  """
  X = jnp.array([[0.0, 1.0], [1.0, 0.0], [2.0, 1.0]])
  n = 3

  x1 = X[:, 0]
  x2 = X[:, 1]
  sigma1 = median_heuristic(x1)
  sigma2 = median_heuristic(x2)
  K1 = gaussian_kernel_matrix(x1, sigma1)
  K2 = gaussian_kernel_matrix(x2, sigma2)

  term1 = K1 * K2
  t1 = jnp.sum(term1) / n ** 2

  sum_K1 = jnp.sum(K1)
  sum_K2 = jnp.sum(K2)
  t2 = (1.0 / n ** 2) * sum_K1 * (1.0 / n ** 2) * sum_K2

  col_sum_K1 = jnp.sum(K1, axis=0)
  col_sum_K2 = jnp.sum(K2, axis=0)
  t3_vec = (2.0 / n) * (1.0 / n) * col_sum_K1 * (1.0 / n) * col_sum_K2
  t3 = jnp.sum(t3_vec)

  expected = t1 + t2 - t3

  actual = dhsic(X)
  assert jnp.allclose(actual, expected, atol=1e-12), (
    f"Expected {expected}, got {actual}"
  )


def test_dhsic_independent_gaussian_near_zero():
  """For independent standard normal components, dHSIC should be small."""
  key = jax.random.PRNGKey(42)
  n, d = 500, 3
  X = jax.random.normal(key, (n, d))
  stat = dhsic(X)
  assert stat < 0.01, f"dHSIC for independent data unexpectedly large: {stat}"


def test_dhsic_dependent_components_large():
  """For strongly dependent components, dHSIC should be large."""
  key = jax.random.PRNGKey(99)
  n = 500
  x1 = jax.random.normal(key, (n,))
  x2 = x1 + 0.01 * jax.random.normal(jax.random.PRNGKey(100), (n,))
  X = jnp.column_stack([x1, x2])
  stat = dhsic(X)
  assert stat > 0.01, f"dHSIC for dependent data unexpectedly small: {stat}"


def test_dhsic_increases_with_dependence():
  """dHSIC should be larger for more dependent data."""
  key = jax.random.PRNGKey(77)
  n = 300
  x1 = jax.random.normal(key, (n,))
  noise = jax.random.normal(jax.random.PRNGKey(78), (n,))

  X_weak = jnp.column_stack([x1, x1 + 2.0 * noise])
  X_strong = jnp.column_stack([x1, x1 + 0.1 * noise])

  stat_weak = dhsic(X_weak)
  stat_strong = dhsic(X_strong)
  assert stat_strong > stat_weak, (
    f"Strong dependence ({stat_strong}) should exceed weak ({stat_weak})"
  )


@pytest.mark.slow
def test_permutation_test_does_not_reject_independent():
  """Permutation test should not reject H0 for independent data."""
  key = jax.random.PRNGKey(55)
  n, d = 200, 3
  X = jax.random.normal(key, (n, d))
  result = dhsic_permutation_test(X, alpha=0.05, num_permutations=200, seed=55)
  assert not result.rejected, (
    f"Incorrectly rejected independence (p={result.p_value})"
  )


@pytest.mark.slow
def test_permutation_test_rejects_dependent():
  """Permutation test should reject H0 for dependent data."""
  key = jax.random.PRNGKey(66)
  n = 200
  x1 = jax.random.normal(key, (n,))
  x2 = jnp.sin(x1) + 0.05 * jax.random.normal(jax.random.PRNGKey(67), (n,))
  X = jnp.column_stack([x1, x2])
  result = dhsic_permutation_test(X, alpha=0.05, num_permutations=200, seed=66)
  assert result.rejected, (
    f"Failed to reject dependence (p={result.p_value})"
  )


class TestDhsicWeighted:

  def test_uniform_weights_matches_unweighted(self):
    """With uniform log-weights, dhsic_weighted should match dhsic."""
    key = jax.random.PRNGKey(200)
    X = jax.random.normal(key, (50, 3))
    log_w = jnp.zeros(50)
    stat_unweighted = dhsic(X)
    stat_weighted = dhsic_weighted(X, log_w)
    assert jnp.allclose(stat_unweighted, stat_weighted, atol=1e-12), (
      f"Unweighted {stat_unweighted} != weighted {stat_weighted}"
    )

  def test_nonnegative(self):
    """Weighted dHSIC should be non-negative for various weight patterns."""
    key = jax.random.PRNGKey(210)
    for i in range(5):
      k1, k2, key = jax.random.split(key, 3)
      X = jax.random.normal(k1, (60, 3))
      log_w = jax.random.normal(k2, (60,))
      stat = dhsic_weighted(X, log_w)
      assert stat >= -1e-10, f"Trial {i}: weighted dHSIC negative: {stat}"

  def test_independent_near_zero(self):
    """Independent components with uniform weights should give small dHSIC."""
    key = jax.random.PRNGKey(220)
    n, d = 500, 3
    X = jax.random.normal(key, (n, d))
    log_w = jnp.zeros(n)
    stat = dhsic_weighted(X, log_w)
    assert stat < 0.01, (
      f"Weighted dHSIC for independent data unexpectedly large: {stat}"
    )

  def test_dependent_large(self):
    """Strongly dependent components should give large weighted dHSIC."""
    key = jax.random.PRNGKey(230)
    n = 500
    x1 = jax.random.normal(key, (n,))
    x2 = x1 + 0.01 * jax.random.normal(jax.random.PRNGKey(231), (n,))
    X = jnp.column_stack([x1, x2])
    log_w = jnp.zeros(n)
    stat = dhsic_weighted(X, log_w)
    assert stat > 0.01, (
      f"Weighted dHSIC for dependent data unexpectedly small: {stat}"
    )


def test_permutation_test_result_fields():
  """PermutationTestResult should have the expected fields."""
  X = jax.random.normal(jax.random.PRNGKey(0), (30, 2))
  result = dhsic_permutation_test(X, num_permutations=10, seed=0)
  assert hasattr(result, "statistic")
  assert hasattr(result, "p_value")
  assert hasattr(result, "rejected")
  assert 0.0 <= result.p_value <= 1.0
