import jax
import jax.numpy as jnp
import numpy as np

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

from local_geometry.rica_jax.distance_graph import build_distance_graph, select_graph_neighborhood
from local_geometry.rica_jax.rnc_pushforward import RefinedTaylorMethod, geodesic_distance


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


def _curved_metric(x):
  x0, x1 = x
  g11 = 1.4 + 0.15 * x0 ** 2 + 0.05 * x1 ** 2
  g22 = 1.1 + 0.07 * x0 ** 2 + 0.18 * x1 ** 2
  g12 = 0.03 * x0 * x1
  return jnp.array([[g11, g12], [g12, g22]], dtype=x.dtype)


def _anisotropic_metric(x):
  x0, _ = x
  return jnp.array([
    [1.0 + 30.0 * (x0 + 1.5) ** 2, 0.0],
    [0.0, 1.0],
  ], dtype=x.dtype)


def test_build_distance_graph_returns_expected_shapes():
  z_points = jnp.array([
    [0.0, 0.0],
    [1.0, 0.0],
    [0.0, 2.0],
    [3.0, 0.0],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_flat_metric, z_points)

  assert graph.z_points.shape == (4, 2)
  assert graph.edge_distance_matrix.shape == (4, 4)
  assert graph.path_distance_matrix.shape == (4, 4)
  assert np.allclose(np.diag(np.asarray(graph.edge_distance_matrix)), 0.0, atol=1e-12)
  assert np.allclose(np.diag(np.asarray(graph.path_distance_matrix)), 0.0, atol=1e-12)


def test_flat_metric_edge_distance_matrix_matches_euclidean_distances():
  z_points = jnp.array([
    [0.0, 0.0],
    [1.0, 0.0],
    [0.0, 2.0],
    [3.0, 0.0],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_flat_metric, z_points)
  expected = np.linalg.norm(
    np.asarray(z_points)[:, None, :] - np.asarray(z_points)[None, :, :],
    axis=-1,
  )

  assert np.allclose(graph.edge_distance_matrix, expected, atol=1e-12)


def test_taylor_string_matches_refined_taylor_method():
  z_points = jnp.array([
    [0.00, 0.00],
    [0.05, -0.02],
    [0.09, 0.04],
    [0.15, -0.01],
  ], dtype=jnp.float64)

  graph_default = build_distance_graph(_curved_metric, z_points, method="taylor")
  graph_refined = build_distance_graph(
    _curved_metric,
    z_points,
    method=RefinedTaylorMethod(),
  )

  assert np.allclose(graph_default.edge_distance_matrix, graph_refined.edge_distance_matrix, atol=1e-12)
  assert np.allclose(graph_default.path_distance_matrix, graph_refined.path_distance_matrix, atol=1e-12)


def test_edge_distance_matrix_is_symmetric():
  z_points = jnp.array([
    [0.00, 0.00],
    [0.05, -0.02],
    [0.09, 0.04],
    [0.15, -0.01],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_curved_metric, z_points, method=RefinedTaylorMethod())
  edge_distance_matrix = np.asarray(graph.edge_distance_matrix)
  expected_01 = 0.5 * (
    float(geodesic_distance(_curved_metric, z_points[0], z_points[1], method=RefinedTaylorMethod()))
    + float(geodesic_distance(_curved_metric, z_points[1], z_points[0], method=RefinedTaylorMethod()))
  )

  assert np.allclose(edge_distance_matrix, edge_distance_matrix.T, atol=1e-12)
  assert np.allclose(edge_distance_matrix[0, 1], expected_01, atol=1e-12)


def test_graph_distance_is_bounded_by_direct_edge_distance():
  z_points = jnp.array([
    [0.00, 0.00],
    [0.05, -0.02],
    [0.09, 0.04],
    [0.15, -0.01],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_curved_metric, z_points, method=RefinedTaylorMethod())

  assert np.all(graph.path_distance_matrix <= graph.edge_distance_matrix + 1e-12)


def test_edge_distance_matrix_is_finite_and_positive_off_diagonal():
  z_points = jnp.array([
    [0.00, 0.00],
    [0.05, -0.02],
    [0.09, 0.04],
    [0.15, -0.01],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_curved_metric, z_points, method=RefinedTaylorMethod())
  edge_distance_matrix = np.asarray(graph.edge_distance_matrix)
  off_diagonal = edge_distance_matrix[~np.eye(edge_distance_matrix.shape[0], dtype=bool)]

  assert np.isfinite(edge_distance_matrix).all()
  assert np.all(off_diagonal > 0.0)


def test_nearest_nodes_matches_flat_metric_k_neighbor_lookup():
  z_points = jnp.array([
    [0.0, 0.0],
    [1.0, 0.0],
    [0.0, 2.0],
    [3.0, 0.0],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_flat_metric, z_points)

  assert np.array_equal(
    np.asarray(graph.nearest_nodes(jnp.array([0.1, 0.2], dtype=jnp.float64), k=2)),
    np.array([0, 1]),
  )
  assert np.array_equal(
    np.asarray(graph.nearest_nodes(jnp.array([2.6, 0.1], dtype=jnp.float64), k=2)),
    np.array([3, 1]),
  )


def test_nearest_nodes_uses_symmetrized_geodesic_distance():
  z_points = jnp.array([
    [-0.35, 0.0],
    [0.0, 0.45],
    [0.4, 0.0],
  ], dtype=jnp.float64)
  z_query = jnp.array([0.0, 0.0], dtype=jnp.float64)

  graph = build_distance_graph(_anisotropic_metric, z_points, method=RefinedTaylorMethod())
  nearest = np.asarray(graph.nearest_nodes(z_query, k=2))

  direct_distances = []
  for z_point in z_points:
    forward = geodesic_distance(
      _anisotropic_metric,
      z_query,
      z_point,
      method=RefinedTaylorMethod(),
    )
    reverse = geodesic_distance(
      _anisotropic_metric,
      z_point,
      z_query,
      method=RefinedTaylorMethod(),
    )
    direct_distances.append(float(0.5 * (forward + reverse)))
  expected = np.argsort(np.asarray(direct_distances))[:2]
  euclidean_order = np.argsort(np.linalg.norm(np.asarray(z_points - z_query), axis=1))[:2]

  assert np.array_equal(nearest, expected)
  assert not np.array_equal(nearest, euclidean_order)


def test_distance_supports_jit_and_vmap():
  z_points = jnp.array([
    [0.0, 0.0],
    [1.0, 0.0],
    [2.1, 0.0],
    [3.3, 0.0],
  ], dtype=jnp.float64)

  graph = build_distance_graph(_flat_metric, z_points)
  starts = jnp.array([0, 1], dtype=jnp.int32)
  ends = jnp.array([3, 3], dtype=jnp.int32)
  expected = jnp.array([3.3, 2.3], dtype=jnp.float64)

  @jax.jit
  def batched_distances(graph, start_indices, end_indices):
    query_fn = lambda start_index, end_index: graph.distance(start_index, end_index)
    return jax.vmap(query_fn)(start_indices, end_indices)

  actual = batched_distances(graph, starts, ends)

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


def test_select_graph_neighborhood_filters_points_within_epsilon():
  z0 = jnp.array([0.0, 0.0], dtype=jnp.float64)
  z_samples = jnp.array([
    [0.10, 0.00],
    [0.24, 0.00],
    [0.31, 0.00],
    [0.00, 0.40],
  ], dtype=jnp.float64)

  retained, mask, distances = select_graph_neighborhood(
    _flat_metric,
    z0,
    z_samples,
    epsilon=0.25,
  )

  assert np.allclose(distances, np.array([0.10, 0.24, 0.31, 0.40]), atol=1e-12)
  assert np.array_equal(np.asarray(mask), np.array([True, True, False, False]))
  assert np.allclose(
    retained,
    np.array([
      [0.10, 0.00],
      [0.24, 0.00],
    ]),
    atol=1e-12,
  )
