import logging
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
from local_geometry.rica_jax.sparse_distance_graph import (
  _all_pairs_sparse_shortest_paths,
  build_sparse_distance_graph,
  select_geodesic_ball_neighborhood,
  select_sparse_graph_neighborhood,
)
import local_geometry.rica_jax.sparse_distance_graph as sparse_distance_graph


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_sparse_complete_graph_matches_dense_graph():
  z_points = jnp.array([
    [0.00, 0.00],
    [0.05, -0.02],
    [0.09, 0.04],
    [0.15, -0.01],
  ], dtype=jnp.float64)
  method = RefinedTaylorMethod()

  dense_graph = build_distance_graph(_curved_metric, z_points, method=method)
  sparse_graph = build_sparse_distance_graph(
    _curved_metric,
    z_points,
    k=z_points.shape[0] - 1,
    method=method,
  )

  assert np.allclose(sparse_graph.edge_distance_matrix, dense_graph.edge_distance_matrix, atol=1e-12)
  assert np.allclose(sparse_graph.path_distance_matrix, dense_graph.path_distance_matrix, atol=1e-12)
  assert np.allclose(sparse_graph.distances_from(1), dense_graph.distances_from(1), atol=1e-12)
  assert np.allclose(sparse_graph.distance(0, 3), dense_graph.distance(0, 3), atol=1e-12)


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

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

  assert np.allclose(sparse_graph.path_distance_matrix, dense_graph.path_distance_matrix, atol=1e-12)
  assert np.allclose(sparse_graph.path_distance_matrix, expected, atol=1e-12)


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

  graph = build_sparse_distance_graph(_flat_metric, z_points, k=1)
  expected_neighbor_indices = np.array([
    [1, -1],
    [0, 2],
    [1, 3],
    [2, -1],
  ], dtype=np.int32)
  expected_neighbor_distances = np.array([
    [1.0, np.inf],
    [1.0, 1.1],
    [1.1, 1.2],
    [1.2, np.inf],
  ], dtype=np.float64)

  assert np.allclose(graph.edge_distance_matrix, graph.edge_distance_matrix.T, atol=1e-12)
  assert np.array_equal(np.asarray(graph.neighbor_indices), expected_neighbor_indices)
  assert np.allclose(np.asarray(graph.neighbor_distances), expected_neighbor_distances, atol=1e-12)


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

  dense_retained, dense_mask, dense_distances = select_graph_neighborhood(
    _flat_metric,
    z0,
    z_samples,
    epsilon=0.25,
  )
  sparse_retained, sparse_mask, sparse_distances = select_sparse_graph_neighborhood(
    _flat_metric,
    z0,
    z_samples,
    epsilon=0.25,
    k=2,
  )

  assert np.allclose(sparse_distances, dense_distances, atol=1e-12)
  assert np.array_equal(np.asarray(sparse_mask), np.asarray(dense_mask))
  assert np.allclose(sparse_retained, dense_retained, atol=1e-12)


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

  retained, mask, distances = select_geodesic_ball_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.24],
    ]),
    atol=1e-12,
  )


def test_select_geodesic_ball_neighborhood_reduces_geodesic_evaluations(monkeypatch):
  z0 = jnp.array([0.0], dtype=jnp.float64)
  z_samples = jnp.arange(64, dtype=jnp.float64)[:, None]
  evaluated_sizes = []

  def fake_symmetrized_query_geodesic_distances(
    gz,
    z_query,
    z_points,
    *,
    method=None,
    frame_rotation=None,
  ):
    del gz, z_query, method, frame_rotation
    evaluated_sizes.append(int(z_points.shape[0]))
    return jnp.linalg.norm(z_points, axis=-1)

  monkeypatch.setattr(
    sparse_distance_graph,
    "_symmetrized_query_geodesic_distances",
    fake_symmetrized_query_geodesic_distances,
  )

  retained, mask, distances = select_geodesic_ball_neighborhood(
    _flat_metric,
    z0,
    z_samples,
    epsilon=0.5,
  )

  assert evaluated_sizes == [16]
  assert retained.shape == (1, 1)
  assert np.array_equal(np.asarray(mask)[1:], np.zeros((63,), dtype=bool))
  assert np.isinf(np.asarray(distances)[16:]).all()


def test_select_geodesic_ball_neighborhood_respects_shortlist_override(monkeypatch):
  z0 = jnp.array([0.0], dtype=jnp.float64)
  z_samples = jnp.arange(8, dtype=jnp.float64)[:, None]
  evaluated_sizes = []

  def fake_symmetrized_query_geodesic_distances(
    gz,
    z_query,
    z_points,
    *,
    method=None,
    frame_rotation=None,
  ):
    del gz, z_query, method, frame_rotation
    evaluated_sizes.append(int(z_points.shape[0]))
    return jnp.linalg.norm(z_points, axis=-1)

  monkeypatch.setattr(
    sparse_distance_graph,
    "_symmetrized_query_geodesic_distances",
    fake_symmetrized_query_geodesic_distances,
  )

  retained, mask, distances = select_geodesic_ball_neighborhood(
    _flat_metric,
    z0,
    z_samples,
    epsilon=0.5,
    shortlist_size=2,
  )

  assert evaluated_sizes == [2]
  assert retained.shape == (1, 1)
  assert np.array_equal(np.asarray(mask)[1:], np.zeros((7,), dtype=bool))
  assert np.isinf(np.asarray(distances)[2:]).all()


def test_select_geodesic_ball_neighborhood_validates_epsilon():
  z0 = jnp.array([0.0], dtype=jnp.float64)
  z_samples = jnp.array([[0.10], [0.24]], dtype=jnp.float64)

  with np.testing.assert_raises(ValueError):
    select_geodesic_ball_neighborhood(
      _flat_metric,
      z0,
      z_samples,
      epsilon=0.0,
    )

  with np.testing.assert_raises(ValueError):
    select_geodesic_ball_neighborhood(
      _flat_metric,
      z0,
      z_samples,
      epsilon=0.25,
      shortlist_size=0,
    )


def test_sparse_retained_edges_use_symmetrized_geodesic_weights():
  z_points = jnp.array([
    [-0.35, 0.0],
    [0.0, 0.45],
    [0.4, 0.0],
    [0.75, -0.1],
  ], dtype=jnp.float64)

  sparse_graph = build_sparse_distance_graph(
    _anisotropic_metric,
    z_points,
    k=1,
    method=RefinedTaylorMethod(),
  )
  edge_distance_matrix = np.asarray(sparse_graph.edge_distance_matrix)
  finite_mask = np.isfinite(edge_distance_matrix) & ~np.eye(edge_distance_matrix.shape[0], dtype=bool)
  source_indices, target_indices = np.nonzero(np.triu(finite_mask, k=1))

  for source_index, target_index in zip(source_indices, target_indices):
    expected = 0.5 * (
      float(
        geodesic_distance(
          _anisotropic_metric,
          z_points[source_index],
          z_points[target_index],
          method=RefinedTaylorMethod(),
        )
      )
      + float(
        geodesic_distance(
          _anisotropic_metric,
          z_points[target_index],
          z_points[source_index],
          method=RefinedTaylorMethod(),
        )
      )
    )
    assert np.allclose(edge_distance_matrix[source_index, target_index], expected, atol=1e-12)


def test_sparse_builder_evaluates_geodesics_only_on_candidate_edges(monkeypatch):
  z_points = jnp.arange(8, dtype=jnp.float64)[:, None]
  batched_call_sizes = []

  def fake_batched_geodesic_distances(
    gz,
    z_sources,
    z_targets,
    *,
    method=None,
    frame_rotation=None,
  ):
    del gz, method, frame_rotation
    batched_call_sizes.append(int(z_sources.shape[0]))
    return jnp.linalg.norm(z_sources - z_targets, axis=-1)

  monkeypatch.setattr(
    sparse_distance_graph,
    "_batched_geodesic_distances",
    fake_batched_geodesic_distances,
  )

  graph = build_sparse_distance_graph(_flat_metric, z_points, k=1)

  total_directed_geodesic_evals = sum(batched_call_sizes)
  full_directed_pair_count = int(z_points.shape[0]) * (int(z_points.shape[0]) - 1)
  assert total_directed_geodesic_evals < full_directed_pair_count
  assert np.allclose(graph.path_distance_matrix[0], np.arange(8, dtype=np.float64), atol=1e-12)


def test_sparse_builder_does_not_use_dense_adjacency_packing(monkeypatch):
  z_points = jnp.arange(6, dtype=jnp.float64)[:, None]

  def fail_if_called(sparse_adjacency):
    del sparse_adjacency
    raise AssertionError("_adjacency_lists should not be used by build_sparse_distance_graph")

  monkeypatch.setattr(sparse_distance_graph, "_adjacency_lists", fail_if_called)

  graph = build_sparse_distance_graph(_flat_metric, z_points, k=1)

  assert np.allclose(graph.path_distance_matrix[0], np.arange(6, dtype=np.float64), atol=1e-12)


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

  graph = build_sparse_distance_graph(_flat_metric, z_points, k=1)
  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_all_pairs_sparse_shortest_paths_wraps_loop_with_tqdm(monkeypatch):
  tqdm_calls = []

  def fake_tqdm(iterable, desc, leave):
    tqdm_calls.append({
      "iterable": list(iterable),
      "desc": desc,
      "leave": leave,
    })
    return tqdm_calls[-1]["iterable"]

  monkeypatch.setattr(sparse_distance_graph.tqdm, "tqdm", fake_tqdm)

  neighbor_indices_list = [
    np.array([1], dtype=np.int32),
    np.array([0, 2], dtype=np.int32),
    np.array([1], dtype=np.int32),
  ]
  neighbor_distances_list = [
    np.array([1.0], dtype=np.float64),
    np.array([1.0, 2.0], dtype=np.float64),
    np.array([2.0], dtype=np.float64),
  ]

  path_distance_matrix = _all_pairs_sparse_shortest_paths(
    neighbor_indices_list,
    neighbor_distances_list,
    dtype=np.dtype(np.float64),
  )

  expected = np.array([
    [0.0, 1.0, 3.0],
    [1.0, 0.0, 2.0],
    [3.0, 2.0, 0.0],
  ], dtype=np.float64)
  assert len(tqdm_calls) == 1
  assert tqdm_calls[0]["iterable"] == [0, 1, 2]
  assert tqdm_calls[0]["desc"] == "Sparse shortest paths"
  assert tqdm_calls[0]["leave"] is False
  assert np.allclose(path_distance_matrix, expected, atol=1e-12)


def test_build_sparse_distance_graph_emits_progress_logs(monkeypatch, caplog):
  z_points = jnp.array([
    [0.0],
    [1.0],
    [2.0],
  ], dtype=jnp.float64)

  def fake_candidate_degree(n_points, k):
    del n_points, k
    return 2

  def fake_euclidean_candidate_neighbors(z_points, candidate_k):
    del z_points, candidate_k
    return np.array([
      [1, 2],
      [0, 2],
      [1, 0],
    ], dtype=np.int32)

  def fake_candidate_edge_weights(
    gz,
    z_points,
    candidate_neighbors,
    *,
    method=None,
    frame_rotation=None,
  ):
    del gz, z_points, candidate_neighbors, method, frame_rotation
    return (
      np.array([0, 1], dtype=np.int32),
      np.array([1, 2], dtype=np.int32),
      np.array([1.0, 1.0], dtype=np.float64),
    )

  def fake_top_k_candidate_lists(n_points, pair_sources, pair_targets, pair_weights, k):
    del n_points, pair_sources, pair_targets, pair_weights, k
    return (
      [
        np.array([1], dtype=np.int32),
        np.array([0, 2], dtype=np.int32),
        np.array([1], dtype=np.int32),
      ],
      [
        np.array([1.0], dtype=np.float64),
        np.array([1.0, 1.0], dtype=np.float64),
        np.array([1.0], dtype=np.float64),
      ],
    )

  def fake_symmetrize_directed_knn(directed_neighbor_indices, directed_neighbor_distances, *, dtype):
    del directed_neighbor_indices, directed_neighbor_distances, dtype
    neighbor_indices_list = [
      np.array([1], dtype=np.int32),
      np.array([0, 2], dtype=np.int32),
      np.array([1], dtype=np.int32),
    ]
    neighbor_distances_list = [
      np.array([1.0], dtype=np.float64),
      np.array([1.0, 1.0], dtype=np.float64),
      np.array([1.0], dtype=np.float64),
    ]
    padded_neighbor_indices = np.array([
      [1, -1],
      [0, 2],
      [1, -1],
    ], dtype=np.int32)
    padded_neighbor_distances = np.array([
      [1.0, np.inf],
      [1.0, 1.0],
      [1.0, np.inf],
    ], dtype=np.float64)
    return (
      neighbor_indices_list,
      neighbor_distances_list,
      padded_neighbor_indices,
      padded_neighbor_distances,
    )

  def fake_all_pairs_sparse_shortest_paths(
    neighbor_indices_list,
    neighbor_distances_list,
    *,
    dtype,
  ):
    del neighbor_indices_list, neighbor_distances_list, dtype
    return jnp.array([
      [0.0, 1.0, 2.0],
      [1.0, 0.0, 1.0],
      [2.0, 1.0, 0.0],
    ], dtype=jnp.float64)

  monkeypatch.setattr(
    sparse_distance_graph,
    "_candidate_degree",
    fake_candidate_degree,
  )
  monkeypatch.setattr(
    sparse_distance_graph,
    "_euclidean_candidate_neighbors",
    fake_euclidean_candidate_neighbors,
  )
  monkeypatch.setattr(
    sparse_distance_graph,
    "_candidate_edge_weights",
    fake_candidate_edge_weights,
  )
  monkeypatch.setattr(
    sparse_distance_graph,
    "_top_k_candidate_lists",
    fake_top_k_candidate_lists,
  )
  monkeypatch.setattr(
    sparse_distance_graph,
    "_symmetrize_directed_knn",
    fake_symmetrize_directed_knn,
  )
  monkeypatch.setattr(
    sparse_distance_graph,
    "_all_pairs_sparse_shortest_paths",
    fake_all_pairs_sparse_shortest_paths,
  )

  with caplog.at_level(logging.INFO, logger=sparse_distance_graph.logger.name):
    graph = build_sparse_distance_graph(_flat_metric, z_points, k=1)

  messages = [record.getMessage() for record in caplog.records]
  assert "Building sparse distance graph with n_points=3 and k=1" in messages
  assert "Selecting Euclidean candidate neighbors with candidate_k=2" in messages
  assert "Computing geodesic edge weights on candidate graph" in messages
  assert "Selecting directed geodesic top-k neighbors from candidate graph" in messages
  assert "Symmetrizing directed KNN graph" in messages
  assert "Running all pairs sparse shortest paths" in messages
  assert "Finished sparse distance graph construction" in messages
  assert graph.k == 1
  assert np.allclose(graph.path_distance_matrix, np.array([
    [0.0, 1.0, 2.0],
    [1.0, 0.0, 1.0],
    [2.0, 1.0, 0.0],
  ]))
