"""Tests for the multi-trial frame optimality C(r) experiment."""
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.frame_optimality_multi_trial import (
  run_single_trial,
  run_multi_trial,
  SingleTrialResult,
  MultiTrialResult,
)


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

@pytest.fixture
def single_trial_2d():
  """Run a single trial with d=2, n=5, small radii grid for fast testing."""
  key = jax.random.PRNGKey(42)
  z0 = jnp.zeros(2)
  radii = jnp.linspace(0.0, 1.0, 5)
  return run_single_trial(
    key, z0, radii, d=2, n=5, eps=0.5, n_samples=100,
  )


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

class TestRunSingleTrial:
  def test_return_type(self, single_trial_2d):
    assert isinstance(single_trial_2d, SingleTrialResult)

  def test_output_shapes(self, single_trial_2d):
    result = single_trial_2d
    assert result.coupling_disentangled.shape == (5,)
    assert result.coupling_metric.shape == (5,)
    assert result.coupling_random.shape == (5,)
    assert result.eigenvalues.shape == (2,)

  def test_disentangled_coupling_at_origin_is_zero(self, single_trial_2d):
    """C(0) for the disentangled frame should be ~0 by Theorem 2."""
    c0_dis = float(single_trial_2d.coupling_disentangled[0])
    assert c0_dis < 1e-10, f"Disentangled C(0) = {c0_dis}, expected ~0"

  def test_metric_coupling_at_origin_is_positive(self, single_trial_2d):
    """C(0) for the metric (SVD) frame should be > 0 in general."""
    c0_met = float(single_trial_2d.coupling_metric[0])
    c0_dis = float(single_trial_2d.coupling_disentangled[0])
    assert c0_met > c0_dis, (
      f"Metric C(0) = {c0_met} should exceed disentangled C(0) = {c0_dis}"
    )

  def test_coupling_grows_with_radius(self, single_trial_2d):
    """C(r) should generally increase with r for the disentangled frame."""
    c_dis = np.asarray(single_trial_2d.coupling_disentangled)
    assert c_dis[-1] > c_dis[0], (
      f"C(r_max)={c_dis[-1]} should exceed C(0)={c_dis[0]}"
    )

  def test_all_values_nonneg(self, single_trial_2d):
    """C(r) is a sum of squares, so it must be non-negative."""
    result = single_trial_2d
    for arr in [result.coupling_disentangled, result.coupling_metric,
                result.coupling_random]:
      assert np.all(np.asarray(arr) >= -1e-14)

  def test_different_seeds_give_different_results(self):
    z0 = jnp.zeros(2)
    radii = jnp.array([0.0, 0.5])
    r1 = run_single_trial(
      jax.random.PRNGKey(0), z0, radii, d=2, n=5, eps=0.5, n_samples=100,
    )
    r2 = run_single_trial(
      jax.random.PRNGKey(99), z0, radii, d=2, n=5, eps=0.5, n_samples=100,
    )
    assert not jnp.allclose(
      r1.coupling_random, r2.coupling_random
    ), "Different seeds should produce different random-frame couplings"


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

class TestRunMultiTrial:
  def test_return_type_and_shapes(self):
    result = run_multi_trial(
      n_trials=2, d=2, n=5, eps=0.5, seed=0,
      n_radii=3, r_max=0.5, n_samples=50,
    )
    assert isinstance(result, MultiTrialResult)
    assert result.coupling_disentangled.shape == (2, 3)
    assert result.coupling_metric.shape == (2, 3)
    assert result.coupling_random.shape == (2, 3)
    assert result.radii.shape == (3,)

  def test_disentangled_mean_c0_near_zero(self):
    result = run_multi_trial(
      n_trials=3, d=2, n=5, eps=0.5, seed=7,
      n_radii=3, r_max=0.5, n_samples=100,
    )
    mean_c0 = float(np.mean(np.asarray(result.coupling_disentangled[:, 0])))
    assert mean_c0 < 1e-8, f"Mean disentangled C(0) = {mean_c0}, expected ~0"
