import jax
import jax.numpy as jnp
import pytest

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

from local_coordinates.basis import get_standard_basis
from local_coordinates.exponential_map import (
  exponential_map_ode,
  logarithmic_map_ode,
  logarithmic_map_taylor_refined,
)
from local_coordinates.jet import Jet, function_to_jet
from local_coordinates.metric import RiemannianMetric
from local_coordinates.normal_coords import get_rnc_jacobians
from local_coordinates.tangent import TangentVector
from local_geometry.rica_jax.rnc_pushforward import (
  ODEMethod,
  RefinedTaylorMethod,
  geodesic_distance,
  s_to_z,
  z_to_s,
)


def _metric_components(x):
  x0, x1 = x
  g11 = 1.5 + 0.2 * x0 ** 2 + 0.1 * x1 ** 2
  g22 = 1.2 + 0.15 * x0 ** 2 + 0.25 * x1 ** 2
  g12 = 0.05 * x0 * x1
  return jnp.array([[g11, g12], [g12, g22]])


def _metric_fn(x):
  basis = get_standard_basis(x)
  components = function_to_jet(_metric_components, x)
  return RiemannianMetric(basis=basis, components=components)


def _build_metric(p):
  basis = get_standard_basis(p)
  components = function_to_jet(_metric_components, p)
  return RiemannianMetric(basis=basis, components=components)


def _create_tangent_vector(p, v):
  dim = p.shape[0]
  return TangentVector(
    p=p,
    basis=get_standard_basis(p),
    components=Jet(
      value=v,
      gradient=jnp.zeros((dim, dim)),
      hessian=jnp.zeros((dim, dim, dim)),
    ),
  )


def _flat_metric(x):
  return jnp.eye(x.shape[0])


def _higher_dim_metric(x):
  dim = x.shape[0]
  g = jnp.eye(dim)
  for i in range(dim):
    g = g.at[i, i].add(0.1 * x[i] ** 2)
  for i in range(dim - 1):
    off_diag = 0.02 * x[i] * x[i + 1]
    g = g.at[i, i + 1].add(off_diag)
    g = g.at[i + 1, i].add(off_diag)
  return g


def _rotation(theta):
  return jnp.array([
    [jnp.cos(theta), -jnp.sin(theta)],
    [jnp.sin(theta), jnp.cos(theta)],
  ])


Z0 = jnp.array([0.17, -0.23])
Z = jnp.array([0.23, -0.19])
FRAME_ROTATION = _rotation(0.37)


def test_refined_taylor_matches_reference():
  metric = _build_metric(Z0)
  method = RefinedTaylorMethod(n_corrections=2)

  expected = logarithmic_map_taylor_refined(
    metric,
    Z,
    n_corrections=method.n_corrections,
  )
  actual = z_to_s(_metric_components, Z0, Z, method=method)

  assert jnp.allclose(actual, expected, atol=1e-10)


def test_ode_matches_reference():
  metric = _build_metric(Z0)
  method = ODEMethod(max_iters=12, tol=1e-7, num_steps=256)

  j_z_to_s, _ = get_rnc_jacobians(metric)
  expected_std = logarithmic_map_ode(
    Z0,
    Z,
    _metric_fn,
    max_iters=method.max_iters,
    tol=method.tol,
  )
  expected = j_z_to_s.value @ expected_std
  actual = z_to_s(_metric_components, Z0, Z, method=method)

  assert jnp.allclose(actual, expected, atol=1e-4)


def test_rotated_refined_taylor_matches_reference():
  metric = _build_metric(Z0)
  method = RefinedTaylorMethod(n_corrections=2)

  expected_std = logarithmic_map_taylor_refined(
    metric,
    Z,
    n_corrections=method.n_corrections,
  )
  expected = FRAME_ROTATION.T @ expected_std
  actual = z_to_s(
    _metric_components,
    Z0,
    Z,
    method=method,
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(actual, expected, atol=1e-10)


def test_rotated_ode_matches_reference():
  metric = _build_metric(Z0)
  method = ODEMethod(max_iters=12, tol=1e-7, num_steps=256)

  j_z_to_s, _ = get_rnc_jacobians(metric, frame_rotation=FRAME_ROTATION)
  expected_std = logarithmic_map_ode(
    Z0,
    Z,
    _metric_fn,
    max_iters=method.max_iters,
    tol=method.tol,
  )
  expected = j_z_to_s.value @ expected_std
  actual = z_to_s(
    _metric_components,
    Z0,
    Z,
    method=method,
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(actual, expected, atol=1e-4)


def test_s_to_z_taylor_matches_reference():
  metric = _build_metric(Z0)
  method = RefinedTaylorMethod(n_corrections=0)
  s = jnp.array([0.03, -0.02])

  _, j_s_to_z = get_rnc_jacobians(metric)
  expected = (
    Z0
    + jnp.einsum("ij,j->i", j_s_to_z.value, s)
    + 0.5 * jnp.einsum("ijk,j,k->i", j_s_to_z.gradient, s, s)
    + (1.0 / 6.0) * jnp.einsum("ijkl,j,k,l->i", j_s_to_z.hessian, s, s, s)
  )
  actual = s_to_z(_metric_components, Z0, s, method=method)

  assert jnp.allclose(actual, expected, atol=1e-10)


def test_s_to_z_ode_matches_reference():
  metric = _build_metric(Z0)
  method = ODEMethod(max_iters=12, tol=1e-7, num_steps=256)
  s = jnp.array([0.03, -0.02])

  _, j_s_to_z = get_rnc_jacobians(metric)
  v_std = j_s_to_z.value @ s
  tangent = _create_tangent_vector(Z0, v_std)
  expected = exponential_map_ode(tangent, _metric_fn)
  actual = s_to_z(_metric_components, Z0, s, method=method)

  assert jnp.allclose(actual, expected, atol=1e-4)


def test_rotated_s_to_z_taylor_matches_reference():
  metric = _build_metric(Z0)
  method = RefinedTaylorMethod(n_corrections=0)
  s = jnp.array([0.03, -0.02])

  _, j_s_to_z = get_rnc_jacobians(metric, frame_rotation=FRAME_ROTATION)
  expected = (
    Z0
    + jnp.einsum("ij,j->i", j_s_to_z.value, s)
    + 0.5 * jnp.einsum("ijk,j,k->i", j_s_to_z.gradient, s, s)
    + (1.0 / 6.0) * jnp.einsum("ijkl,j,k,l->i", j_s_to_z.hessian, s, s, s)
  )
  actual = s_to_z(
    _metric_components,
    Z0,
    s,
    method=method,
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(actual, expected, atol=1e-10)


def test_rotated_s_to_z_ode_matches_reference():
  metric = _build_metric(Z0)
  method = ODEMethod(max_iters=12, tol=1e-7, num_steps=256)
  s = jnp.array([0.03, -0.02])

  _, j_s_to_z = get_rnc_jacobians(metric, frame_rotation=FRAME_ROTATION)
  v_std = j_s_to_z.value @ s
  tangent = _create_tangent_vector(Z0, v_std)
  expected = exponential_map_ode(tangent, _metric_fn)
  actual = s_to_z(
    _metric_components,
    Z0,
    s,
    method=method,
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(actual, expected, atol=1e-4)


def test_flat_metric_reduces_to_linear_map():
  z0 = jnp.array([0.3, -0.4, 0.1])
  z = jnp.array([0.45, -0.2, -0.05])
  expected = z - z0
  s = jnp.array([0.15, 0.2, -0.15])

  actual_refined = z_to_s(_flat_metric, z0, z, method=RefinedTaylorMethod(n_corrections=3))
  actual_ode = z_to_s(_flat_metric, z0, z, method=ODEMethod(num_steps=32))
  forward_refined = s_to_z(_flat_metric, z0, s, method=RefinedTaylorMethod(n_corrections=3))
  forward_ode = s_to_z(_flat_metric, z0, s, method=ODEMethod(num_steps=32))

  assert jnp.allclose(actual_refined, expected, atol=1e-12)
  assert jnp.allclose(actual_ode, expected, atol=1e-8)
  assert jnp.allclose(forward_refined, z0 + s, atol=1e-12)
  assert jnp.allclose(forward_ode, z0 + s, atol=1e-8)


def test_geodesic_distance_flat_metric_matches_euclidean_distance():
  z0 = jnp.array([0.3, -0.4, 0.1], dtype=jnp.float64)
  z1 = jnp.array([0.45, -0.2, -0.05], dtype=jnp.float64)
  expected = jnp.linalg.norm(z1 - z0)

  actual = geodesic_distance(_flat_metric, z0, z1, method="taylor")

  assert jnp.allclose(actual, expected, atol=1e-12)


def test_geodesic_distance_improved_alias_matches_refined_taylor_method():
  z_local = Z0 + jnp.array([0.015, -0.01], dtype=jnp.float64)

  actual = geodesic_distance(_metric_components, Z0, z_local, method="taylor")
  expected = geodesic_distance(
    _metric_components,
    Z0,
    z_local,
    method=RefinedTaylorMethod(),
  )

  assert jnp.allclose(actual, expected, atol=1e-12)


def test_geodesic_distance_is_invariant_to_frame_rotation():
  z_local = Z0 + jnp.array([0.015, -0.01], dtype=jnp.float64)

  unrotated = geodesic_distance(_metric_components, Z0, z_local, method="taylor")
  rotated = geodesic_distance(
    _metric_components,
    Z0,
    z_local,
    method="taylor",
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(rotated, unrotated, atol=1e-12)


def test_geodesic_distance_curved_metric_is_finite_and_positive():
  z_local = Z0 + jnp.array([0.015, -0.01], dtype=jnp.float64)

  actual = geodesic_distance(_metric_components, Z0, z_local, method="taylor")

  assert jnp.isfinite(actual)
  assert actual > 0.0


def test_refined_and_ode_agree_locally():
  z_local = Z0 + jnp.array([0.015, -0.01])
  refined = z_to_s(_metric_components, Z0, z_local, method=RefinedTaylorMethod(n_corrections=2))
  exact = z_to_s(_metric_components, Z0, z_local, method=ODEMethod(max_iters=12, tol=1e-8, num_steps=256))
  s_local = jnp.array([0.015, -0.01])
  z_refined = s_to_z(_metric_components, Z0, s_local, method=RefinedTaylorMethod(n_corrections=2))
  z_exact = s_to_z(_metric_components, Z0, s_local, method=ODEMethod(max_iters=12, tol=1e-8, num_steps=256))

  assert jnp.allclose(refined, exact, atol=1e-5)
  assert jnp.allclose(z_refined, z_exact, atol=1e-5)


def test_z_to_s_and_s_to_z_roundtrip_on_s():
  s = jnp.array([0.02, -0.015])
  z = s_to_z(_metric_components, Z0, s, method=RefinedTaylorMethod(n_corrections=1))
  s_recovered = z_to_s(_metric_components, Z0, z, method=RefinedTaylorMethod(n_corrections=1))

  assert jnp.allclose(s_recovered, s, atol=1e-6)


def test_s_to_z_and_z_to_s_roundtrip_on_z():
  z = Z0 + jnp.array([0.02, -0.015])
  s = z_to_s(_metric_components, Z0, z, method=RefinedTaylorMethod(n_corrections=1))
  z_recovered = s_to_z(_metric_components, Z0, s, method=RefinedTaylorMethod(n_corrections=1))

  assert jnp.allclose(z_recovered, z, atol=1e-6)


def test_rotated_roundtrip_on_s():
  s = jnp.array([0.02, -0.015])
  z = s_to_z(
    _metric_components,
    Z0,
    s,
    method=RefinedTaylorMethod(n_corrections=1),
    frame_rotation=FRAME_ROTATION,
  )
  s_recovered = z_to_s(
    _metric_components,
    Z0,
    z,
    method=RefinedTaylorMethod(n_corrections=1),
    frame_rotation=FRAME_ROTATION,
  )

  assert jnp.allclose(s_recovered, s, atol=1e-6)


def test_higher_dim_refined_matches_reference():
  dim = 4
  z0 = 0.1 * jnp.arange(dim, dtype=jnp.float64)
  z = z0 + jnp.array([0.04, -0.02, 0.03, -0.01], dtype=jnp.float64)
  method = RefinedTaylorMethod(n_corrections=1)

  metric = RiemannianMetric(
    basis=get_standard_basis(z0),
    components=function_to_jet(_higher_dim_metric, z0),
  )
  expected = logarithmic_map_taylor_refined(
    metric,
    z,
    n_corrections=method.n_corrections,
  )
  actual = z_to_s(_higher_dim_metric, z0, z, method=method)

  assert actual.shape == (dim,)
  assert jnp.allclose(actual, expected, atol=1e-10)


def test_invalid_method_configs_raise():
  with pytest.raises(ValueError, match="nonnegative"):
    z_to_s(_metric_components, Z0, Z, method=RefinedTaylorMethod(n_corrections=-1))

  with pytest.raises(ValueError, match="positive"):
    z_to_s(_metric_components, Z0, Z, method=ODEMethod(num_steps=0))
