"""Tests for the disentanglement radius validation figure script."""
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from unittest.mock import patch

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

from local_geometry.paper.figures.disentanglement_radius_validation import (
  run_single_trial,
  run_validation_sweep,
)


@pytest.fixture(scope="module")
def single_trial_result():
  """Run one trial with small parameters for fast testing."""
  key = jax.random.PRNGKey(42)
  z0 = jnp.array([0.3, 0.3])
  normalized_radii = np.array([0.2, 0.5, 1.0, 2.0, 3.0])
  r_D, raw_radii, dhsic_vals = run_single_trial(
    key, z0, d=2, n=5, eps=1.0,
    normalized_radii=normalized_radii,
    n_buffer=50_000, n_resample=300,
  )
  return r_D, raw_radii, dhsic_vals, normalized_radii


class TestSingleTrial:

  def test_r_D_finite_and_positive(self, single_trial_result):
    r_D, _, _, _ = single_trial_result
    assert np.isfinite(r_D), f"r_D should be finite, got {r_D}"
    assert r_D > 0.0, f"r_D should be positive, got {r_D}"

  def test_dhsic_values_finite(self, single_trial_result):
    _, _, dhsic_vals, _ = single_trial_result
    assert np.all(np.isfinite(dhsic_vals)), (
      f"All dHSIC values should be finite: {dhsic_vals}"
    )

  def test_dhsic_values_nonnegative(self, single_trial_result):
    _, _, dhsic_vals, _ = single_trial_result
    assert np.all(dhsic_vals >= -1e-10), (
      f"dHSIC values should be non-negative: {dhsic_vals}"
    )

  def test_raw_radii_scale_with_r_D(self, single_trial_result):
    r_D, raw_radii, _, normalized_radii = single_trial_result
    expected = normalized_radii * r_D
    np.testing.assert_allclose(raw_radii, expected, rtol=1e-10)

  def test_dhsic_higher_outside_r_D(self, single_trial_result):
    """dHSIC at r=3*r_D should be larger than at r=0.2*r_D."""
    _, _, dhsic_vals, normalized_radii = single_trial_result
    idx_small = np.argmin(np.abs(normalized_radii - 0.2))
    idx_large = np.argmin(np.abs(normalized_radii - 3.0))
    assert dhsic_vals[idx_large] > dhsic_vals[idx_small], (
      f"dHSIC at r/r_D=3.0 ({dhsic_vals[idx_large]:.6f}) should exceed "
      f"dHSIC at r/r_D=0.2 ({dhsic_vals[idx_small]:.6f})"
    )

  def test_output_shapes(self, single_trial_result):
    r_D, raw_radii, dhsic_vals, normalized_radii = single_trial_result
    n_radii = len(normalized_radii)
    assert raw_radii.shape == (n_radii,)
    assert dhsic_vals.shape == (n_radii,)


class TestValidationSweep:

  def test_result_shapes(self):
    result = run_validation_sweep(
      n_trials=2, d=2, n=5, seed=0,
      eps_values=np.array([0.5, 1.0]),
      normalized_radii=np.array([0.5, 1.0, 2.0]),
      n_buffer=30_000, n_resample=200,
    )
    assert result.eps_values.shape == (2,)
    assert result.normalized_radii.shape == (3,)
    for eps_f in [0.5, 1.0]:
      assert result.dhsic_all[eps_f].shape == (2, 3)
      assert result.r_D_all[eps_f].shape == (2,)
      assert result.raw_radii_all[eps_f].shape == (2, 3)

class TestWeightedSingleTrial:

  @pytest.fixture(scope="class")
  def weighted_result(self):
    """Run one trial with weighted dHSIC for testing."""
    key = jax.random.PRNGKey(42)
    z0 = jnp.array([0.3, 0.3])
    normalized_radii = np.array([0.2, 0.5, 1.0, 2.0, 3.0])
    r_D, raw_radii, dhsic_vals = run_single_trial(
      key, z0, d=2, n=5, eps=1.0,
      normalized_radii=normalized_radii,
      n_buffer=50_000, n_weighted=800, weighted=True,
    )
    return r_D, raw_radii, dhsic_vals, normalized_radii

  def test_r_D_finite_and_positive(self, weighted_result):
    r_D, _, _, _ = weighted_result
    assert np.isfinite(r_D)
    assert r_D > 0.0

  def test_dhsic_values_finite(self, weighted_result):
    _, _, dhsic_vals, _ = weighted_result
    assert np.all(np.isfinite(dhsic_vals))

  def test_dhsic_higher_outside_r_D(self, weighted_result):
    """Weighted dHSIC at r=3*r_D should exceed r=0.2*r_D."""
    _, _, dhsic_vals, normalized_radii = weighted_result
    idx_small = np.argmin(np.abs(normalized_radii - 0.2))
    idx_large = np.argmin(np.abs(normalized_radii - 3.0))
    assert dhsic_vals[idx_large] > dhsic_vals[idx_small], (
      f"Weighted dHSIC at r/r_D=3.0 ({dhsic_vals[idx_large]:.6f}) should "
      f"exceed r/r_D=0.2 ({dhsic_vals[idx_small]:.6f})"
    )

  def test_output_shapes(self, weighted_result):
    _, raw_radii, dhsic_vals, normalized_radii = weighted_result
    n_radii = len(normalized_radii)
    assert raw_radii.shape == (n_radii,)
    assert dhsic_vals.shape == (n_radii,)


class TestValidationSweep:

  def test_r_D_decreases_with_eps(self):
    result = run_validation_sweep(
      n_trials=2, d=2, n=5, seed=0,
      eps_values=np.array([0.3, 1.5]),
      normalized_radii=np.array([1.0]),
      n_buffer=30_000, n_resample=200,
    )
    mean_r_D_low = np.mean(result.r_D_all[0.3])
    mean_r_D_high = np.mean(result.r_D_all[1.5])
    assert mean_r_D_low > mean_r_D_high, (
      f"r_D should decrease with eps: {mean_r_D_low:.3f} vs {mean_r_D_high:.3f}"
    )


def test_weighted_path_uses_n_weighted_buffer():
  key = jax.random.PRNGKey(0)
  z0 = jnp.array([0.3, 0.3])
  normalized_radii = np.array([1.0])
  calls = []

  def fake_direct_dhsic_weighted(
    key, size, s_to_z_fn, dz_ds_fn, n_buffer, d, use_sphere=False
  ):
    calls.append((n_buffer, d))
    return jnp.array(0.0)

  with patch(
    "local_geometry.paper.figures.disentanglement_radius_validation.direct_dhsic_weighted",
    new=fake_direct_dhsic_weighted,
  ):
    run_single_trial(
      key, z0, d=2, n=5, eps=1.0, normalized_radii=normalized_radii,
      n_buffer=50_000, n_weighted=777, weighted=True,
    )

  assert calls == [(777, 2)]


def test_exact_method_runs_single_trial():
  key = jax.random.PRNGKey(123)
  z0 = jnp.array([0.3, 0.3])
  normalized_radii = np.array([0.5])
  r_D, raw_radii, dhsic_vals = run_single_trial(
    key, z0, d=2, n=4, eps=0.8, normalized_radii=normalized_radii,
    n_buffer=300, n_resample=80, method="exact",
  )
  assert np.isfinite(r_D)
  assert raw_radii.shape == (1,)
  assert dhsic_vals.shape == (1,)
  assert np.isfinite(dhsic_vals[0])
