"""Tests for AbstractMetricLocalModel (PCF), NeuralMetricModel, and PCF losses."""
import jax
import jax.numpy as jnp
import jax.random as random
import numpy as np
import equinox as eqx
import optax
import pytest
from jaxtyping import Array, Float

from local_disentanglement.models.metric import AbstractMetricLocalModel
from local_disentanglement.losses.pcf_losses import (
  density_loss,
  density_loss_batch,
  regularization_loss,
  regularization_loss_batch,
  disentanglement_radius_loss,
  disentanglement_radius_loss_batch,
  pcf_loss,
  pcf_loss_batch,
)
from local_geometry.rica_jax.riemannian_ica import (
  taylor_data,
  christoffel,
  christoffel_gradient,
  ricci_from_christoffel,
  orthonormal_frame,
)


class ConstantMetricModel(AbstractMetricLocalModel):
  """Test model with a fixed diagonal metric g = diag(eigenvalues)."""
  _eigenvalues: Float[Array, "D"]

  def __init__(self, eigenvalues):
    self._eigenvalues = jnp.array(eigenvalues, dtype=jnp.float64)

  def metric_components(self, x, sigma=None):
    return jnp.diag(self._eigenvalues)


class IdentityMetricModel(AbstractMetricLocalModel):
  """Test model with g = I everywhere."""
  dim: int = eqx.field(static=True)

  def __init__(self, dim):
    self.dim = dim

  def metric_components(self, x, sigma=None):
    return jnp.eye(self.dim)


class VaryingMetricModel(AbstractMetricLocalModel):
  """Test model with spatially varying diagonal metric."""

  def metric_components(self, x, sigma=None):
    return jnp.diag(1.0 + x ** 2)


class TestRicciEigendecomposition:
  """Tests that _eigh eigendecomposes the Ricci tensor correctly."""

  def test_constant_metric_zero_ricci_eigenvalues(self):
    """A constant metric has zero Ricci tensor, so eigenvalues are zero."""
    model = ConstantMetricModel([1.0, 4.0])
    x = jnp.array([0.0, 0.0])
    eigenvalues, _ = model._eigh(x)
    assert jnp.allclose(eigenvalues, jnp.zeros(2), atol=1e-10)

  def test_directions_are_g_orthonormal(self):
    """Directions U from _eigh satisfy U^T g U = I."""
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    _, directions = model._eigh(x)
    g = model.metric_components(x)
    assert jnp.allclose(directions.T @ g @ directions, jnp.eye(2), atol=1e-6)

  def test_directions_are_g_orthonormal_3d(self):
    """G-orthonormality in 3D."""
    class Varying3D(AbstractMetricLocalModel):
      def metric_components(self, x, sigma=None):
        d = 1.0 + x ** 2
        off = 0.1 * x[0] * x[1]
        return jnp.array([
          [d[0], off, 0.0],
          [off, d[1], 0.0],
          [0.0, 0.0, d[2]],
        ])

    model = Varying3D()
    x = jnp.array([0.5, -0.3, 1.0])
    _, directions = model._eigh(x)
    g = model.metric_components(x)
    assert jnp.allclose(directions.T @ g @ directions, jnp.eye(3), atol=1e-5)

  def test_eigenvalues_match_manual_ricci(self):
    """Eigenvalues from _eigh match manual Ricci computation."""
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])

    td = taylor_data(model.metric_components, x)
    g, dg, d2g = td.value, td.gradient, td.hessian
    g_inv = jnp.linalg.inv(g)
    gamma = christoffel(g_inv, dg)
    dgamma = christoffel_gradient(g_inv, dg, d2g)
    ric = ricci_from_christoffel(gamma, dgamma)
    E = orthonormal_frame(g)
    ric_rnc = E.T @ ric @ E
    expected_eigenvalues, _ = jnp.linalg.eigh(ric_rnc)

    actual_eigenvalues, _ = model._eigh(x)
    assert jnp.allclose(actual_eigenvalues, expected_eigenvalues, atol=1e-10)

  def test_ricci_eigenvalues_are_real(self):
    """Ricci eigenvalues from eigh of a symmetric matrix should be real."""
    model = VaryingMetricModel()
    x = jnp.array([0.5, -0.3])
    eigenvalues, _ = model._eigh(x)
    assert jnp.all(jnp.isfinite(eigenvalues))

  def test_flat_metric_log_s_is_neg_inf(self):
    """For a constant metric, Ricci = 0 and log(|0|) = -inf."""
    model = ConstantMetricModel([1.0, 4.0])
    x = jnp.array([0.0, 0.0])
    log_s = model.log_s(x)
    assert jnp.all(log_s == -jnp.inf)


class TestUMethods:
  """Tests for U_mvp, U_mvp_inv, get_U_matrix consistency."""

  def test_U_mvp_matches_matrix_multiply(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    v = jnp.array([1.0, 2.0])
    U = model.get_U_matrix(x)
    assert jnp.allclose(model.U_mvp(x, v), U @ v, atol=1e-6)

  def test_U_mvp_inv_matches_transpose_multiply(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    v = jnp.array([1.0, 2.0])
    U = model.get_U_matrix(x)
    assert jnp.allclose(model.U_mvp_inv(x, v), U.T @ v, atol=1e-6)

  def test_U_g_orthonormal_round_trip(self):
    """For g-orthonormal U, U^T g (U v) should recover v."""
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    v = jnp.array([0.7, -1.3])
    g = model.metric_components(x)
    Uv = model.U_mvp(x, v)
    recovered = model.U_mvp_inv(x, g @ Uv)
    assert jnp.allclose(recovered, v, atol=1e-6)


class TestScoreModel:
  """Tests for score_model correctness (unchanged by PCF refactor)."""

  def test_constant_metric_zero_score(self):
    model = ConstantMetricModel([1.0, 4.0])
    x = jnp.array([0.5, -0.3])
    score = model.score_model(x)
    assert jnp.allclose(score, jnp.zeros(2), atol=1e-6)

  def test_varying_metric_nonzero_score(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    score = model.score_model(x)
    assert not jnp.allclose(score, jnp.zeros(2), atol=1e-3)

  def test_score_matches_manual_grad(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])

    def log_vol(x):
      g = model.metric_components(x)
      return 0.5 * jnp.linalg.slogdet(g)[1]

    expected_score = jax.grad(log_vol)(x)
    actual_score = model.score_model(x)
    assert jnp.allclose(actual_score, expected_score, atol=1e-6)

  def test_score_matches_manual_grad_3d(self):
    class VaryingMetric3D(AbstractMetricLocalModel):
      def metric_components(self, x, sigma=None):
        d = 1.0 + x ** 2
        off = 0.1 * x[0] * x[1]
        return jnp.array([
          [d[0], off, 0.0],
          [off, d[1], 0.0],
          [0.0, 0.0, d[2]],
        ])

    model = VaryingMetric3D()
    x = jnp.array([0.5, -0.3, 1.0])

    def log_vol(x):
      g = model.metric_components(x)
      return 0.5 * jnp.linalg.slogdet(g)[1]

    expected_score = jax.grad(log_vol)(x)
    actual_score = model.score_model(x)
    assert jnp.allclose(actual_score, expected_score, atol=1e-6)


class TestGetPullbackMetric:
  """Test that get_pullback_metric returns the correct metric."""

  def test_pullback_metric_value_matches_g(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    metric = model.get_pullback_metric(x)
    g = model.metric_components(x)
    assert jnp.allclose(metric.components.value, g, atol=1e-6)

  def test_pullback_metric_gradient_shape(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    metric = model.get_pullback_metric(x)
    assert metric.components.gradient.shape == (2, 2, 2)

  def test_pullback_metric_hessian_shape(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    metric = model.get_pullback_metric(x)
    assert metric.components.hessian.shape == (2, 2, 2, 2)


class TestNeuralMetricModel:
  """Tests for NeuralMetricModel."""

  def test_output_is_symmetric(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    g = model.metric_components(x)
    assert jnp.allclose(g, g.T, atol=1e-7)

  def test_output_is_positive_definite(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=4)
    x = jnp.array([0.5, -0.3, 1.0, 0.2])
    g = model.metric_components(x)
    eigenvalues = jnp.linalg.eigvalsh(g)
    assert jnp.all(eigenvalues > 0)

  def test_output_shape(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    for dim in [2, 4, 8]:
      model = NeuralMetricModel(key=random.PRNGKey(0), dim=dim)
      x = jnp.zeros(dim)
      g = model.metric_components(x)
      assert g.shape == (dim, dim)

  def test_score_model_no_nan(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    score = model.score_model(x)
    assert not jnp.any(jnp.isnan(score))
    assert score.shape == (3,)

  def test_score_model_differentiable(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    hessian = jax.jacobian(model.score_model)(x)
    assert not jnp.any(jnp.isnan(hessian))
    assert hessian.shape == (2, 2)

  def test_ricci_directions_g_orthonormal(self):
    """Ricci eigenvectors from a neural metric are g-orthonormal."""
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(42), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    U = model.get_U_matrix(x)
    g = model.metric_components(x)
    assert jnp.allclose(U.T @ g @ U, jnp.eye(3), atol=1e-4)

  def test_euclidean_at_init(self):
    """At init_scale=0, metric should be identity for any input."""
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    for dim in [2, 3, 5]:
      model = NeuralMetricModel(key=random.PRNGKey(0), dim=dim)
      x = random.normal(random.PRNGKey(1), (dim,))
      g = model.metric_components(x)
      assert jnp.allclose(g, jnp.eye(dim), atol=1e-7), (
        f"Metric not identity at init for dim={dim}"
      )

  def test_score_zero_at_init(self):
    """At init_scale=0, the metric is constant I, so score should be zero."""
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    score = model.score_model(x)
    assert jnp.allclose(score, jnp.zeros(3), atol=1e-7)

  def test_different_inputs_different_outputs(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    g1 = model.metric_components(jnp.array([0.0, 0.0]))
    g2 = model.metric_components(jnp.array([1.0, 1.0]))
    assert not jnp.allclose(g1, g2, atol=1e-3)


class TestNoiseConditionalNeuralMetricModel:
  """Tests for NoiseConditionalNeuralMetricModel."""

  def test_output_is_symmetric(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    g = model.metric_components(x)
    assert jnp.allclose(g, g.T, atol=1e-7)

  def test_output_is_positive_definite(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=4)
    x = jnp.array([0.5, -0.3, 1.0, 0.2])
    g = model.metric_components(x)
    eigenvalues = jnp.linalg.eigvalsh(g)
    assert jnp.all(eigenvalues > 0)

  def test_spd_at_various_sigma(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    for sigma_val in [0.0, 0.01, 0.1, 0.5, 1.0]:
      sigma = jnp.array(sigma_val)
      g = model._metric_components_with_sigma(x, sigma)
      assert jnp.allclose(g, g.T, atol=1e-7), f"Not symmetric at sigma={sigma_val}"
      eigenvalues = jnp.linalg.eigvalsh(g)
      assert jnp.all(eigenvalues > 0), f"Not PD at sigma={sigma_val}"

  def test_output_shape(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    for dim in [2, 4, 8]:
      model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=dim)
      x = jnp.zeros(dim)
      g = model.metric_components(x)
      assert g.shape == (dim, dim)

  def test_metric_components_equals_sigma_zero(self):
    """metric_components(x) should equal _metric_components_with_sigma(x, 0)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    g_clean = model.metric_components(x)
    g_sigma0 = model._metric_components_with_sigma(x, jnp.array(0.0))
    assert jnp.allclose(g_clean, g_sigma0, atol=1e-10)

  def test_euclidean_at_init(self):
    """At init_scale=0, metric should be identity for any input and sigma."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    g = model.metric_components(x)
    assert jnp.allclose(g, jnp.eye(3), atol=1e-7)
    for sigma_val in [0.0, 0.1, 1.0]:
      g_sigma = model._metric_components_with_sigma(x, jnp.array(sigma_val))
      assert jnp.allclose(g_sigma, jnp.eye(3), atol=1e-7), (
        f"Metric not identity at init for sigma={sigma_val}"
      )

  def test_score_zero_at_init(self):
    """At init_scale=0, score should be zero everywhere."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    score = model.score_model(x)
    assert jnp.allclose(score, jnp.zeros(2), atol=1e-7)
    score_sigma = model.score_model(x, sigma=jnp.array(0.5))
    assert jnp.allclose(score_sigma, jnp.zeros(2), atol=1e-7)

  def test_different_sigma_different_metric(self):
    """Different sigma values should produce different metric tensors."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3, 1.0])
    g_low = model._metric_components_with_sigma(x, jnp.array(0.01))
    g_high = model._metric_components_with_sigma(x, jnp.array(1.0))
    assert not jnp.allclose(g_low, g_high, atol=1e-3)

  def test_score_model_no_nan(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    score = model.score_model(x)
    assert not jnp.any(jnp.isnan(score))
    assert score.shape == (3,)

  def test_score_model_with_sigma_no_nan(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    for sigma_val in [0.01, 0.1, 1.0]:
      sigma = jnp.array(sigma_val)
      score = model.score_model(x, sigma=sigma)
      assert not jnp.any(jnp.isnan(score)), f"NaN at sigma={sigma_val}"
      assert score.shape == (3,)

  def test_score_model_none_sigma_matches_zero(self):
    """score_model(x) should match score_model(x, sigma=0)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    score_none = model.score_model(x)
    score_zero = model.score_model(x, sigma=jnp.array(0.0))
    assert jnp.allclose(score_none, score_zero, atol=1e-10)

  def test_score_model_different_sigma_different_scores(self):
    """Different sigma should produce different scores."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    score_low = model.score_model(x, sigma=jnp.array(0.01))
    score_high = model.score_model(x, sigma=jnp.array(1.0))
    assert not jnp.allclose(score_low, score_high, atol=1e-3)

  def test_ncsm_gradient_no_nan(self):
    """Gradient of NCSM loss through the model should be finite."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    from local_disentanglement.losses.dsm import ncsm_loss

    model = NoiseConditionalNeuralMetricModel(
      key=random.PRNGKey(0), dim=2, working_size=32, hidden_size=32, n_blocks=2,
    )
    x = jnp.array([0.5, -0.3])
    sigma = jnp.array(0.1)
    key = random.PRNGKey(1)

    def loss_fn(model):
      return ncsm_loss(x, model, sigma, key)

    grads = eqx.filter_grad(loss_fn)(model)
    leaves = jax.tree_util.tree_leaves(eqx.filter(grads, eqx.is_inexact_array))
    for leaf in leaves:
      assert not jnp.any(jnp.isnan(leaf)), "NaN found in NCSM gradients"

  def test_noise_levels_stored(self):
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    from local_disentanglement.models.neural import IMANoiseScheduleConfig
    config = IMANoiseScheduleConfig(sigma_min=0.01, sigma_max=1.0, num_noise_levels=5)
    model = NoiseConditionalNeuralMetricModel(
      key=random.PRNGKey(0), dim=2, noise_config=config,
    )
    assert len(model.get_noise_levels()) == 5
    assert model.get_noise_levels()[0] == pytest.approx(0.01, rel=1e-4)

  def test_score_matches_manual_grad_with_sigma(self):
    """score_model(x, sigma) should equal grad 0.5 log det g(x, sigma)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    sigma = jnp.array(0.5)

    def log_vol(x):
      g = model._metric_components_with_sigma(x, sigma)
      return 0.5 * jnp.linalg.slogdet(g)[1]

    expected_score = jax.grad(log_vol)(x)
    actual_score = model.score_model(x, sigma=sigma)
    assert jnp.allclose(actual_score, expected_score, atol=1e-6)


class TestSigmaAwareMetricComponents:
  """Tests that metric_components(x, sigma) properly forwards sigma."""

  def test_metric_components_with_sigma_matches_internal(self):
    """metric_components(x, sigma) should match _metric_components_with_sigma."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=3)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3, 1.0])
    for sigma_val in [0.0, 0.01, 0.1, 0.5, 1.0]:
      sigma = jnp.array(sigma_val)
      g_public = model.metric_components(x, sigma=sigma)
      g_internal = model._metric_components_with_sigma(x, sigma)
      assert jnp.allclose(g_public, g_internal, atol=1e-10), (
        f"metric_components mismatch at sigma={sigma_val}"
      )

  def test_metric_components_none_sigma_is_clean(self):
    """metric_components(x) and metric_components(x, sigma=None) are sigma=0."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    g_none = model.metric_components(x)
    g_zero = model.metric_components(x, sigma=jnp.array(0.0))
    assert jnp.allclose(g_none, g_zero, atol=1e-10)

  def test_eigh_respects_sigma(self):
    """_eigh(x, sigma) should differ from _eigh(x) for nonzero sigma."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    eig_clean, _ = model._eigh(x)
    eig_noisy, _ = model._eigh(x, sigma=jnp.array(0.5))
    assert not jnp.allclose(eig_clean, eig_noisy, atol=1e-3), (
      "sigma should change the Ricci eigenvalues"
    )

  def test_get_U_matrix_respects_sigma(self):
    """get_U_matrix(x, sigma) should differ from get_U_matrix(x) for nonzero sigma."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    U_clean = model.get_U_matrix(x)
    U_noisy = model.get_U_matrix(x, sigma=jnp.array(0.5))
    assert not jnp.allclose(U_clean, U_noisy, atol=1e-3), (
      "sigma should change the principal directions"
    )

  def test_pullback_metric_respects_sigma(self):
    """get_pullback_metric(x, sigma) value should match metric_components(x, sigma)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    sigma = jnp.array(0.3)
    metric = model.get_pullback_metric(x, sigma=sigma)
    g = model.metric_components(x, sigma=sigma)
    assert jnp.allclose(metric.components.value, g, atol=1e-6)

  def test_non_conditional_ignores_sigma(self):
    """NeuralMetricModel should return the same metric regardless of sigma."""
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    g_none = model.metric_components(x)
    g_sigma = model.metric_components(x, sigma=jnp.array(0.5))
    assert jnp.allclose(g_none, g_sigma, atol=1e-10)


class TestSigmaAwarePCFLosses:
  """Tests that PCF losses properly use sigma."""

  def test_regularization_loss_with_sigma(self):
    """regularization_loss with sigma should use g(x, sigma)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    loss_clean = regularization_loss(x, model)
    loss_noisy = regularization_loss(x, model, sigma=jnp.array(0.5))
    assert not jnp.allclose(loss_clean, loss_noisy, atol=1e-3), (
      "sigma should change the regularization loss"
    )

  def test_regularization_loss_sigma_none_backward_compat(self):
    """regularization_loss(x, model) should be the same as sigma=None."""
    model = VaryingMetricModel()
    x = jnp.array([0.5, -0.3])
    loss_default = regularization_loss(x, model)
    loss_none = regularization_loss(x, model, sigma=None)
    assert jnp.allclose(loss_default, loss_none, atol=1e-10)

  def test_disentanglement_loss_with_sigma(self):
    """disentanglement_radius_loss with sigma should use g(x, sigma)."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x = jnp.array([0.5, -0.3])
    loss_clean = disentanglement_radius_loss(x, model)
    loss_noisy = disentanglement_radius_loss(x, model, sigma=jnp.array(0.5))
    assert not jnp.allclose(loss_clean, loss_noisy, atol=1e-3), (
      "sigma should change the disentanglement radius loss"
    )

  def test_regularization_batch_with_sigma(self):
    """Batch regularization loss with sigma should match manual mean."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2], [-0.5, 0.8]])
    sigma = jnp.array(0.3)
    batch_loss = regularization_loss_batch(x_batch, model, sigma=sigma)
    manual_losses = jnp.array([
      regularization_loss(x_batch[i], model, sigma=sigma)
      for i in range(3)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)

  def test_disentanglement_batch_with_sigma(self):
    """Batch disentanglement loss with sigma should match manual mean."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    model = NoiseConditionalNeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2]])
    sigma = jnp.array(0.3)
    batch_loss = disentanglement_radius_loss_batch(x_batch, model, sigma=sigma)
    manual_losses = jnp.array([
      disentanglement_radius_loss(x_batch[i], model, sigma=sigma)
      for i in range(2)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)


class TestExpLogRoundTrip:
  """Tests for exponential/logarithmic map round-trip."""

  def test_log_then_exp_round_trip_no_rotation(self):
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    x = jnp.array([0.6, 0.35])
    s = model.logarithmic_map(x_star, x, method="taylor")
    x_rec = model.exponential_map(x_star, s, method="taylor")
    assert jnp.allclose(x_rec, x, atol=1e-3)

  def test_exp_then_log_round_trip_no_rotation(self):
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    delta = jnp.array([0.05, -0.02])
    x = model.exponential_map(x_star, delta, method="taylor")
    delta_rec = model.logarithmic_map(x_star, x, method="taylor")
    assert jnp.allclose(delta_rec, delta, atol=1e-3)

  def test_log_then_exp_with_frame_rotation(self):
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    x = jnp.array([0.55, 0.32])

    theta = 0.7
    Q = jnp.array([
      [jnp.cos(theta), -jnp.sin(theta)],
      [jnp.sin(theta), jnp.cos(theta)],
    ])

    s = model.logarithmic_map(x_star, x, method="taylor", frame_rotation=Q)
    x_rec = model.exponential_map(x_star, s, method="taylor", frame_rotation=Q)
    assert jnp.allclose(x_rec, x, atol=1e-3)

  def test_rotation_changes_rnc_coords(self):
    """Different frame_rotation should give different RNC coordinates."""
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    x = jnp.array([0.55, 0.32])

    s_no_rot = model.logarithmic_map(x_star, x, method="taylor")

    theta = 0.7
    Q = jnp.array([
      [jnp.cos(theta), -jnp.sin(theta)],
      [jnp.sin(theta), jnp.cos(theta)],
    ])
    s_with_rot = model.logarithmic_map(
      x_star, x, method="taylor", frame_rotation=Q
    )

    assert not jnp.allclose(s_no_rot, s_with_rot, atol=1e-4)

  def test_identity_at_base_point(self):
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    s = model.logarithmic_map(x_star, x_star, method="taylor")
    assert jnp.allclose(s, jnp.zeros(2), atol=1e-6)

  def test_exp_of_zero_is_base_point(self):
    model = VaryingMetricModel()
    x_star = jnp.array([0.5, 0.3])
    x_rec = model.exponential_map(x_star, jnp.zeros(2), method="taylor")
    assert jnp.allclose(x_rec, x_star, atol=1e-6)


class TestDensityLoss:
  """Tests for the density score matching loss."""

  def test_loss_finite_positive(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    score_fn = lambda x: -x
    loss = density_loss(x, model, score_fn)
    assert jnp.isfinite(loss)
    assert loss >= 0.0

  def test_loss_zero_when_scores_match(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    perfect_score_fn = lambda x: model.score_model(x)
    loss = density_loss(x, model, perfect_score_fn)
    assert jnp.allclose(loss, 0.0, atol=1e-10)

  def test_loss_decreases_toward_correct_score(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    true_score = model.score_model(x)

    close_score_fn = lambda x: true_score + 0.01 * jnp.ones(2)
    far_score_fn = lambda x: true_score + 10.0 * jnp.ones(2)

    loss_close = density_loss(x, model, close_score_fn)
    loss_far = density_loss(x, model, far_score_fn)
    assert loss_close < loss_far

  def test_batch_matches_manual_mean(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2], [-0.5, 0.8]])
    score_fn = lambda x: -x

    batch_loss = density_loss_batch(x_batch, model, score_fn)
    manual_losses = jnp.array([
      density_loss(x_batch[i], model, score_fn)
      for i in range(3)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)

  def test_gradient_no_nan(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x = jnp.array([0.5, -0.3])
    score_fn = lambda x: -x

    def loss_fn(model):
      return density_loss(x, model, score_fn)

    grads = eqx.filter_grad(loss_fn)(model)
    leaves = jax.tree_util.tree_leaves(eqx.filter(grads, eqx.is_inexact_array))
    for leaf in leaves:
      assert not jnp.any(jnp.isnan(leaf)), "NaN found in gradients"


class TestRegularizationLoss:
  """Tests for the metric regularity loss."""

  def test_zero_when_identity(self):
    model = IdentityMetricModel(dim=3)
    x = jnp.array([0.5, -0.3, 1.0])
    loss = regularization_loss(x, model)
    assert jnp.allclose(loss, 0.0, atol=1e-10)

  def test_positive_when_not_identity(self):
    model = ConstantMetricModel([2.0, 3.0])
    x = jnp.zeros(2)
    loss = regularization_loss(x, model)
    expected = jnp.sum((jnp.diag(jnp.array([2.0, 3.0])) - jnp.eye(2)) ** 2)
    assert jnp.allclose(loss, expected, atol=1e-10)

  def test_batch_matches_manual_mean(self):
    model = VaryingMetricModel()
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2], [-0.5, 0.8]])
    batch_loss = regularization_loss_batch(x_batch, model)
    manual_losses = jnp.array([
      regularization_loss(x_batch[i], model)
      for i in range(3)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)


class TestDisentanglementRadiusLoss:
  """Tests for the disentanglement radius loss."""

  def test_zero_for_constant_metric(self):
    """A constant metric has constant Ric = 0, so grad Ric = 0."""
    model = ConstantMetricModel([1.0, 4.0])
    x = jnp.array([0.0, 0.0])
    loss = disentanglement_radius_loss(x, model)
    assert jnp.allclose(loss, 0.0, atol=1e-10)

  def test_finite_for_varying_metric(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    loss = disentanglement_radius_loss(x, model)
    assert jnp.isfinite(loss)
    assert loss >= 0.0

  def test_batch_matches_manual_mean(self):
    model = VaryingMetricModel()
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2]])
    batch_loss = disentanglement_radius_loss_batch(x_batch, model)
    manual_losses = jnp.array([
      disentanglement_radius_loss(x_batch[i], model)
      for i in range(2)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)


class TestPCFLoss:
  """Tests for the combined PCF loss."""

  def test_combined_equals_sum_of_parts(self):
    model = VaryingMetricModel()
    x = jnp.array([1.0, 2.0])
    score_fn = lambda x: -x
    c1, c2 = 0.5, 0.1

    combined = pcf_loss(x, model, score_fn, c1=c1, c2=c2)
    l_d = density_loss(x, model, score_fn)
    l_r = regularization_loss(x, model)
    l_dr = disentanglement_radius_loss(x, model)
    expected = l_d + c1 * l_r + c2 * l_dr
    assert jnp.allclose(combined, expected, atol=1e-10)

  def test_batch_matches_manual_mean(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    x_batch = jnp.array([[0.5, -0.3], [1.0, 0.2]])
    score_fn = lambda x: -x

    batch_loss = pcf_loss_batch(x_batch, model, score_fn, c1=0.5, c2=0.1)
    manual_losses = jnp.array([
      pcf_loss(x_batch[i], model, score_fn, c1=0.5, c2=0.1)
      for i in range(2)
    ])
    assert jnp.allclose(batch_loss, jnp.mean(manual_losses), atol=1e-6)


class TestIntegrationWithConformalFrame:
  """Tests for compatibility with conformal_frame evaluation."""

  def test_evaluate_model_runs(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    from local_disentanglement.experiments.conformal_frame import evaluate_model

    model = NeuralMetricModel(key=random.PRNGKey(0), dim=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    obs = random.normal(random.PRNGKey(1), (50, 2))
    score_fn = lambda x: -x
    eval_key = random.PRNGKey(2)
    metrics = evaluate_model(model, obs, score_fn, eval_key, subsample=10)
    assert "eval/fisher_divergence" in metrics
    assert all(np.isfinite(v) for v in metrics.values())

  def test_funnel_alignment_runs(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel
    from local_disentanglement.experiments.conformal_frame import (
      evaluate_funnel_alignment,
    )
    from local_disentanglement.energies import RotatedFunnel

    dim = 4
    energy = RotatedFunnel(ndim=dim, sigma_0=3.0, seed=0)
    model = NeuralMetricModel(key=random.PRNGKey(0), dim=dim)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    obs = energy.sample(random.PRNGKey(1), 100)
    eval_key = random.PRNGKey(2)
    metrics = evaluate_funnel_alignment(
      model, energy, obs, eval_key, subsample=20
    )
    assert "eval/funnel_alignment_mean" in metrics


class TestTrainingSmokeTest:
  """Short training run to verify the full PCF pipeline."""

  def test_density_loss_decreases_over_training(self):
    from local_disentanglement.models.neural_metric import NeuralMetricModel

    dim = 2
    key = random.PRNGKey(42)
    key, data_key, model_key, train_key = random.split(key, 4)

    obs = random.normal(data_key, (2000, dim))
    score_fn = lambda x: -x

    model = NeuralMetricModel(key=model_key, dim=dim, working_size=32, hidden_size=32, n_blocks=2)
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(0.1))
    optimizer = optax.adam(3e-3)
    opt_state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))

    @eqx.filter_jit
    def step(model, opt_state, batch):
      def loss_fn(model):
        return density_loss_batch(batch, model, score_fn)
      loss, grads = eqx.filter_value_and_grad(loss_fn)(model)
      updates, new_opt_state = optimizer.update(grads, opt_state, model)
      new_model = eqx.apply_updates(model, updates)
      return new_model, new_opt_state, loss

    n_steps = 200
    batch_size = 64
    losses = []

    for i in range(n_steps):
      train_key, batch_key = random.split(train_key)
      idx = random.choice(batch_key, obs.shape[0], shape=(batch_size,))
      batch = obs[idx]
      model, opt_state, loss = step(model, opt_state, batch)
      losses.append(float(loss))

    early_loss = np.mean(losses[:20])
    late_loss = np.mean(losses[-20:])
    assert late_loss < early_loss, (
      f"Loss did not decrease: early={early_loss:.4f}, late={late_loss:.4f}"
    )

  def test_ncsm_loss_decreases_over_training(self):
    """DSM training with NoiseConditionalNeuralMetricModel should reduce loss."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    from local_disentanglement.models.neural import IMANoiseScheduleConfig
    from local_disentanglement.losses.dsm import ncsm_loss_batch_continuous

    dim = 2
    key = random.PRNGKey(42)
    key, data_key, model_key, train_key = random.split(key, 4)

    obs = random.normal(data_key, (2000, dim))

    noise_config = IMANoiseScheduleConfig(
      sigma_min=0.01, sigma_max=1.0, num_noise_levels=5,
    )
    model = NoiseConditionalNeuralMetricModel(
      key=model_key, dim=dim, noise_config=noise_config,
      working_size=32, hidden_size=32, n_blocks=2,
    )
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(0.1))
    optimizer = optax.adam(3e-3)
    opt_state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))

    sigma_min, sigma_max = 0.01, 1.0

    @eqx.filter_jit
    def step(model, opt_state, batch, step_key):
      def loss_fn(model):
        return ncsm_loss_batch_continuous(
          batch, model, sigma_min, sigma_max, step_key,
        )
      loss, grads = eqx.filter_value_and_grad(loss_fn)(model)
      updates, new_opt_state = optimizer.update(grads, opt_state, model)
      new_model = eqx.apply_updates(model, updates)
      return new_model, new_opt_state, loss

    n_steps = 200
    batch_size = 64
    losses = []

    for i in range(n_steps):
      train_key, batch_key, step_key = random.split(train_key, 3)
      idx = random.choice(batch_key, obs.shape[0], shape=(batch_size,))
      batch = obs[idx]
      model, opt_state, loss = step(model, opt_state, batch, step_key)
      losses.append(float(loss))

    early_loss = np.mean(losses[:20])
    late_loss = np.mean(losses[-20:])
    assert late_loss < early_loss, (
      f"NCSM loss did not decrease: early={early_loss:.4f}, late={late_loss:.4f}"
    )

  def test_evaluate_model_runs_with_noise_conditional(self):
    """evaluate_model should work with NoiseConditionalNeuralMetricModel."""
    from local_disentanglement.models.neural_metric import NoiseConditionalNeuralMetricModel
    from local_disentanglement.experiments.conformal_frame import evaluate_model

    model = NoiseConditionalNeuralMetricModel(
      key=random.PRNGKey(0), dim=2,
    )
    model = eqx.tree_at(lambda m: m.init_scale, model, jnp.array(1.0))
    obs = random.normal(random.PRNGKey(1), (50, 2))
    score_fn = lambda x: -x
    eval_key = random.PRNGKey(2)
    eval_sigma = model.get_noise_levels()[1]
    metrics = evaluate_model(
      model, obs, score_fn, eval_key, subsample=10, sigma=eval_sigma,
    )
    assert "eval/fisher_divergence" in metrics
    assert all(np.isfinite(v) for v in metrics.values())
