"""Tests for make_fixed_decoder in analytic_verification.py.

Verifies that the hardcoded decoder parameters satisfy the structural
invariants (orthonormality, symmetry, unit norm) and that the resulting
surface has nonzero Riemann curvature.
"""
import jax
import jax.numpy as jnp
import pytest

from local_coordinates.metric import pullback_metric, get_euclidean_metric
from local_geometry.experiments.synthetic_models.decoder import make_fixed_decoder

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


@pytest.fixture
def decoder():
  return make_fixed_decoder(d=2, n=3, eps=1.0)


def test_A_orthonormal(decoder):
  """Columns of A should be orthonormal."""
  _, A, _, _ = decoder
  gram = A.T @ A
  assert jnp.allclose(gram, jnp.eye(2), atol=1e-12)


def test_B_symmetric(decoder):
  """B should be fully symmetric in its last three indices."""
  _, _, B, _ = decoder
  perms = [
    (0, 1, 2, 3), (0, 1, 3, 2), (0, 2, 1, 3),
    (0, 2, 3, 1), (0, 3, 1, 2), (0, 3, 2, 1),
  ]
  for p in perms:
    assert jnp.allclose(B, B.transpose(p), atol=1e-14)


def test_B_unit_frobenius_norm(decoder):
  """B should have unit Frobenius norm."""
  _, _, B, _ = decoder
  assert jnp.allclose(jnp.linalg.norm(B), 1.0, atol=1e-12)


def test_deterministic():
  """Two calls should produce identical outputs."""
  f1, A1, B1, c1 = make_fixed_decoder(d=2, n=3, eps=0.5)
  f2, A2, B2, c2 = make_fixed_decoder(d=2, n=3, eps=0.5)
  assert jnp.allclose(A1, A2, atol=0.0)
  assert jnp.allclose(B1, B2, atol=0.0)
  assert jnp.allclose(c1, c2, atol=0.0)

  z = jnp.array([0.3, -0.2])
  assert jnp.allclose(f1(z), f2(z), atol=0.0)


def test_nonzero_curvature(decoder):
  """The surface should have nonzero Riemann curvature at the base point.

  For a 2D manifold the full curvature is captured by the scalar curvature.
  We check this is nonzero by computing it via finite differences of the
  metric determinant (the Christoffel symbol route).
  """
  f, _, _, _ = decoder
  z0 = jnp.array([0.3, -0.2])
  eucl = get_euclidean_metric(f(z0))
  metric = pullback_metric(z0, f, eucl)

  g = metric.components.value
  dg = metric.components.gradient
  d2g = metric.components.hessian

  g_inv = jnp.linalg.inv(g)
  christoffel = 0.5 * jnp.einsum(
    "il,ljk->ijk", g_inv,
    dg[:, :, :] + jnp.einsum("jki->ijk", dg) - jnp.einsum("kji->ijk", dg),
  )

  R1_212 = (
    d2g[0, 1, 0, 1] - 0.5 * d2g[0, 0, 1, 1] - 0.5 * d2g[1, 1, 0, 0]
    + jnp.einsum("p,p->", g_inv[0, :],
                  jnp.einsum("pq,q->p", g[:, :],
                             christoffel[0, 1, :] * christoffel[1, 0, :] -
                             christoffel[0, 0, :] * christoffel[1, 1, :]))
  )
  assert jnp.abs(R1_212) > 1e-4, f"Expected nonzero curvature, got {R1_212}"


def test_invalid_dimensions():
  """make_fixed_decoder should reject d != 2 or n != 3."""
  with pytest.raises(ValueError, match="d=2, n=3"):
    make_fixed_decoder(d=3, n=3, eps=1.0)
  with pytest.raises(ValueError, match="d=2, n=3"):
    make_fixed_decoder(d=2, n=4, eps=1.0)
