"""Tests for the Monge patch decoder in analytic_verification.py."""
import jax
import jax.numpy as jnp
import numpy as np
import pytest

from local_coordinates.metric import pullback_metric, get_euclidean_metric
from local_geometry.experiments.synthetic_models.decoder import (
  MONGE_HEIGHT_FUNCTIONS,
  apply_decoder_offset,
  asymmetric_saddle,
  make_monge_decoder,
  monkey_saddle,
  twisted_paraboloid,
)
from local_geometry.independent_components import riemannian_ica
from local_geometry.paper.figures.analytic_verification import (
  log_pz,
  s_to_z,
  dz_ds,
  log_ps,
  make_coord_fns,
)

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


ALL_HEIGHT_FNS = list(MONGE_HEIGHT_FUNCTIONS.values())


@pytest.fixture(params=ALL_HEIGHT_FNS, ids=list(MONGE_HEIGHT_FUNCTIONS.keys()))
def monge_setup(request):
  """Build a monge decoder, metric, and Jacobians for each height function."""
  h_fn = request.param
  z0 = jnp.array([0.3, -0.2])
  f = make_monge_decoder(h_fn)
  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()
  return f, h_fn, z0, metric, j_z_to_s, j_s_to_z


def test_monge_decoder_output_shape(monge_setup):
  """Monge decoder should map R^2 to R^3."""
  f, _, z0, *_ = monge_setup
  y = f(z0)
  assert y.shape == (3,)


def test_monge_decoder_identity_coordinates(monge_setup):
  """First two output coordinates of a Monge decoder must equal the input."""
  f, _, z0, *_ = monge_setup
  y = f(z0)
  assert jnp.allclose(y[:2], z0, atol=1e-14)


def test_monge_decoder_third_coord_equals_height(monge_setup):
  """Third output coordinate must equal h(z)."""
  f, h_fn, z0, *_ = monge_setup
  y = f(z0)
  assert jnp.allclose(y[2], h_fn(z0), atol=1e-14)


def test_monge_metric_structure(monge_setup):
  """The pullback metric of a Monge patch should be g = I + grad(h) grad(h)^T."""
  f, h_fn, z0, metric, *_ = monge_setup
  g = metric.components.value

  grad_h = jax.grad(h_fn)(z0)
  g_expected = jnp.eye(2) + jnp.outer(grad_h, grad_h)
  assert jnp.allclose(g, g_expected, atol=1e-10)


def test_monge_metric_positive_definite(monge_setup):
  """The pullback metric must be positive definite."""
  _, _, _, metric, *_ = monge_setup
  g = metric.components.value
  eigvals = jnp.linalg.eigvalsh(g)
  assert jnp.all(eigvals > 0)


def test_s_to_z_origin_is_z0(monge_setup):
  """Mapping s=0 through the RNC Jacobian should return z0."""
  _, _, z0, _, _, j_s_to_z = monge_setup
  z_mapped = s_to_z(jnp.zeros(2), z0, j_s_to_z)
  assert jnp.allclose(z_mapped, z0, atol=1e-12)


def test_round_trip_near_origin(monge_setup):
  """Mapping z0 -> s -> z should approximately recover small s."""
  _, _, z0, _, j_z_to_s, j_s_to_z = monge_setup
  small_s = jnp.array([0.05, -0.03])
  z = s_to_z(small_s, z0, j_s_to_z)
  dz = z - z0

  s_recovered = jnp.einsum("ia,a->i", j_z_to_s.value, dz)
  s_recovered += 0.5 * jnp.einsum("iab,a,b->i", j_z_to_s.gradient, dz, dz)
  s_recovered += (1.0 / 6.0) * jnp.einsum(
    "iabc,a,b,c->i", j_z_to_s.hessian, dz, dz, dz
  )
  assert jnp.allclose(s_recovered, small_s, atol=1e-6)


def test_disentangled_hessian_is_diagonal(monge_setup):
  """The Hessian of log p_s at the origin should be diagonal."""
  f, _, z0, _, _, j_s_to_z = monge_setup

  s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_s_to_z)

  def _log_ps(s):
    return log_ps(s, s_to_z_fn, dz_ds_fn)

  hessian = jax.hessian(_log_ps)(jnp.zeros(2))
  off_diag = hessian - jnp.diag(jnp.diag(hessian))
  assert jnp.allclose(off_diag, 0.0, atol=1e-8)


def test_coordinate_fns_pipeline(monge_setup):
  """The full coordinate function pipeline should produce finite outputs."""
  f, _, z0, _, j_z_to_s, j_s_to_z = monge_setup
  s_to_z_fn, dz_ds_fn = make_coord_fns(
    "taylor", z0, j_s_to_z,
  )
  s_test = jnp.array([0.2, -0.1])
  z = s_to_z_fn(s_test)
  jac = dz_ds_fn(s_test)
  assert jnp.isfinite(z).all()
  assert jnp.isfinite(jac).all()
  assert z.shape == (2,)
  assert jac.shape == (2, 2)


def test_monge_decoder_at_various_points():
  """The Monge decoder should work at several different points."""
  f = make_monge_decoder(asymmetric_saddle)
  test_points = jnp.array([
    [0.0, 0.0],
    [1.0, 0.0],
    [0.0, 1.0],
    [-0.5, 0.5],
    [0.3, -0.7],
  ])
  for z in test_points:
    y = f(z)
    assert y.shape == (3,)
    assert jnp.allclose(y[:2], z, atol=1e-14)
    assert jnp.isfinite(y[2])


def test_height_fn_at_origin():
  """All height functions should be zero at the origin (no constant term)."""
  z0 = jnp.zeros(2)
  for name, h_fn in MONGE_HEIGHT_FUNCTIONS.items():
    val = h_fn(z0)
    assert jnp.allclose(val, 0.0, atol=1e-14), (
      f"Height function {name} is nonzero at origin: {val}"
    )


def test_height_fn_gradient_at_origin():
  """All height functions should have zero gradient at the origin
  (no linear term), so the metric is identity there."""
  z0 = jnp.zeros(2)
  for name, h_fn in MONGE_HEIGHT_FUNCTIONS.items():
    grad_h = jax.grad(h_fn)(z0)
    assert jnp.allclose(grad_h, 0.0, atol=1e-14), (
      f"Height function {name} has nonzero gradient at origin: {grad_h}"
    )


def test_log_det_g_monge():
  """Log det g for a Monge patch should equal log(1 + ||grad h||^2)."""
  f = make_monge_decoder(asymmetric_saddle)
  z_test = jnp.array([0.5, -0.3])

  J_f = jax.jacfwd(f)(z_test)
  g = J_f.T @ J_f
  log_det_g = jnp.linalg.slogdet(g)[1]

  grad_h = jax.grad(asymmetric_saddle)(z_test)
  expected = jnp.log(1.0 + jnp.dot(grad_h, grad_h))
  assert jnp.allclose(log_det_g, expected, atol=1e-10)


# ---------------------------------------------------------------------------
# Offset tests
# ---------------------------------------------------------------------------

def test_offset_shifts_decoder_evaluation():
  """With offset c, apply_decoder_offset should evaluate f at z - c."""
  offset = jnp.array([1.0, -0.5])
  f = apply_decoder_offset(make_monge_decoder(asymmetric_saddle), offset)
  z_test = jnp.array([0.3, 0.4])

  y = f(z_test)
  z_shifted = z_test - offset
  expected_height = asymmetric_saddle(z_shifted)
  assert jnp.allclose(y[2], expected_height, atol=1e-14)
  assert jnp.allclose(y[:2], z_shifted, atol=1e-14)


def test_offset_equals_evaluation_at_shifted_point():
  """apply_decoder_offset(f, c)(z) should equal f(z - c) for any decoder."""
  offset = jnp.array([2.0, 3.0])
  f_base = make_monge_decoder(twisted_paraboloid)
  f_offset = apply_decoder_offset(f_base, offset)
  z_test = jnp.array([-0.5, 0.7])

  assert jnp.allclose(f_offset(z_test), f_base(z_test - offset), atol=1e-14)


def test_zero_offset_matches_no_offset():
  """Passing offset=zeros(2) should produce the same decoder as no offset."""
  f_no_offset = make_monge_decoder(asymmetric_saddle)
  f_zero_offset = apply_decoder_offset(f_no_offset, jnp.zeros(2))
  z_test = jnp.array([0.5, -0.3])

  assert jnp.allclose(f_no_offset(z_test), f_zero_offset(z_test), atol=1e-14)


def test_offset_produces_nontrivial_metric_at_origin():
  """With a non-zero offset, the metric at z=0 should differ from the
  identity because grad h is nonzero at z=0 when the polynomial center
  is shifted away."""
  offset = jnp.array([1.0, 0.5])
  f = apply_decoder_offset(make_monge_decoder(asymmetric_saddle), offset)
  z0 = jnp.zeros(2)

  eucl = get_euclidean_metric(f(z0))
  metric = pullback_metric(z0, f, eucl)
  g = metric.components.value

  assert not jnp.allclose(g, jnp.eye(2), atol=1e-6), (
    "Metric at origin with offset should not be identity"
  )


def test_offset_metric_equals_shifted_monge_formula():
  """The pullback metric with offset should still satisfy the Monge
  formula g = I + grad_z(h(z-c)) grad_z(h(z-c))^T."""
  offset = jnp.array([0.7, -0.3])
  f = apply_decoder_offset(make_monge_decoder(twisted_paraboloid), offset)
  z0 = jnp.array([0.2, 0.1])

  eucl = get_euclidean_metric(f(z0))
  metric = pullback_metric(z0, f, eucl)
  g = metric.components.value

  def h_shifted(z):
    return twisted_paraboloid(z - offset)

  grad_h = jax.grad(h_shifted)(z0)
  g_expected = jnp.eye(2) + jnp.outer(grad_h, grad_h)
  assert jnp.allclose(g, g_expected, atol=1e-10)


def test_offset_full_pipeline():
  """The full RNC pipeline should work with an offset decoder."""
  offset = jnp.array([0.5, -0.5])
  f = apply_decoder_offset(make_monge_decoder(asymmetric_saddle), offset)
  z0 = jnp.array([0.3, -0.2])

  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()

  s_to_z_fn, dz_ds_fn = make_coord_fns("taylor", z0, j_s_to_z)
  s_test = jnp.array([0.1, -0.1])
  z = s_to_z_fn(s_test)
  jac = dz_ds_fn(s_test)
  assert jnp.isfinite(z).all()
  assert jnp.isfinite(jac).all()
  assert z.shape == (2,)
  assert jac.shape == (2, 2)
