import jax
import jax.numpy as jnp

from local_geometry.experiments.sanity.von_mises import (
  ODEMethod,
  RefinedTaylorMethod,
  VonMisesFisher,
  build_chart_metric_fn,
  build_sphere_chart_context,
  build_tangent_frame,
  log_ps_via_inverse_map,
  make_stereographic_chart,
  make_tangent_graph_chart,
  sphere_exp_map,
  sphere_log_map,
  vmf_entanglement_eigenvalue,
)
from local_geometry.rica_jax.rnc_pushforward import s_to_z, z_to_s


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


def _normalize(x):
  return x / jnp.linalg.norm(x)


X0 = _normalize(jnp.array([0.35, -0.25, 0.9], dtype=jnp.float64))
MU = _normalize(jnp.array([-0.3, 0.6, 0.74], dtype=jnp.float64))
TANGENT_FRAME = build_tangent_frame(X0)


def test_vmf_samples_lie_on_unit_sphere():
  vmf = VonMisesFisher(mean_direction=MU, concentration=9.0)
  key = jax.random.PRNGKey(0)

  samples = vmf.sample(key, sample_shape=(256,))
  norms = jnp.linalg.norm(samples, axis=-1)

  assert samples.shape == (256, 3)
  assert jnp.allclose(norms, 1.0, atol=1e-6)


def test_vmf_log_prob_depends_only_on_mean_inner_product():
  vmf = VonMisesFisher(mean_direction=jnp.array([0.0, 0.0, 1.0]), concentration=5.0)
  target_dot = 0.4
  radial = jnp.sqrt(1.0 - target_dot ** 2)
  x1 = jnp.array([radial, 0.0, target_dot])
  x2 = jnp.array([0.0, radial, target_dot])

  assert jnp.allclose(vmf.log_prob(x1), vmf.log_prob(x2), atol=1e-10)


def test_sphere_charts_roundtrip():
  tangent_chart = make_tangent_graph_chart(X0, tangent_frame=TANGENT_FRAME)
  stereographic_chart = make_stereographic_chart(X0, tangent_frame=TANGENT_FRAME)
  s = jnp.array([0.16, -0.09], dtype=jnp.float64)
  x = sphere_exp_map(X0, TANGENT_FRAME @ s)

  for chart in (tangent_chart, stereographic_chart):
    z = chart.x_to_z_fn(x)
    x_recovered = chart.z_to_x_fn(z)
    assert jnp.allclose(x_recovered, x, atol=1e-8)


def test_chart_metric_at_origin_matches_known_scaling():
  tangent_chart = make_tangent_graph_chart(X0, tangent_frame=TANGENT_FRAME)
  stereographic_chart = make_stereographic_chart(X0, tangent_frame=TANGENT_FRAME)
  tangent_metric = build_chart_metric_fn(tangent_chart)
  stereographic_metric = build_chart_metric_fn(stereographic_chart)

  assert jnp.allclose(tangent_metric(tangent_chart.z0), jnp.eye(2), atol=1e-10)
  assert jnp.allclose(stereographic_metric(stereographic_chart.z0), 4.0 * jnp.eye(2), atol=1e-10)


def test_rnc_pushforward_matches_closed_form_sphere_maps():
  method = ODEMethod(max_iters=20, tol=1e-9, num_steps=256)
  tangent_chart = make_tangent_graph_chart(X0, tangent_frame=TANGENT_FRAME)
  stereographic_chart = make_stereographic_chart(X0, tangent_frame=TANGENT_FRAME)
  s_expected = jnp.array([0.12, -0.07], dtype=jnp.float64)
  x = sphere_exp_map(X0, TANGENT_FRAME @ s_expected)

  for chart in (tangent_chart, stereographic_chart):
    metric_fn = build_chart_metric_fn(chart)
    z = chart.x_to_z_fn(x)
    s_actual = z_to_s(metric_fn, chart.z0, z, method=method)
    tangent_log = sphere_log_map(X0, x)
    s_closed_form = TANGENT_FRAME.T @ tangent_log
    z_recovered = s_to_z(metric_fn, chart.z0, s_expected, method=method)

    assert jnp.allclose(s_actual, s_closed_form, atol=2e-5)
    assert jnp.allclose(z_recovered, z, atol=2e-5)


def test_chart_context_recovers_chart_invariant_local_density_and_eigenvalues():
  vmf = VonMisesFisher(mean_direction=MU, concentration=7.5)
  tangent_context = build_sphere_chart_context(
    vmf,
    X0,
    chart_name="tangent_graph",
    tangent_frame=TANGENT_FRAME,
    method=RefinedTaylorMethod(n_corrections=2),
  )
  stereographic_context = build_sphere_chart_context(
    vmf,
    X0,
    chart_name="stereographic",
    tangent_frame=TANGENT_FRAME,
    method=RefinedTaylorMethod(n_corrections=2),
  )
  s_points = jnp.array([
    [0.0, 0.0],
    [0.08, -0.04],
    [-0.06, 0.07],
  ], dtype=jnp.float64)
  analytic = vmf_entanglement_eigenvalue(vmf, X0)

  tangent_log_ps = jax.vmap(
    lambda s: log_ps_via_inverse_map(
      s,
      tangent_context.s_to_z_fn,
      tangent_context.z_to_s_fn,
      tangent_context.log_pz_fn,
    )
  )(s_points)
  stereographic_log_ps = jax.vmap(
    lambda s: log_ps_via_inverse_map(
      s,
      stereographic_context.s_to_z_fn,
      stereographic_context.z_to_s_fn,
      stereographic_context.log_pz_fn,
    )
  )(s_points)

  assert jnp.allclose(tangent_log_ps, stereographic_log_ps, atol=3e-5)
  assert jnp.allclose(tangent_context.eigenvalues, analytic * jnp.ones(2), atol=1e-8)
  assert jnp.allclose(stereographic_context.eigenvalues, analytic * jnp.ones(2), atol=1e-8)
  assert jnp.allclose(
    tangent_context.directions.T @ tangent_context.metric_fn(tangent_context.z0) @ tangent_context.directions,
    jnp.eye(2),
    atol=1e-8,
  )
  assert jnp.allclose(
    stereographic_context.directions.T @ stereographic_context.metric_fn(stereographic_context.z0) @ stereographic_context.directions,
    jnp.eye(2),
    atol=1e-8,
  )
