"""Tests for the disentanglement radius figure script."""
import jax
import jax.numpy as jnp
import numpy as np
import pytest

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

from local_geometry.paper.figures.disentanglement_radius import (
  run_disentanglement_radius_sweep,
)


class TestDisentanglementRadiusSweep:

  def test_r_D_decreases_with_epsilon(self):
    """r_D should generally decrease as epsilon increases."""
    result = run_disentanglement_radius_sweep(
      n_trials=3, d=2, n=5, seed=0,
      eps_values=np.array([0.1, 0.5, 1.0, 2.0]),
    )
    assert result.r_D_mean[0] > result.r_D_mean[-1], (
      f"r_D should decrease: {result.r_D_mean}"
    )

  def test_kappa_increases_with_epsilon(self):
    """kappa should generally increase as epsilon increases."""
    result = run_disentanglement_radius_sweep(
      n_trials=3, d=2, n=5, seed=0,
      eps_values=np.array([0.1, 0.5, 1.0, 2.0]),
    )
    assert result.kappa_mean[0] < result.kappa_mean[-1], (
      f"kappa should increase: {result.kappa_mean}"
    )

  def test_result_shapes(self):
    """Output arrays should have the right shapes."""
    eps_values = np.array([0.1, 0.5, 1.0])
    result = run_disentanglement_radius_sweep(
      n_trials=2, d=2, n=5, seed=0,
      eps_values=eps_values,
    )
    n_eps = len(eps_values)
    assert result.r_D_mean.shape == (n_eps,)
    assert result.r_D_std.shape == (n_eps,)
    assert result.kappa_mean.shape == (n_eps,)
    assert result.kappa_std.shape == (n_eps,)
    assert result.n_trials == 2
    assert result.d == 2
    assert result.n == 5

  def test_all_r_D_positive(self):
    """All r_D values should be positive at a generic base point."""
    result = run_disentanglement_radius_sweep(
      n_trials=2, d=2, n=5, seed=0,
      eps_values=np.array([0.5, 1.0]),
    )
    assert np.all(result.r_D_mean > 0), f"Expected all positive: {result.r_D_mean}"
