"""Tests for the SSL disentanglement experiment utility functions."""
import jax
import jax.numpy as jnp
import pytest

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

from local_geometry.experiments.ssl_disentanglement_design import (
  log_gaussian_mixture,
  linear_probe_accuracy,
)


class TestLogGaussianMixture:

  def test_single_component_matches_gaussian(self):
    """A single-component mixture should match a standard Gaussian."""
    d = 3
    z = jnp.array([0.5, -0.3, 0.1])
    means = jnp.zeros((1, d))
    log_mix = log_gaussian_mixture(z, means)
    log_gauss = -0.5 * d * jnp.log(2 * jnp.pi) - 0.5 * jnp.sum(z ** 2)
    assert float(log_mix) == pytest.approx(float(log_gauss), abs=1e-12)

  def test_uniform_mixture_at_origin(self):
    """At the origin with symmetric means, the mixture has a known value."""
    d = 2
    means = jnp.array([[1.0, 0.0], [-1.0, 0.0]])
    z = jnp.zeros(d)
    log_p = log_gaussian_mixture(z, means)
    expected = jnp.log(
      0.5 * jnp.exp(-0.5) / (2 * jnp.pi)
      + 0.5 * jnp.exp(-0.5) / (2 * jnp.pi)
    )
    assert float(log_p) == pytest.approx(float(expected), abs=1e-10)

  def test_gradient_computable(self):
    """The gradient of the log mixture should be computable via JAX."""
    d = 3
    means = jnp.array([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    z = jnp.array([0.2, -0.1, 0.3])
    grad = jax.grad(log_gaussian_mixture)(z, means)
    assert grad.shape == (d,)
    assert jnp.all(jnp.isfinite(grad))

  def test_hessian_computable(self):
    """The Hessian should be computable for the r_D calculation."""
    d = 3
    means = jnp.array([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    z = jnp.array([0.2, -0.1, 0.3])
    H = jax.hessian(log_gaussian_mixture)(z, means)
    assert H.shape == (d, d)
    assert jnp.allclose(H, H.T, atol=1e-10)

  def test_hessian_offdiagonal_nonzero_for_mixture(self):
    """A mixture whose means differ along multiple axes should have
    non-zero off-diagonal Hessian (the density does not factorize)."""
    d = 2
    means = jnp.array([[1.0, 1.0], [-1.0, -1.0]])
    z = jnp.array([0.5, 0.3])
    H = jax.hessian(log_gaussian_mixture)(z, means)
    off_diag = jnp.abs(H[0, 1])
    assert float(off_diag) > 1e-8, "Mixture Hessian should have off-diagonal terms"

  def test_custom_mixing_weights(self):
    """Non-uniform mixing weights should change the density at a
    point that is not equidistant from the means."""
    d = 2
    means = jnp.array([[1.0, 0.0], [-1.0, 0.0]])
    z = jnp.array([0.5, 0.0])
    log_uniform = log_gaussian_mixture(z, means)
    log_weighted = log_gaussian_mixture(
      z, means, log_mixing_weights=jnp.array([jnp.log(0.9), jnp.log(0.1)])
    )
    assert float(log_uniform) != pytest.approx(float(log_weighted), abs=1e-5)


class TestLinearProbe:

  def test_perfect_separation(self):
    """Linear probe should achieve near-perfect accuracy on linearly separable data."""
    key = jax.random.PRNGKey(42)
    k1, k2 = jax.random.split(key)
    n_per_class = 200
    d = 5

    class_0 = jax.random.normal(k1, (n_per_class, d)) + jnp.array([3.0] * d)
    class_1 = jax.random.normal(k2, (n_per_class, d)) - jnp.array([3.0] * d)

    features = jnp.concatenate([class_0, class_1])
    labels = jnp.concatenate([jnp.zeros(n_per_class, dtype=jnp.int32),
                              jnp.ones(n_per_class, dtype=jnp.int32)])

    perm = jax.random.permutation(jax.random.PRNGKey(0), len(features))
    features, labels = features[perm], labels[perm]
    split = 300
    acc = linear_probe_accuracy(
      features[:split], labels[:split],
      features[split:], labels[split:],
      num_classes=2, lr=0.01, num_steps=500,
    )
    assert acc > 0.9, f"Expected near-perfect accuracy, got {acc}"

  def test_random_features_near_chance(self):
    """Random features should give near-chance accuracy."""
    key = jax.random.PRNGKey(7)
    n = 200
    d = 5
    features = jax.random.normal(key, (n, d))
    labels = jax.random.randint(jax.random.PRNGKey(8), (n,), 0, 3)

    acc = linear_probe_accuracy(
      features[:150], labels[:150],
      features[150:], labels[150:],
      num_classes=3, lr=0.01, num_steps=200,
    )
    assert acc < 0.6, f"Expected near-chance accuracy, got {acc}"
