"""Tests for the dHSIC neighborhood independence experiment."""
import jax
import jax.numpy as jnp
import pytest

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

from local_geometry.paper.figures.dhsic_independence import (
  sir_resample,
  compute_log_weights,
  direct_dhsic,
  _single_trial_from_key,
  run_single_trial,
  run_multi_trial,
)
from local_geometry.paper.figures.neighborhood_independence import (
  build_metric,
  get_disentangled_frame_rotation,
  make_random_rotation,
  sample_uniform_region,
)
from local_geometry.paper.figures.analytic_verification import (
  make_decoder,
  make_coord_fns,
)
from local_geometry.independence_test_paper.dhsic import dhsic
from local_coordinates.normal_coords import get_rnc_jacobians
from local_coordinates.jacobian import Jacobian


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def small_decoder():
  """A d=2, n=5 decoder for fast testing."""
  key = jax.random.PRNGKey(42)
  f, _, _, _ = make_decoder(key, d=2, n=5, eps=0.8)
  z0 = jnp.zeros(2)
  return f, z0


@pytest.fixture
def identity_jacobian():
  """An identity Jacobian with zero higher-order terms."""
  d = 2
  return Jacobian(
    value=jnp.eye(d),
    gradient=jnp.zeros((d, d, d)),
    hessian=jnp.zeros((d, d, d, d)),
  )


# ---------------------------------------------------------------------------
# Test: sir_resample
# ---------------------------------------------------------------------------

class TestSirResample:
  def test_output_shape(self):
    key = jax.random.PRNGKey(0)
    s = jax.random.normal(key, (1000, 3))
    log_w = jnp.zeros(1000)
    resampled = sir_resample(jax.random.PRNGKey(1), s, log_w, 200)
    assert resampled.shape == (200, 3)

  def test_uniform_weights_preserve_distribution(self):
    """With uniform weights, the resampled marginal means should match."""
    key = jax.random.PRNGKey(0)
    s = jax.random.normal(key, (10_000, 2))
    log_w = jnp.zeros(10_000)
    resampled = sir_resample(jax.random.PRNGKey(1), s, log_w, 5000)
    original_mean = jnp.mean(s, axis=0)
    resampled_mean = jnp.mean(resampled, axis=0)
    assert jnp.allclose(original_mean, resampled_mean, atol=0.1)

  def test_concentrated_weights_select_specific_points(self):
    """When one weight dominates, resampled points should cluster there."""
    s = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0], [3.0, 3.0]])
    log_w = jnp.array([-1000.0, -1000.0, 0.0, -1000.0])
    resampled = sir_resample(jax.random.PRNGKey(0), s, log_w, 100)
    assert jnp.allclose(resampled, jnp.array([2.0, 2.0]), atol=1e-10)

  def test_jit_compatible(self):
    key = jax.random.PRNGKey(0)
    s = jax.random.normal(key, (500, 2))
    log_w = jnp.zeros(500)
    jitted = jax.jit(sir_resample, static_argnums=(3,))
    resampled = jitted(jax.random.PRNGKey(1), s, log_w, 100)
    assert resampled.shape == (100, 2)


# ---------------------------------------------------------------------------
# Test: compute_log_weights
# ---------------------------------------------------------------------------

class TestComputeLogWeights:
  def test_output_shape(self, identity_jacobian):
    s = jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    log_w = compute_log_weights(s, s_to_z_fn, dz_ds_fn)
    assert log_w.shape == (100,)

  def test_finite_values(self, identity_jacobian):
    """Weights should be finite for reasonable inputs."""
    s = 0.1 * jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    log_w = compute_log_weights(s, s_to_z_fn, dz_ds_fn)
    assert jnp.all(jnp.isfinite(log_w))

  def test_identity_map_weights_are_gaussian(self, identity_jacobian):
    """With identity Jacobian, log weights should equal the Gaussian log-pdf
    up to a constant (the log-det is zero)."""
    s = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    log_w = compute_log_weights(s, s_to_z_fn, dz_ds_fn)
    expected = -0.5 * jnp.sum(s ** 2, axis=-1)
    assert jnp.allclose(log_w - log_w[0], expected - expected[0], atol=1e-10)

  def test_with_real_decoder(self, small_decoder):
    f, z0 = small_decoder
    metric = build_metric(f, z0)
    _, j_s_to_z = get_rnc_jacobians(metric, frame_rotation=None)
    s = 0.1 * jax.random.normal(jax.random.PRNGKey(0), (50, 2))
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_s_to_z)
    log_w = compute_log_weights(s, s_to_z_fn, dz_ds_fn)
    assert log_w.shape == (50,)
    assert jnp.all(jnp.isfinite(log_w))


# ---------------------------------------------------------------------------
# Test: direct_dhsic
# ---------------------------------------------------------------------------

class TestDirectDhsic:
  def test_identity_map_gives_small_dhsic(self, identity_jacobian):
    """With identity Jacobian and Gaussian prior, the components are
    independent, so dHSIC should be small."""
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    stat = direct_dhsic(
      key, size=2.0, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_buffer=10_000, n_resample=500, d=2,
    )
    assert float(stat) < 0.01

  def test_returns_nonnegative(self, identity_jacobian):
    key = jax.random.PRNGKey(1)
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    stat = direct_dhsic(
      key, size=1.0, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_buffer=5_000, n_resample=300, d=2,
    )
    assert float(stat) >= -1e-10

  def test_with_real_decoder(self, small_decoder):
    """Smoke test with a real decoder."""
    f, z0 = small_decoder
    metric = build_metric(f, z0)
    _, j_s_to_z = get_rnc_jacobians(metric, frame_rotation=None)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_s_to_z)
    stat = direct_dhsic(
      jax.random.PRNGKey(0), size=0.5, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_buffer=5_000, n_resample=300, d=2,
    )
    assert jnp.isfinite(stat)
    assert float(stat) >= -1e-10

  def test_sphere_mode(self, identity_jacobian):
    key = jax.random.PRNGKey(2)
    z0 = jnp.zeros(2)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, identity_jacobian)
    stat = direct_dhsic(
      key, size=2.0, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_buffer=5_000, n_resample=300, d=2, use_sphere=True,
    )
    assert jnp.isfinite(stat)
    assert float(stat) >= -1e-10


# ---------------------------------------------------------------------------
# Test: _single_trial_from_key
# ---------------------------------------------------------------------------

class TestSingleTrialFromKey:
  def test_output_shapes(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 5)
    d_dis, d_met, d_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5,
      n_buffer=5_000, n_resample=200,
    )
    assert d_dis.shape == (5,)
    assert d_met.shape == (5,)
    assert d_rand.shape == (5,)

  def test_all_finite(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 0.5, 3)
    d_dis, d_met, d_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5,
      n_buffer=5_000, n_resample=200,
    )
    assert jnp.all(jnp.isfinite(d_dis))
    assert jnp.all(jnp.isfinite(d_met))
    assert jnp.all(jnp.isfinite(d_rand))

  def test_sphere_output_shapes(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 5)
    d_dis, d_met, d_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5,
      n_buffer=5_000, n_resample=200, use_sphere=True,
    )
    assert d_dis.shape == (5,)
    assert d_met.shape == (5,)
    assert d_rand.shape == (5,)


# ---------------------------------------------------------------------------
# Test: run_single_trial
# ---------------------------------------------------------------------------

class TestRunSingleTrial:
  def test_output_shapes(self):
    n_sizes = 5
    result = run_single_trial(
      d=2, n=5, eps=0.5, seed=0,
      n_sizes=n_sizes, max_size=1.0,
      n_buffer=5_000, n_resample=200,
    )
    assert result.sizes.shape == (n_sizes,)
    assert result.dhsic_disentangled.shape == (n_sizes,)
    assert result.dhsic_metric.shape == (n_sizes,)
    assert result.dhsic_random.shape == (n_sizes,)
    assert result.eigenvalues.shape == (2,)


# ---------------------------------------------------------------------------
# Test: run_multi_trial
# ---------------------------------------------------------------------------

class TestRunMultiTrial:
  def test_output_shapes(self):
    n_trials = 2
    n_sizes = 3
    result = run_multi_trial(
      n_trials=n_trials, d=2, n=5, eps=0.5, seed=0,
      n_sizes=n_sizes, max_size=1.0,
      n_buffer=5_000, n_resample=200,
    )
    assert result.sizes.shape == (n_sizes,)
    assert result.dhsic_disentangled.shape == (n_trials, n_sizes)
    assert result.dhsic_metric.shape == (n_trials, n_sizes)
    assert result.dhsic_random.shape == (n_trials, n_sizes)


# ---------------------------------------------------------------------------
# Test: disentangled frame gives lower dHSIC than random
# ---------------------------------------------------------------------------

class TestDisentangledIsBetter:
  def test_disentangled_lower_dhsic_than_random_on_average(self):
    """Averaged over multiple decoders and random rotations, the
    disentangled frame should produce lower dHSIC than random frames."""
    n_trials = 3
    n_sizes = 3
    result = run_multi_trial(
      n_trials=n_trials, d=2, n=5, eps=0.8, seed=42,
      n_sizes=n_sizes, max_size=1.0,
      n_buffer=50_000, n_resample=500,
    )
    last = -1
    mean_dis = float(jnp.mean(result.dhsic_disentangled[:, last]))
    mean_rand = float(jnp.mean(result.dhsic_random[:, last]))
    assert mean_dis <= mean_rand, (
      f"Mean disentangled dHSIC ({mean_dis:.6f}) should be "
      f"<= mean random dHSIC ({mean_rand:.6f})"
    )
