"""Tests for the coordinate mapping functions in analytic_verification.py.

Verifies that the Taylor, improved, and exact methods all produce consistent
results and that the factory function wires them correctly.
"""
import jax
import jax.numpy as jnp
import pytest

from local_coordinates.metric import pullback_metric, get_euclidean_metric
from local_geometry.independent_components import riemannian_ica
from local_geometry.paper.figures.analytic_verification import (
  make_decoder,
  log_pz,
  s_to_z,
  dz_ds,
  z_to_s_taylor,
  s_to_z_improved,
  s_to_z_exact,
  make_coord_fns,
)


@pytest.fixture
def setup():
  key = jax.random.PRNGKey(15)
  z0 = jnp.array([0.3, -0.2])
  f, _, _, _ = make_decoder(key, d=2, n=3, eps=0.6)
  eucl = get_euclidean_metric(f(z0))
  metric = pullback_metric(z0, f, eucl)
  j_z_to_s = riemannian_ica(metric, log_pz)
  j_s_to_z = j_z_to_s.get_inverse()

  def metric_fn(z):
    return pullback_metric(z, f, get_euclidean_metric(f(z)))

  return z0, f, j_s_to_z, j_z_to_s, metric_fn


def test_taylor_at_origin_is_base_point(setup):
  """s=0 should map back to z0 for all methods."""
  z0, _, j_s_to_z, j_z_to_s, _ = setup
  s_origin = jnp.zeros(2)
  z = s_to_z(s_origin, z0, j_s_to_z)
  assert jnp.allclose(z, z0, atol=1e-12)


def test_improved_at_origin_is_base_point(setup):
  z0, _, j_s_to_z, j_z_to_s, _ = setup
  s_origin = jnp.zeros(2)
  z = s_to_z_improved(s_origin, z0, j_s_to_z, j_z_to_s)
  assert jnp.allclose(z, z0, atol=1e-12)


def test_z_to_s_taylor_roundtrip(setup):
  """Applying z_to_s_taylor(s_to_z(s)) should approximately recover s."""
  z0, _, j_s_to_z, j_z_to_s, _ = setup
  s_test = jnp.array([0.1, 0.1])
  z = s_to_z(s_test, z0, j_s_to_z)
  s_recovered = z_to_s_taylor(z, z0, j_z_to_s)
  assert jnp.allclose(s_test, s_recovered, atol=1e-3)


def test_improved_closer_to_exact_than_taylor(setup):
  """The improved method should have smaller error relative to exact."""
  z0, _, j_s_to_z, j_z_to_s, metric_fn = setup
  s_test = jnp.array([0.3, 0.4])

  z_taylor = s_to_z(s_test, z0, j_s_to_z)
  z_improved = s_to_z_improved(s_test, z0, j_s_to_z, j_z_to_s)
  z_exact = s_to_z_exact(s_test, z0, j_s_to_z, metric_fn)

  err_taylor = jnp.linalg.norm(z_taylor - z_exact)
  err_improved = jnp.linalg.norm(z_improved - z_exact)
  assert err_improved < err_taylor


def test_make_coord_fns_taylor(setup):
  """Factory with 'taylor' should match the raw s_to_z and dz_ds."""
  z0, _, j_s_to_z, j_z_to_s, _ = setup
  s_test = jnp.array([0.2, -0.1])

  s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_s_to_z)
  z_factory = s_to_z_fn(s_test)
  jac_factory = dz_ds_fn(s_test)

  z_direct = s_to_z(s_test, z0, j_s_to_z)
  jac_direct = dz_ds(s_test, j_s_to_z)

  assert jnp.allclose(z_factory, z_direct, atol=1e-14)
  assert jnp.allclose(jac_factory, jac_direct, atol=1e-14)


def test_make_coord_fns_improved(setup):
  """Factory with 'improved' should match the raw s_to_z_improved."""
  z0, _, j_s_to_z, j_z_to_s, _ = setup
  s_test = jnp.array([0.2, -0.1])

  s_to_z_fn, dz_ds_fn = make_coord_fns(
    "improved", z0, j_s_to_z, j_z_to_s=j_z_to_s,
  )
  z_factory = s_to_z_fn(s_test)
  z_direct = s_to_z_improved(s_test, z0, j_s_to_z, j_z_to_s)
  assert jnp.allclose(z_factory, z_direct, atol=1e-14)

  jac = dz_ds_fn(s_test)
  assert jac.shape == (2, 2)
  assert jnp.isfinite(jac).all()


def test_make_coord_fns_exact(setup):
  """Factory with 'exact' should match the raw s_to_z_exact."""
  z0, _, j_s_to_z, j_z_to_s, metric_fn = setup
  s_test = jnp.array([0.2, -0.1])

  s_to_z_fn, dz_ds_fn = make_coord_fns(
    "exact", z0, j_s_to_z, metric_fn=metric_fn,
  )
  z_factory = s_to_z_fn(s_test)
  z_direct = s_to_z_exact(s_test, z0, j_s_to_z, metric_fn)
  assert jnp.allclose(z_factory, z_direct, atol=1e-14)

  jac = dz_ds_fn(s_test)
  assert jac.shape == (2, 2)
  assert jnp.isfinite(jac).all()


def test_make_coord_fns_invalid_method(setup):
  z0, _, j_s_to_z, _, _ = setup
  with pytest.raises(ValueError, match="Unknown method"):
    make_coord_fns("bogus", z0, j_s_to_z)


def test_make_coord_fns_improved_requires_j_z_to_s(setup):
  z0, _, j_s_to_z, _, _ = setup
  with pytest.raises(ValueError, match="j_z_to_s"):
    make_coord_fns("improved", z0, j_s_to_z)


def test_make_coord_fns_exact_requires_metric_fn(setup):
  z0, _, j_s_to_z, _, _ = setup
  with pytest.raises(ValueError, match="metric_fn"):
    make_coord_fns("exact", z0, j_s_to_z)
