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

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

from local_geometry.paper.figures.neighborhood_independence import (
  mean_correlation_coefficient,
  filter_to_square,
  filter_to_sphere,
  compute_mcc_at_size,
  masked_mcc,
  weighted_mcc,
  sample_uniform_region,
  direct_mcc,
  _single_trial_from_key,
  build_metric,
  get_disentangled_frame_rotation,
  make_random_rotation,
  build_z_to_s_fn,
  run_neighborhood_independence,
  run_multi_trial,
  MIN_ACCEPTED,
)
from local_geometry.paper.figures.analytic_verification import (
  make_decoder,
  make_coord_fns,
)
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


# ---------------------------------------------------------------------------
# Test: mean_correlation_coefficient (legacy helper)
# ---------------------------------------------------------------------------

class TestMeanCorrelationCoefficient:
  def test_uncorrelated_gives_near_zero(self):
    key = jax.random.PRNGKey(0)
    samples = jax.random.normal(key, (10_000, 3))
    mcc = mean_correlation_coefficient(samples)
    assert mcc < 0.05, f"MCC of iid Gaussian should be ~0, got {mcc}"

  def test_perfectly_correlated(self):
    x = jnp.linspace(-1, 1, 500).reshape(-1, 1)
    samples = jnp.hstack([x, x, -x])
    mcc = mean_correlation_coefficient(samples)
    assert mcc > 0.95, f"MCC of perfectly correlated data should be ~1, got {mcc}"

  def test_d1_gives_zero(self):
    samples = jnp.ones((100, 1))
    assert mean_correlation_coefficient(samples) == 0.0

  def test_two_dimensions(self):
    key = jax.random.PRNGKey(1)
    x = jax.random.normal(key, (5000,))
    noise = 0.1 * jax.random.normal(jax.random.PRNGKey(2), (5000,))
    samples = jnp.stack([x, x + noise], axis=-1)
    mcc = mean_correlation_coefficient(samples)
    assert mcc > 0.9, f"Strongly correlated 2d data should give high MCC, got {mcc}"


# ---------------------------------------------------------------------------
# Test: filter_to_square (legacy helper)
# ---------------------------------------------------------------------------

class TestFilterToSquare:
  def test_all_inside(self):
    s = jnp.array([[0.0, 0.0], [0.1, -0.1], [0.2, 0.2]])
    result = filter_to_square(s, size=1.0)
    assert result.shape == (3, 2)

  def test_some_outside(self):
    s = jnp.array([[0.0, 0.0], [0.6, 0.0], [0.0, -0.6]])
    result = filter_to_square(s, size=1.0)
    assert result.shape == (1, 2)
    assert jnp.allclose(result[0], jnp.array([0.0, 0.0]))

  def test_boundary_included(self):
    s = jnp.array([[0.5, 0.5], [-0.5, -0.5]])
    result = filter_to_square(s, size=1.0)
    assert result.shape == (2, 2)

  def test_empty_result(self):
    s = jnp.array([[10.0, 10.0]])
    result = filter_to_square(s, size=1.0)
    assert result.shape[0] == 0


# ---------------------------------------------------------------------------
# Test: filter_to_sphere
# ---------------------------------------------------------------------------

class TestFilterToSphere:
  def test_all_inside(self):
    s = jnp.array([[0.0, 0.0], [0.1, -0.1], [0.2, 0.2]])
    result = filter_to_sphere(s, size=2.0)
    assert result.shape == (3, 2)

  def test_some_outside(self):
    s = jnp.array([[0.0, 0.0], [0.3, 0.3], [0.6, 0.0]])
    result = filter_to_sphere(s, size=1.0)
    assert result.shape[0] == 2
    assert jnp.allclose(result[0], jnp.array([0.0, 0.0]))
    assert jnp.allclose(result[1], jnp.array([0.3, 0.3]))

  def test_boundary_included(self):
    s = jnp.array([[0.5, 0.0], [0.0, -0.5]])
    result = filter_to_sphere(s, size=1.0)
    assert result.shape == (2, 2)

  def test_empty_result(self):
    s = jnp.array([[10.0, 10.0]])
    result = filter_to_sphere(s, size=1.0)
    assert result.shape[0] == 0

  def test_corner_of_square_excluded(self):
    """A point at the corner of [-0.5, 0.5]^2 lies outside the ball."""
    s = jnp.array([[0.5, 0.5]])
    assert filter_to_square(s, size=1.0).shape[0] == 1
    assert filter_to_sphere(s, size=1.0).shape[0] == 0


# ---------------------------------------------------------------------------
# Test: compute_mcc_at_size (legacy helper)
# ---------------------------------------------------------------------------

class TestComputeMccAtSize:
  def test_nan_for_tiny_neighborhood(self):
    identity_fn = lambda z: z
    samples = jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    mcc, n_acc = compute_mcc_at_size(identity_fn, size=0.001, prior_samples=samples)
    assert n_acc < MIN_ACCEPTED
    assert jnp.isnan(mcc)

  def test_large_neighborhood_gives_valid_mcc(self):
    identity_fn = lambda z: z
    samples = jax.random.normal(jax.random.PRNGKey(0), (50_000, 2))
    mcc, n_acc = compute_mcc_at_size(identity_fn, size=2.0, prior_samples=samples)
    assert n_acc >= MIN_ACCEPTED
    assert not jnp.isnan(mcc)
    assert 0.0 <= mcc <= 1.0


# ---------------------------------------------------------------------------
# Test: masked_mcc (jit-compatible version)
# ---------------------------------------------------------------------------

class TestMaskedMcc:
  def test_uncorrelated_near_zero(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (10_000, 3))
    mcc = masked_mcc(samples, size=100.0)
    assert float(mcc) < 0.05

  def test_nan_for_tiny_neighborhood(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    mcc = masked_mcc(samples, size=0.001)
    assert jnp.isnan(mcc)

  def test_jit_compatible(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (1000, 2))
    jitted = jax.jit(masked_mcc)
    mcc = jitted(samples, 2.0)
    assert not jnp.isnan(mcc)
    assert 0.0 <= float(mcc) <= 1.0

  def test_vmap_over_sizes(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (5000, 2))
    sizes = jnp.array([0.5, 1.0, 2.0, 4.0])
    mccs = jax.vmap(lambda s: masked_mcc(samples, s))(sizes)
    assert mccs.shape == (4,)

  def test_sphere_uncorrelated_near_zero(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (10_000, 3))
    mcc = masked_mcc(samples, size=100.0, use_sphere=True)
    assert float(mcc) < 0.05

  def test_sphere_nan_for_tiny_neighborhood(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    mcc = masked_mcc(samples, size=0.001, use_sphere=True)
    assert jnp.isnan(mcc)

  def test_sphere_jit_compatible(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (1000, 2))
    jitted = jax.jit(masked_mcc, static_argnames=("use_sphere",))
    mcc = jitted(samples, 2.0, use_sphere=True)
    assert not jnp.isnan(mcc)
    assert 0.0 <= float(mcc) <= 1.0

  def test_sphere_fewer_accepted_than_square(self):
    """The sphere is inscribed in the square, so it should accept fewer samples."""
    samples = jax.random.normal(jax.random.PRNGKey(0), (50_000, 2))
    mcc_sq = masked_mcc(samples, size=2.0, use_sphere=False)
    mcc_sp = masked_mcc(samples, size=2.0, use_sphere=True)
    assert not jnp.isnan(mcc_sq)
    assert not jnp.isnan(mcc_sp)


# ---------------------------------------------------------------------------
# Test: weighted_mcc
# ---------------------------------------------------------------------------

class TestWeightedMcc:
  def test_uniform_weights_uncorrelated_near_zero(self):
    """With uniform weights, uncorrelated samples should give low MCC."""
    samples = jax.random.normal(jax.random.PRNGKey(0), (10_000, 3))
    log_w = jnp.zeros(10_000)
    mcc = weighted_mcc(samples, log_w)
    assert float(mcc) < 0.05

  def test_uniform_weights_correlated_high(self):
    """With uniform weights, strongly correlated data should give high MCC."""
    x = jnp.linspace(-1, 1, 500)
    noise = 0.05 * jax.random.normal(jax.random.PRNGKey(0), (500,))
    samples = jnp.stack([x, x + noise], axis=-1)
    log_w = jnp.zeros(500)
    mcc = weighted_mcc(samples, log_w)
    assert float(mcc) > 0.9

  def test_low_n_eff_gives_nan(self):
    """When one weight dominates, n_eff is small and result should be NaN."""
    samples = jax.random.normal(jax.random.PRNGKey(0), (100, 2))
    log_w = jnp.full(100, -1000.0).at[0].set(0.0)
    mcc = weighted_mcc(samples, log_w)
    assert jnp.isnan(mcc)

  def test_jit_compatible(self):
    samples = jax.random.normal(jax.random.PRNGKey(0), (1000, 2))
    log_w = jnp.zeros(1000)
    jitted = jax.jit(weighted_mcc)
    mcc = jitted(samples, log_w)
    assert not jnp.isnan(mcc)
    assert 0.0 <= float(mcc) <= 1.0

  def test_agrees_with_masked_mcc_for_large_region(self):
    """With all samples inside, uniform weights should roughly match masked_mcc."""
    samples = jax.random.normal(jax.random.PRNGKey(0), (5000, 2)) * 0.1
    log_w = jnp.zeros(5000)
    mcc_w = float(weighted_mcc(samples, log_w))
    mcc_m = float(masked_mcc(samples, size=100.0))
    assert abs(mcc_w - mcc_m) < 0.05


# ---------------------------------------------------------------------------
# Test: bias correction
# ---------------------------------------------------------------------------

class TestBiasCorrection:
  def test_mean_corr_coeff_average_near_zero_small_n(self):
    """Averaged over many draws of 50 iid samples, corrected MCC should be near zero."""
    mccs = []
    for seed in range(50):
      samples = jax.random.normal(jax.random.PRNGKey(seed), (50, 3))
      mccs.append(mean_correlation_coefficient(samples))
    avg = sum(mccs) / len(mccs)
    assert avg < 0.04, f"Mean bias-corrected MCC over 50 draws should be ~0, got {avg}"

  def test_masked_mcc_average_near_zero_small_count(self):
    """Averaged over many draws of 80 iid samples, corrected masked_mcc is near zero."""
    mccs = []
    for seed in range(50):
      samples = jax.random.normal(jax.random.PRNGKey(100 + seed), (80, 2))
      mccs.append(float(masked_mcc(samples, size=100.0)))
    avg = sum(mccs) / len(mccs)
    assert avg < 0.04, (
      f"Mean bias-corrected masked_mcc over 50 draws should be ~0, got {avg}"
    )

  def test_weighted_mcc_average_near_zero_small_n(self):
    """Averaged over many draws, corrected weighted_mcc with 80 samples is near zero."""
    mccs = []
    for seed in range(50):
      samples = jax.random.normal(jax.random.PRNGKey(200 + seed), (80, 2))
      log_w = jnp.zeros(80)
      mccs.append(float(weighted_mcc(samples, log_w)))
    avg = sum(mccs) / len(mccs)
    assert avg < 0.04, (
      f"Mean bias-corrected weighted_mcc over 50 draws should be ~0, got {avg}"
    )

  def test_correlated_data_still_detected(self):
    """Bias correction should not suppress genuine correlations."""
    x = jnp.linspace(-1, 1, 100)
    noise = 0.1 * jax.random.normal(jax.random.PRNGKey(3), (100,))
    samples = jnp.stack([x, x + noise], axis=-1)
    mcc = mean_correlation_coefficient(samples)
    assert mcc > 0.8, f"Strongly correlated data should still be detected, got {mcc}"

  def test_masked_mcc_correlated_still_detected(self):
    """Bias correction should not suppress genuine correlations in masked_mcc."""
    x = jnp.linspace(-1, 1, 100)
    noise = 0.1 * jax.random.normal(jax.random.PRNGKey(4), (100,))
    samples = jnp.stack([x, x + noise], axis=-1)
    mcc = masked_mcc(samples, size=100.0)
    assert float(mcc) > 0.8, (
      f"Strongly correlated data should still be detected, got {float(mcc)}"
    )

  def test_bias_correction_nonnegative(self):
    """Bias-corrected MCC should never go negative."""
    for seed in range(10):
      samples = jax.random.normal(jax.random.PRNGKey(seed), (40, 3))
      mcc = mean_correlation_coefficient(samples)
      assert mcc >= 0.0, f"MCC should be >= 0 after clamping, got {mcc}"


# ---------------------------------------------------------------------------
# Test: sample_uniform_region
# ---------------------------------------------------------------------------

class TestSampleUniformRegion:
  def test_cube_bounds(self):
    key = jax.random.PRNGKey(0)
    s = sample_uniform_region(key, 10_000, 2, 0.5, use_sphere=False)
    assert s.shape == (10_000, 2)
    assert jnp.all(jnp.abs(s) <= 0.5)

  def test_sphere_bounds(self):
    key = jax.random.PRNGKey(0)
    s = sample_uniform_region(key, 10_000, 2, 0.5, use_sphere=True)
    assert s.shape == (10_000, 2)
    radii = jnp.linalg.norm(s, axis=-1)
    assert jnp.all(radii <= 0.5 + 1e-10)

  def test_sphere_fills_volume(self):
    """Samples should not all cluster near the origin."""
    key = jax.random.PRNGKey(1)
    s = sample_uniform_region(key, 50_000, 3, 1.0, use_sphere=True)
    radii = jnp.linalg.norm(s, axis=-1)
    assert float(jnp.mean(radii > 0.5)) > 0.5


# ---------------------------------------------------------------------------
# Test: direct_mcc
# ---------------------------------------------------------------------------

class TestDirectMcc:
  def test_identity_map_uncorrelated(self):
    """With the identity Jacobian, direct_mcc on iid Gaussian should give low MCC."""
    d = 2
    j_identity = Jacobian(
      value=jnp.eye(d),
      gradient=jnp.zeros((d, d, d)),
      hessian=jnp.zeros((d, d, d, d)),
    )
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(d)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_identity)
    mcc = direct_mcc(
      key, size=2.0, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_samples=10_000, d=d, use_sphere=False,
    )
    assert not jnp.isnan(mcc)
    assert float(mcc) < 0.1

  def test_small_neighborhood_gives_valid_result(self):
    """direct_mcc should produce a valid result even for tiny neighborhoods
    where rejection sampling would return NaN."""
    d = 2
    j_identity = Jacobian(
      value=jnp.eye(d),
      gradient=jnp.zeros((d, d, d)),
      hessian=jnp.zeros((d, d, d, d)),
    )
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(d)
    s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_identity)
    mcc = direct_mcc(
      key, size=0.01, s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn,
      n_samples=10_000, d=d, use_sphere=False,
    )
    assert not jnp.isnan(mcc), "direct_mcc should not NaN for small neighborhoods"

  def test_with_real_decoder(self):
    """Smoke test using a real decoder and its inverse Jacobian."""
    key = jax.random.PRNGKey(42)
    d, n_net = 2, 5
    f, _, _, _ = make_decoder(key, d=d, n=n_net, eps=0.8)
    z0 = jnp.zeros(d)
    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)
    mcc = direct_mcc(
      jax.random.PRNGKey(1), size=0.5,
      s_to_z_fn=s_to_z_fn, dz_ds_fn=dz_ds_fn, n_samples=5_000, d=d,
    )
    assert not jnp.isnan(mcc)
    assert 0.0 <= float(mcc) <= 1.0


# ---------------------------------------------------------------------------
# Test: _single_trial_from_key (jit smoke test)
# ---------------------------------------------------------------------------

class TestSingleTrialFromKey:
  def test_output_shapes(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 5)
    mcc_dis, mcc_met, mcc_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5, n_prior=5_000,
    )
    assert mcc_dis.shape == (5,)
    assert mcc_met.shape == (5,)
    assert mcc_rand.shape == (5,)

  def test_jit_compatible(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 3)
    jitted = jax.jit(
      _single_trial_from_key,
      static_argnames=("d", "n", "eps", "n_prior"),
    )
    mcc_dis, mcc_met, mcc_rand = jitted(
      key, z0, sizes, d=2, n=5, eps=0.5, n_prior=5_000,
    )
    assert mcc_dis.shape == (3,)

  def test_sphere_output_shapes(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 5)
    mcc_dis, mcc_met, mcc_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5, n_prior=5_000, use_sphere=True,
    )
    assert mcc_dis.shape == (5,)
    assert mcc_met.shape == (5,)
    assert mcc_rand.shape == (5,)

  def test_direct_output_shapes(self):
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.1, 1.0, 5)
    mcc_dis, mcc_met, mcc_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5, n_prior=5_000, use_direct=True,
    )
    assert mcc_dis.shape == (5,)
    assert mcc_met.shape == (5,)
    assert mcc_rand.shape == (5,)

  def test_direct_no_nans_for_small_sizes(self):
    """Direct sampling should produce valid results for small neighborhoods
    where rejection sampling would give NaN."""
    key = jax.random.PRNGKey(0)
    z0 = jnp.zeros(2)
    sizes = jnp.linspace(0.01, 0.05, 3)
    mcc_dis, mcc_met, mcc_rand = _single_trial_from_key(
      key, z0, sizes, d=2, n=5, eps=0.5, n_prior=10_000, use_direct=True,
    )
    assert not jnp.any(jnp.isnan(mcc_dis)), "direct sampling should not NaN at small sizes"
    assert not jnp.any(jnp.isnan(mcc_met))
    assert not jnp.any(jnp.isnan(mcc_rand))


# ---------------------------------------------------------------------------
# Test: disentangled MCC <= random MCC (legacy path)
# ---------------------------------------------------------------------------

class TestDisentangledIsBetter:
  def test_disentangled_lower_mcc_than_random(self, small_decoder):
    """At a moderate neighborhood size, the disentangled frame should
    produce lower MCC than a random frame."""
    f, z0 = small_decoder
    metric = build_metric(f, z0)

    Q_dis, _ = get_disentangled_frame_rotation(metric)
    Q_rand = make_random_rotation(jax.random.PRNGKey(7), 2)

    z_to_s_dis = build_z_to_s_fn(metric, z0, frame_rotation=Q_dis)
    z_to_s_rand = build_z_to_s_fn(metric, z0, frame_rotation=Q_rand)

    prior = jax.random.normal(jax.random.PRNGKey(99), (200_000, 2))
    size = 0.8

    mcc_dis, n_dis = compute_mcc_at_size(z_to_s_dis, size, prior)
    mcc_rand, n_rand = compute_mcc_at_size(z_to_s_rand, size, prior)

    assert n_dis >= MIN_ACCEPTED, f"Not enough disentangled samples ({n_dis})"
    assert n_rand >= MIN_ACCEPTED, f"Not enough random samples ({n_rand})"
    assert mcc_dis <= mcc_rand, (
      f"Disentangled MCC ({mcc_dis:.4f}) should be <= random MCC ({mcc_rand:.4f})"
    )


# ---------------------------------------------------------------------------
# Test: run_neighborhood_independence (single trial)
# ---------------------------------------------------------------------------

class TestRunNeighborhoodIndependence:
  def test_output_shapes(self):
    n_sizes = 5
    result = run_neighborhood_independence(
      d=2, n=5, eps=0.5, seed=0,
      n_sizes=n_sizes, max_size=1.0, n_prior=50_000,
    )
    assert result.sizes.shape == (n_sizes,)
    assert result.mcc_disentangled.shape == (n_sizes,)
    assert result.mcc_metric.shape == (n_sizes,)
    assert result.mcc_random.shape == (n_sizes,)
    assert result.eigenvalues.shape == (2,)


# ---------------------------------------------------------------------------
# Test: multi-trial experiment (vmapped)
# ---------------------------------------------------------------------------

class TestRunMultiTrial:
  def test_output_shapes(self):
    n_trials = 3
    n_sizes = 5
    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_prior=50_000,
    )
    assert result.sizes.shape == (n_sizes,)
    assert result.mcc_disentangled.shape == (n_trials, n_sizes)
    assert result.mcc_metric.shape == (n_trials, n_sizes)
    assert result.mcc_random.shape == (n_trials, n_sizes)

  def test_mean_disentangled_below_random(self):
    """Averaged over multiple decoders, the disentangled frame should
    achieve lower MCC than the random frame at a large neighborhood size."""
    result = run_multi_trial(
      n_trials=3, d=2, n=5, eps=0.8, seed=42,
      n_sizes=5, max_size=1.0, n_prior=200_000,
    )
    last = -1
    dis_vals = result.mcc_disentangled[:, last]
    rand_vals = result.mcc_random[:, last]
    valid_dis = dis_vals[~jnp.isnan(dis_vals)]
    valid_rand = rand_vals[~jnp.isnan(rand_vals)]
    assert len(valid_dis) > 0 and len(valid_rand) > 0
    assert float(jnp.mean(valid_dis)) <= float(jnp.mean(valid_rand)), (
      f"Mean disentangled MCC ({jnp.mean(valid_dis):.4f}) should be "
      f"<= mean random MCC ({jnp.mean(valid_rand):.4f})"
    )
