import jax
import jax.numpy as jnp
import pytest

from local_geometry.experiments.synthetic_models.decoder import (
  ContextSpec,
  build_context,
  build_context_from_spec,
  build_decoder_from_spec,
  build_log_pz_from_spec,
  log_ps_via_inverse_map,
  make_fixed_decoder,
)


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


def _build_fixed_context(log_pz_fn=None):
  decoder, _, _, _ = make_fixed_decoder(d=2, n=3, eps=1.0)
  z0 = jnp.array([0.3, -0.2])
  return build_context(decoder, z0, method="taylor", log_pz_fn=log_pz_fn)


def test_build_context_roundtrip_is_local_identity():
  context = _build_fixed_context()
  s = jnp.array([0.02, -0.015])
  z = context.s_to_z_fn(s)
  s_recovered = context.z_to_s_fn(z)
  assert jnp.allclose(s_recovered, s, atol=1e-6)


def test_build_context_diagonalizes_log_ps_hessian():
  context = _build_fixed_context()

  def log_ps_fn(s):
    return log_ps_via_inverse_map(
      s,
      context.s_to_z_fn,
      context.z_to_s_fn,
      log_pz_fn=context.log_pz_fn,
    )

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


def test_build_context_uses_custom_prior():
  mu = jnp.array([0.4, -0.1])

  def shifted_log_pz(z):
    centered = z - mu
    return -0.5 * jnp.dot(centered, centered)

  default_context = _build_fixed_context()
  custom_context = _build_fixed_context(log_pz_fn=shifted_log_pz)
  s = jnp.array([0.05, -0.03])

  log_ps_default = log_ps_via_inverse_map(
    s,
    default_context.s_to_z_fn,
    default_context.z_to_s_fn,
    log_pz_fn=default_context.log_pz_fn,
  )
  log_ps_custom = log_ps_via_inverse_map(
    s,
    custom_context.s_to_z_fn,
    custom_context.z_to_s_fn,
    log_pz_fn=custom_context.log_pz_fn,
  )

  assert custom_context.log_pz_fn is shifted_log_pz
  assert not jnp.allclose(log_ps_default, log_ps_custom)


def test_build_decoder_from_spec_fixed_matches_manual_decoder():
  spec = ContextSpec(decoder_name="fixed", eps=1.0)
  decoder_from_spec = build_decoder_from_spec(spec)
  decoder_manual, _, _, _ = make_fixed_decoder(d=2, n=3, eps=1.0)
  z = jnp.array([0.1, -0.2])
  assert jnp.allclose(decoder_from_spec(z), decoder_manual(z), atol=1e-12)


def test_build_decoder_from_spec_polynomial_is_seeded():
  spec = ContextSpec(decoder_name="polynomial", seed=17, d=2, n=3, eps=0.7)
  decoder_a = build_decoder_from_spec(spec)
  decoder_b = build_decoder_from_spec(spec)
  z = jnp.array([0.2, -0.05])
  assert jnp.allclose(decoder_a(z), decoder_b(z), atol=1e-12)


def test_build_decoder_from_spec_named_monge_surface():
  spec = ContextSpec(decoder_name="monkey_saddle")
  decoder = build_decoder_from_spec(spec)
  z = jnp.array([0.2, -0.1])
  expected_height = 0.3 * (z[0] ** 3 - 3.0 * z[0] * z[1] ** 2)
  expected = jnp.array([z[0], z[1], expected_height])
  assert jnp.allclose(decoder(z), expected, atol=1e-12)


def test_build_context_from_spec_matches_manual_fixed_context():
  spec = ContextSpec(
    decoder_name="fixed",
    z0=(0.3, -0.2),
    method="taylor",
    eps=1.0,
  )
  context_from_spec = build_context_from_spec(spec)
  decoder_manual, _, _, _ = make_fixed_decoder(d=2, n=3, eps=1.0)
  context_manual = build_context(
    decoder_manual,
    jnp.array([0.3, -0.2]),
    method="taylor",
  )
  s = jnp.array([0.03, -0.02])
  assert jnp.allclose(
    context_from_spec.s_to_z_fn(s),
    context_manual.s_to_z_fn(s),
    atol=1e-8,
  )
  assert jnp.allclose(
    context_from_spec.frame_rotation,
    context_manual.frame_rotation,
    atol=1e-10,
  )


def test_build_log_pz_from_spec_returns_standard_normal():
  spec = ContextSpec(prior_name="standard_normal")
  log_pz_fn = build_log_pz_from_spec(spec)
  z = jnp.array([0.2, -0.1])
  assert jnp.allclose(log_pz_fn(z), -0.5 * jnp.dot(z, z), atol=1e-12)


def test_context_spec_rejects_unknown_decoder_name():
  with pytest.raises(ValueError, match="Unknown decoder_name"):
    ContextSpec(decoder_name="banana")


def test_context_spec_rejects_unknown_prior_name():
  with pytest.raises(ValueError, match="Unknown prior_name"):
    ContextSpec(prior_name="laplace")


def test_context_spec_requires_named_monge_height():
  with pytest.raises(ValueError, match="Unknown monge height name"):
    ContextSpec(decoder_name="monge")
