"""Tests for the analytic verification experiments."""
import jax
import jax.numpy as jnp
import pytest

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

from local_geometry.paper.figures.analytic_experiments import (
  sample_sphere,
  off_diagonal_frobenius_sq,
  hessian_log_ps,
  make_coupling_fn,
  compute_coupling,
  build_metric_and_log_det,
  get_disentangled_frame_rotation,
  build_frame_coord_fns,
  make_random_rotation,
  get_entanglement_eigenvalues,
  get_scalar_curvature,
  compute_coupling_rate,
  compute_disentanglement_radius,
  compute_effective_radius,
  pushforward_directions,
  principal_angles,
  run_invariance_test,
  run_density_contours,
)
from local_geometry.paper.figures.analytic_verification import (
  make_decoder,
  log_pz,
)


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

@pytest.fixture
def small_decoder_2d():
  """A small d=2, n=3 decoder for fast testing."""
  key = jax.random.PRNGKey(42)
  f, _, _, _ = make_decoder(key, d=2, n=3, eps=0.5)
  z0 = jnp.array([0.3, -0.2])
  return f, z0


@pytest.fixture
def small_decoder_3d():
  """A small d=3, n=6 decoder for 3d tests."""
  key = jax.random.PRNGKey(7)
  f, _, _, _ = make_decoder(key, d=3, n=6, eps=0.5)
  z0 = jnp.zeros(3)
  return f, z0


# ---------------------------------------------------------------------------
# Test: sphere sampling
# ---------------------------------------------------------------------------

class TestSphereSampling:
  def test_correct_shape(self):
    pts = sample_sphere(jax.random.PRNGKey(0), d=5, r=1.0, n_samples=100)
    assert pts.shape == (100, 5)

  def test_correct_radius(self):
    pts = sample_sphere(jax.random.PRNGKey(0), d=5, r=2.5, n_samples=200)
    norms = jnp.linalg.norm(pts, axis=-1)
    assert jnp.allclose(norms, 2.5, atol=1e-12)

  def test_different_keys_give_different_points(self):
    p1 = sample_sphere(jax.random.PRNGKey(0), d=3, r=1.0, n_samples=10)
    p2 = sample_sphere(jax.random.PRNGKey(1), d=3, r=1.0, n_samples=10)
    assert not jnp.allclose(p1, p2)


# ---------------------------------------------------------------------------
# Test: off-diagonal Frobenius
# ---------------------------------------------------------------------------

class TestOffDiagonalFrobenius:
  def test_diagonal_matrix_gives_zero(self):
    H = jnp.diag(jnp.array([1.0, 2.0, 3.0]))
    assert jnp.isclose(off_diagonal_frobenius_sq(H), 0.0, atol=1e-15)

  def test_known_value(self):
    H = jnp.array([[1.0, 0.5], [0.5, 2.0]])
    expected = 2 * 0.5**2
    assert jnp.isclose(off_diagonal_frobenius_sq(H), expected, atol=1e-15)

  def test_identity_gives_zero(self):
    H = jnp.eye(4)
    assert jnp.isclose(off_diagonal_frobenius_sq(H), 0.0, atol=1e-15)


# ---------------------------------------------------------------------------
# Test: C(0) = 0 for disentangled frame
# ---------------------------------------------------------------------------

class TestCouplingAtOrigin:
  def test_disentangled_coupling_zero_at_origin(self, small_decoder_2d):
    """The disentangled frame should have C(0) = 0 by construction."""
    f, z0 = small_decoder_2d
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    s_to_z_fn, dz_ds_fn, _ = build_frame_coord_fns(
      metric, z0, frame_rotation=Q_dis
    )
    cfn = make_coupling_fn(s_to_z_fn, dz_ds_fn, d=2)
    key = jax.random.PRNGKey(0)
    c0 = compute_coupling(cfn, 0.0, key, d=2)
    assert float(c0) < 1e-12, f"C(0) for disentangled frame should be ~0, got {c0}"

  def test_metric_frame_coupling_nonzero_at_origin(self, small_decoder_2d):
    """The metric frame (no disentangling rotation) should generally have C(0) > 0."""
    f, z0 = small_decoder_2d
    metric, _ = build_metric_and_log_det(f, z0)
    s_to_z_fn, dz_ds_fn, _ = build_frame_coord_fns(
      metric, z0, frame_rotation=None
    )
    cfn = make_coupling_fn(s_to_z_fn, dz_ds_fn, d=2)
    key = jax.random.PRNGKey(0)
    c0 = compute_coupling(cfn, 0.0, key, d=2)
    assert float(c0) > 1e-6, f"C(0) for metric frame should be > 0, got {c0}"


# ---------------------------------------------------------------------------
# Test: C(r) increases with r for disentangled frame
# ---------------------------------------------------------------------------

class TestCouplingGrowth:
  def test_coupling_grows_with_radius(self, small_decoder_2d):
    """C(r) should grow with r for the disentangled frame."""
    f, z0 = small_decoder_2d
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    s_to_z_fn, dz_ds_fn, _ = build_frame_coord_fns(
      metric, z0, frame_rotation=Q_dis
    )
    cfn = make_coupling_fn(s_to_z_fn, dz_ds_fn, d=2)
    key = jax.random.PRNGKey(0)
    c_small = compute_coupling(cfn, 0.3, key, d=2, n_samples=200)
    c_large = compute_coupling(cfn, 1.0, key, d=2, n_samples=200)
    assert float(c_large) > float(c_small), (
      f"C(1.0)={c_large} should exceed C(0.3)={c_small}"
    )


# ---------------------------------------------------------------------------
# Test: Hessian of log p_s is diagonal at origin for disentangled frame
# ---------------------------------------------------------------------------

class TestHessianDiagonality:
  def test_hessian_diagonal_at_origin(self, small_decoder_2d):
    f, z0 = small_decoder_2d
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, eigenvalues = get_disentangled_frame_rotation(metric)
    s_to_z_fn, dz_ds_fn, _ = build_frame_coord_fns(
      metric, z0, frame_rotation=Q_dis
    )
    s0 = jnp.zeros(2)
    H = hessian_log_ps(s0, s_to_z_fn, dz_ds_fn)
    off_diag = H - jnp.diag(jnp.diag(H))
    assert jnp.allclose(off_diag, 0.0, atol=1e-8), (
      f"Off-diagonal Hessian should be ~0, max={jnp.max(jnp.abs(off_diag))}"
    )

  def test_hessian_eigenvalues_match_entanglement(self, small_decoder_2d):
    """Diagonal entries of the Hessian should equal eigenvalues of D."""
    f, z0 = small_decoder_2d
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, eigenvalues = get_disentangled_frame_rotation(metric)
    s_to_z_fn, dz_ds_fn, _ = build_frame_coord_fns(
      metric, z0, frame_rotation=Q_dis
    )
    s0 = jnp.zeros(2)
    H = hessian_log_ps(s0, s_to_z_fn, dz_ds_fn)
    diag_sorted = jnp.sort(jnp.diag(H))
    eigs_sorted = jnp.sort(eigenvalues)
    assert jnp.allclose(diag_sorted, eigs_sorted, atol=1e-6), (
      f"Hessian diagonal {diag_sorted} should match eigenvalues {eigs_sorted}"
    )


# ---------------------------------------------------------------------------
# Test: random rotation is orthogonal
# ---------------------------------------------------------------------------

class TestRandomRotation:
  def test_orthogonality(self):
    Q = make_random_rotation(jax.random.PRNGKey(0), 5)
    assert jnp.allclose(Q @ Q.T, jnp.eye(5), atol=1e-12)
    assert jnp.allclose(Q.T @ Q, jnp.eye(5), atol=1e-12)

  def test_determinant_one_or_minus_one(self):
    Q = make_random_rotation(jax.random.PRNGKey(0), 5)
    det = jnp.linalg.det(Q)
    assert jnp.isclose(jnp.abs(det), 1.0, atol=1e-12)


# ---------------------------------------------------------------------------
# Test: entanglement eigenvalues
# ---------------------------------------------------------------------------

class TestEntanglementEigenvalues:
  def test_linear_decoder_gives_degenerate_eigenvalues(self):
    """When eps=0, the decoder is linear, curvature vanishes, and D = -I."""
    key = jax.random.PRNGKey(0)
    d, n = 3, 6
    f, _, _, _ = make_decoder(key, d=d, n=n, eps=0.0)
    z0 = jnp.zeros(d)
    metric, _ = build_metric_and_log_det(f, z0)
    eigs = get_entanglement_eigenvalues(metric)
    assert jnp.allclose(eigs, eigs[0] * jnp.ones(d), atol=1e-8), (
      f"Linear decoder should have degenerate eigenvalues, got {eigs}"
    )

  def test_nonzero_eps_gives_nondegenerate(self):
    """Nonzero eps should generally give distinct eigenvalues."""
    key = jax.random.PRNGKey(0)
    d, n = 3, 6
    f, _, _, _ = make_decoder(key, d=d, n=n, eps=1.0)
    z0 = jnp.zeros(d)
    metric, _ = build_metric_and_log_det(f, z0)
    eigs = get_entanglement_eigenvalues(metric)
    gaps = jnp.diff(eigs)
    assert jnp.any(jnp.abs(gaps) > 1e-4), (
      f"Non-linear decoder should have non-degenerate eigenvalues, got {eigs}"
    )


# ---------------------------------------------------------------------------
# Test: scalar curvature
# ---------------------------------------------------------------------------

class TestScalarCurvature:
  def test_flat_metric_zero_curvature(self):
    """Linear decoder (eps=0) gives a flat metric with zero scalar curvature."""
    key = jax.random.PRNGKey(0)
    d, n = 3, 6
    f, _, _, _ = make_decoder(key, d=d, n=n, eps=0.0)
    z0 = jnp.zeros(d)
    metric, _ = build_metric_and_log_det(f, z0)
    R = get_scalar_curvature(metric)
    assert jnp.isclose(R, 0.0, atol=1e-8), (
      f"Linear decoder should have zero curvature, got {R}"
    )

  def test_curved_metric_nonzero(self):
    """Nonlinear decoder at a non-origin point should have nonzero scalar curvature."""
    key = jax.random.PRNGKey(0)
    d, n = 3, 6
    f, _, _, _ = make_decoder(key, d=d, n=n, eps=1.0)
    z0 = jnp.array([0.3, -0.2, 0.1])
    metric, _ = build_metric_and_log_det(f, z0)
    R = get_scalar_curvature(metric)
    assert jnp.abs(R) > 1e-4, (
      f"Non-linear decoder at non-origin should have nonzero curvature, got {R}"
    )


# ---------------------------------------------------------------------------
# Test: effective radius
# ---------------------------------------------------------------------------

class TestEffectiveRadius:
  def test_returns_last_if_never_exceeded(self):
    radii = jnp.linspace(0, 2, 10)
    couplings = jnp.zeros(10)
    eigs = jnp.array([-2.0, -1.0])
    r_star = compute_effective_radius(radii, couplings, eigs)
    assert r_star == float(radii[-1])

  def test_returns_first_exceeding_radius(self):
    radii = jnp.array([0.0, 0.5, 1.0, 1.5, 2.0])
    couplings = jnp.array([0.0, 0.2, 0.5, 1.0, 2.0])
    eigs = jnp.array([-1.0, -1.0])
    r_star = compute_effective_radius(radii, couplings, eigs, threshold_frac=0.1)
    assert r_star == 0.5


# ---------------------------------------------------------------------------
# Test: principal angles
# ---------------------------------------------------------------------------

class TestPrincipalAngles:
  def test_same_subspace_gives_zero_angles(self):
    U = jnp.eye(5, 3)
    angles = principal_angles(U, U)
    assert jnp.allclose(angles, 0.0, atol=1e-12)

  def test_orthogonal_subspaces(self):
    U = jnp.eye(4)[:, :2]
    V = jnp.eye(4)[:, 2:]
    angles = principal_angles(U, V)
    assert jnp.allclose(angles, jnp.pi / 2, atol=1e-12)


# ---------------------------------------------------------------------------
# Test: parameterization invariance (small scale)
# ---------------------------------------------------------------------------

class TestParameterizationInvariance:
  def test_linear_invariance(self):
    """Image-space directions should agree under orthogonal reparameterization."""
    result = run_invariance_test(
      d=3, n=6, eps=0.5, seed=42, reparam_type="linear"
    )
    assert jnp.allclose(result.angles, 0.0, atol=1e-5), (
      f"Principal angles should be ~0, got {result.angles}"
    )

  def test_nonlinear_invariance(self):
    """Image-space directions should agree under nonlinear reparameterization."""
    result = run_invariance_test(
      d=3, n=6, eps=0.5, seed=42, reparam_type="nonlinear"
    )
    assert jnp.allclose(result.angles, 0.0, atol=1e-4), (
      f"Principal angles should be ~0, got {result.angles}"
    )


# ---------------------------------------------------------------------------
# Test: density contours
# ---------------------------------------------------------------------------

class TestDensityContours:
  def test_output_shapes(self):
    result = run_density_contours(d=2, n=3, eps=0.5, seed=0, n_grid=10)
    assert result.S1.shape == (10, 10)
    assert result.log_ps_disentangled.shape == (10, 10)
    assert result.log_ps_rotated.shape == (10, 10)

  def test_disentangled_contours_more_axis_aligned(self):
    """The Hessian at the origin in the disentangled frame should be more
    diagonal than in the rotated frame."""
    result = run_density_contours(d=2, n=5, eps=0.8, seed=0, n_grid=20, R=0.5)
    center = result.S1.shape[0] // 2

    def hessian_off_diag_at_center(lps):
      s_1d = result.S1[0, :]
      ds = float(s_1d[1] - s_1d[0])
      lps_np = jnp.array(lps)
      d2_ds1ds2 = (
        lps_np[center+1, center+1] - lps_np[center+1, center-1]
        - lps_np[center-1, center+1] + lps_np[center-1, center-1]
      ) / (4 * ds * ds)
      return float(jnp.abs(d2_ds1ds2))

    off_diag_dis = hessian_off_diag_at_center(result.log_ps_disentangled)
    off_diag_rot = hessian_off_diag_at_center(result.log_ps_rotated)
    assert off_diag_dis < off_diag_rot, (
      f"Disentangled off-diag ({off_diag_dis}) should be < rotated ({off_diag_rot})"
    )


# ---------------------------------------------------------------------------
# Test: disentanglement radius
# ---------------------------------------------------------------------------

class TestDisentanglementRadius:
  """Tests for compute_coupling_rate and compute_disentanglement_radius."""

  def test_kappa_matches_coupling_ratio(self, small_decoder_2d):
    """kappa should match C(r)/r^2 for small r."""
    f, z0 = small_decoder_2d
    d = z0.shape[0]
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)
    s_to_z_fn, dz_ds_fn = fns[0], fns[1]

    kappa, T = compute_coupling_rate(s_to_z_fn, dz_ds_fn, d)
    coupling_fn = make_coupling_fn(s_to_z_fn, dz_ds_fn, d)

    r_small = 0.01
    key = jax.random.PRNGKey(99)
    C_r = compute_coupling(coupling_fn, r_small, key, d, n_samples=5000)
    kappa_empirical = float(C_r) / (r_small ** 2)

    assert float(kappa) == pytest.approx(kappa_empirical, rel=0.1), (
      f"Analytic kappa={float(kappa):.6f} vs empirical C(r)/r^2={kappa_empirical:.6f}"
    )

  def test_kappa_positive_on_curved_manifold(self, small_decoder_2d):
    """kappa should be positive when the manifold has curvature."""
    f, z0 = small_decoder_2d
    d = z0.shape[0]
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)

    kappa, _ = compute_coupling_rate(fns[0], fns[1], d)
    assert float(kappa) > 0.0, "kappa should be positive on a curved manifold"

  def test_r_D_finite_on_curved_manifold(self, small_decoder_2d):
    """r_D should be finite and positive on a curved manifold."""
    f, z0 = small_decoder_2d
    d = z0.shape[0]
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)

    r_D, kappa, _ = compute_disentanglement_radius(fns[0], fns[1], d)
    assert jnp.isfinite(r_D), "r_D should be finite on a curved manifold"
    assert float(r_D) > 0.0, "r_D should be positive"

  def test_kappa_zero_at_symmetric_point(self):
    """At z0=0 the polynomial decoder has vanishing third derivatives of
    log p_s due to the special structure of the cubic term, so kappa=0."""
    key = jax.random.PRNGKey(42)
    d, n = 2, 5
    z0 = jnp.zeros(d)
    f, _, _, _ = make_decoder(key, d=d, n=n, eps=1.0)
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)
    kappa, _ = compute_coupling_rate(fns[0], fns[1], d)
    assert float(kappa) == pytest.approx(0.0, abs=1e-10)

  def test_kappa_decreases_with_less_curvature(self):
    """kappa should decrease as curvature (eps) decreases at a generic point."""
    key = jax.random.PRNGKey(42)
    d, n = 2, 5
    z0 = jnp.array([0.3, -0.2])

    kappas = []
    for eps in [1.0, 0.5, 0.1]:
      f, _, _, _ = make_decoder(key, d=d, n=n, eps=eps)
      metric, _ = build_metric_and_log_det(f, z0)
      Q_dis, _ = get_disentangled_frame_rotation(metric)
      fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)
      kappa, _ = compute_coupling_rate(fns[0], fns[1], d)
      kappas.append(float(kappa))

    assert kappas[0] > kappas[1] > kappas[2], (
      f"kappa should decrease with eps: got {kappas}"
    )

  def test_T_tensor_shape(self, small_decoder_3d):
    """T tensor should have shape (d, d, d)."""
    f, z0 = small_decoder_3d
    d = z0.shape[0]
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)

    _, T = compute_coupling_rate(fns[0], fns[1], d)
    assert T.shape == (d, d, d), f"Expected ({d},{d},{d}), got {T.shape}"

  def test_T_tensor_symmetric_in_first_two_indices(self, small_decoder_2d):
    """T_{ijk} should be symmetric in i,j since the Hessian is symmetric."""
    f, z0 = small_decoder_2d
    d = z0.shape[0]
    metric, _ = build_metric_and_log_det(f, z0)
    Q_dis, _ = get_disentangled_frame_rotation(metric)
    fns = build_frame_coord_fns(metric, z0, frame_rotation=Q_dis)

    _, T = compute_coupling_rate(fns[0], fns[1], d)
    assert jnp.allclose(T, jnp.transpose(T, (1, 0, 2)), atol=1e-10), (
      "T should be symmetric in first two indices"
    )
