"""Tests verifying that the foreground (RNC-parametrized) and background (regular
z-grid) surfaces in analytic_verification.py lie on the same manifold.

All tests use the exact logarithmic map (geodesic ODE) so that the RNC-to-z
mapping is not limited by Taylor truncation error.
"""
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.independent_components import riemannian_ica
from local_geometry.paper.figures.analytic_verification import (
  make_decoder,
  log_pz,
  make_coord_fns,
  build_rnc_grid,
  map_grid_to_z,
  map_grid_to_surface,
  make_coordinate_curves,
  make_boundary,
  map_curve_to_z,
  map_curve_to_surface,
)


R = 1.0
N_GRID = 20
N_CURVES = 3
N_CURVE_PTS = 15


@pytest.fixture
def setup():
  """Shared fixture providing the exact coordinate functions, decoder, and
  metric objects used by the analytic verification figure."""
  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)))

  s_to_z_fn, dz_ds_fn = make_coord_fns(
    "exact", z0, j_s_to_z, j_z_to_s=j_z_to_s, metric_fn=metric_fn,
  )

  def log_det_g_fn(z):
    J_f = jax.jacfwd(f)(z)
    g = J_f.T @ J_f
    return jnp.linalg.slogdet(g)[1]

  return dict(
    z0=z0,
    f=f,
    j_s_to_z=j_s_to_z,
    j_z_to_s=j_z_to_s,
    metric_fn=metric_fn,
    metric=metric,
    s_to_z_fn=s_to_z_fn,
    dz_ds_fn=dz_ds_fn,
    log_det_g_fn=log_det_g_fn,
  )


# ------------------------------------------------------------------
# Test 1: Base point maps correctly through both pipelines
# ------------------------------------------------------------------

def test_base_point_exact_maps_to_z0(setup):
  """s = 0 in RNC must map back to z0 through the exact geodesic map."""
  s_to_z_fn = setup["s_to_z_fn"]
  z0 = setup["z0"]

  z_mapped = s_to_z_fn(jnp.zeros(2))
  assert jnp.allclose(z_mapped, z0, atol=1e-5), (
    f"s=0 mapped to {z_mapped}, expected {z0}"
  )


def test_base_point_surface_agreement(setup):
  """f(s_to_z(0)) must equal f(z0), so the foreground and background
  surfaces share the same base point in R^3."""
  s_to_z_fn = setup["s_to_z_fn"]
  z0 = setup["z0"]
  f = setup["f"]

  x_via_rnc = f(s_to_z_fn(jnp.zeros(2)))
  x_direct = f(z0)
  assert jnp.allclose(x_via_rnc, x_direct, atol=1e-5), (
    f"Surface base points differ: {x_via_rnc} vs {x_direct}"
  )


# ------------------------------------------------------------------
# Test 2: Foreground z-range contained within background z-range
# ------------------------------------------------------------------

def test_foreground_z_range_within_background(setup):
  """All foreground z-values must lie within the background grid extent
  (which adds a 12.5% margin on each side)."""
  s_to_z_fn = setup["s_to_z_fn"]

  S1, S2 = build_rnc_grid(R, N_GRID)
  Z1_mapped, Z2_mapped = map_grid_to_z(S1, S2, s_to_z_fn)

  z1_span = Z1_mapped.max() - Z1_mapped.min()
  z2_span = Z2_mapped.max() - Z2_mapped.min()
  bg_z1_min = Z1_mapped.min() - 0.125 * z1_span
  bg_z1_max = Z1_mapped.max() + 0.125 * z1_span
  bg_z2_min = Z2_mapped.min() - 0.125 * z2_span
  bg_z2_max = Z2_mapped.max() + 0.125 * z2_span

  assert np.all(Z1_mapped >= bg_z1_min), "Foreground z1 below background min"
  assert np.all(Z1_mapped <= bg_z1_max), "Foreground z1 above background max"
  assert np.all(Z2_mapped >= bg_z2_min), "Foreground z2 below background min"
  assert np.all(Z2_mapped <= bg_z2_max), "Foreground z2 above background max"


# ------------------------------------------------------------------
# Test 3: Surface points agree for shared z-values
# ------------------------------------------------------------------

def test_surface_points_agree_for_shared_z(setup):
  """For several RNC test points, f(s_to_z(s)) must equal f applied directly
  to the z-value returned by s_to_z. This confirms both surfaces use the
  same decoder and the pipeline does not introduce inconsistencies."""
  s_to_z_fn = setup["s_to_z_fn"]
  f = setup["f"]

  test_points = jnp.array([
    [0.0, 0.0],
    [0.3, 0.4],
    [-0.2, 0.5],
    [0.5, -0.3],
    [-0.4, -0.4],
  ])

  for s in test_points:
    z = s_to_z_fn(s)
    x_pipeline = f(z)
    x_direct = f(jnp.asarray(z))
    assert jnp.allclose(x_pipeline, x_direct, atol=1e-12), (
      f"At s={s}, pipeline and direct f(z) disagree: {x_pipeline} vs {x_direct}"
    )


# ------------------------------------------------------------------
# Test 4: Initial velocity Riemannian norm equals RNC Euclidean norm
# ------------------------------------------------------------------

def test_velocity_riemannian_norm_equals_rnc_norm(setup):
  """The disentangled RNC Jacobian maps s to a tangent vector v_std whose
  Riemannian length (in the pullback metric at z0) equals ||s||. This is
  the defining property of normal coordinates at the origin."""
  j_s_to_z = setup["j_s_to_z"]
  metric = setup["metric"]
  g_z0 = metric.components.value

  test_points = jnp.array([
    [0.3, 0.4],
    [-0.5, 0.1],
    [0.0, 0.7],
    [0.6, -0.2],
  ])

  for s in test_points:
    v_std = j_s_to_z.value @ s
    riemannian_norm_sq = v_std @ g_z0 @ v_std
    euclidean_norm_sq = s @ s
    assert jnp.allclose(riemannian_norm_sq, euclidean_norm_sq, atol=1e-10), (
      f"At s={s}, ||v||_g^2={riemannian_norm_sq:.8f} != ||s||^2={euclidean_norm_sq:.8f}"
    )


# ------------------------------------------------------------------
# Test 5: Coordinate curves from foreground lie on the manifold
# ------------------------------------------------------------------

def test_coordinate_curves_lie_on_manifold(setup):
  """Every point on a foreground coordinate curve, mapped through f, must
  equal f evaluated at the corresponding z-value. This confirms the curves
  are drawn on the manifold surface, not floating off it."""
  s_to_z_fn = setup["s_to_z_fn"]
  f = setup["f"]

  h_curves, v_curves = make_coordinate_curves(R, N_CURVES, N_CURVE_PTS)

  for s1_arr, s2_arr in h_curves + v_curves:
    z1c, z2c = map_curve_to_z(s1_arr, s2_arr, s_to_z_fn)
    pts_surface = map_curve_to_surface(z1c, z2c, f)

    Z_curve = jnp.stack([z1c, z2c], axis=-1)
    pts_direct = np.asarray(jax.vmap(f)(Z_curve))

    assert np.allclose(pts_surface, pts_direct, atol=1e-12), (
      "Coordinate curve point on surface differs from direct f(z) evaluation"
    )


# ------------------------------------------------------------------
# Test 6: Foreground and background grids agree through f
# ------------------------------------------------------------------

def test_foreground_and_background_grids_agree_through_f(setup):
  """The foreground surface grid (RNC -> z -> f) and a direct evaluation of f
  at the same mapped z-values must produce identical R^3 points."""
  s_to_z_fn = setup["s_to_z_fn"]
  f = setup["f"]

  S1, S2 = build_rnc_grid(R, N_GRID)
  Z1_mapped, Z2_mapped = map_grid_to_z(S1, S2, s_to_z_fn)
  surf_foreground = map_grid_to_surface(Z1_mapped, Z2_mapped, f)

  Z_stack = jnp.stack([Z1_mapped, Z2_mapped], axis=-1)
  shape = Z_stack.shape[:-1]
  Z_flat = Z_stack.reshape(-1, 2)
  X_flat = jax.vmap(f)(Z_flat)
  surf_direct = np.asarray(X_flat.reshape(*shape, -1))

  assert np.allclose(surf_foreground, surf_direct, atol=1e-12), (
    "Foreground surface grid and direct f(z) grid differ"
  )
