import logging
from types import SimpleNamespace

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

import local_geometry.experiments.disentanglement_neighborhood as disentanglement_neighborhood
from local_geometry.experiments.disentanglement_neighborhood import (
  compute_density_grid,
  map_s_samples_to_z,
  sample_z0s_from_prior,
  select_graph_local_neighborhood,
)


def _standard_normal_log_pz(z):
  return -0.5 * jnp.dot(z, z)


def test_sample_z0s_from_prior_is_reproducible():
  key = jax.random.PRNGKey(7)

  samples_a = sample_z0s_from_prior(key, 4, 2)
  samples_b = sample_z0s_from_prior(key, 4, 2)

  assert samples_a.shape == (4, 2)
  assert jnp.allclose(samples_a, samples_b)


def test_map_s_samples_to_z_preserves_shape_and_is_finite():
  z0 = jnp.array([0.2, -0.1], dtype=jnp.float64)
  s_samples = jnp.array([
    [0.0, 0.0],
    [0.1, -0.2],
    [-0.15, 0.05],
  ], dtype=jnp.float64)

  z_samples = map_s_samples_to_z(s_samples, lambda s: z0 + s)

  assert z_samples.shape == s_samples.shape
  assert jnp.isfinite(z_samples).all()
  assert jnp.allclose(z_samples, z0 + s_samples)


def test_compute_density_grid_returns_expected_shape_and_values():
  z0 = jnp.array([0.25, -0.15], dtype=jnp.float64)
  s_1d = jnp.linspace(-0.2, 0.2, 5)
  S1, S2 = jnp.meshgrid(s_1d, s_1d)

  log_density_grid = compute_density_grid(
    S1,
    S2,
    lambda s: z0 + s,
    lambda z: z - z0,
    _standard_normal_log_pz,
  )

  assert log_density_grid.shape == (5, 5)
  assert np.isfinite(log_density_grid).all()
  assert jnp.allclose(log_density_grid[2, 2], _standard_normal_log_pz(z0))


def test_select_graph_local_neighborhood_returns_variable_counts_inside_epsilon():
  def flat_metric(z):
    return jnp.eye(z.shape[0], dtype=z.dtype)

  z_prior_samples = jnp.array([
    [0.05, 0.0],
    [0.18, 0.0],
    [0.40, 0.0],
    [1.10, 0.0],
  ], dtype=jnp.float64)
  epsilon = 0.25

  context_a = SimpleNamespace(
    metric_fn=flat_metric,
    z0=jnp.array([0.0, 0.0], dtype=jnp.float64),
    frame_rotation=None,
    z_to_s_fn=lambda z: z - jnp.array([0.0, 0.0], dtype=jnp.float64),
  )
  context_b = SimpleNamespace(
    metric_fn=flat_metric,
    z0=jnp.array([1.0, 0.0], dtype=jnp.float64),
    frame_rotation=None,
    z_to_s_fn=lambda z: z - jnp.array([1.0, 0.0], dtype=jnp.float64),
  )

  s_samples_a, z_samples_a = select_graph_local_neighborhood(
    context_a,
    z_prior_samples,
    epsilon,
    ball_shortlist_size=4,
  )
  s_samples_b, z_samples_b = select_graph_local_neighborhood(
    context_b,
    z_prior_samples,
    epsilon,
    ball_shortlist_size=4,
  )

  assert z_samples_a.shape == (2, 2)
  assert z_samples_b.shape == (1, 2)
  assert np.isfinite(s_samples_a).all()
  assert np.isfinite(s_samples_b).all()
  assert jnp.all(jnp.linalg.norm(s_samples_a, axis=-1) <= epsilon + 1e-12)
  assert jnp.all(jnp.linalg.norm(s_samples_b, axis=-1) <= epsilon + 1e-12)


def test_build_visualization_figure_emits_progress_logs(monkeypatch, caplog):
  def fake_sample_z0s_from_prior(key, n_samples, d):
    del key
    return jnp.zeros((n_samples, d), dtype=jnp.float64)

  def fake_select_graph_local_neighborhood(
    context,
    z_prior_samples,
    epsilon,
    *,
    ball_shortlist_size=None,
    method=None,
  ):
    del z_prior_samples, epsilon, ball_shortlist_size, method
    z0 = jnp.asarray(context.z0)
    z_samples = jnp.stack([z0, z0 + 0.1], axis=0)
    s_samples = jnp.array([[0.0, 0.0], [0.1, 0.1]], dtype=jnp.float64)
    return s_samples, z_samples

  def fake_compute_density_grid(S1, S2, s_to_z_fn, z_to_s_fn, log_pz_fn):
    del s_to_z_fn, z_to_s_fn, log_pz_fn
    return np.zeros_like(S1 + S2)

  def fake_compute_pz_grid(Z1, Z2, log_pz_fn):
    del log_pz_fn
    return np.zeros_like(Z1 + Z2)

  monkeypatch.setattr(disentanglement_neighborhood, "sample_z0s_from_prior", fake_sample_z0s_from_prior)
  monkeypatch.setattr(
    disentanglement_neighborhood,
    "select_graph_local_neighborhood",
    fake_select_graph_local_neighborhood,
  )
  monkeypatch.setattr(disentanglement_neighborhood, "compute_density_grid", fake_compute_density_grid)
  monkeypatch.setattr(disentanglement_neighborhood, "compute_pz_grid", fake_compute_pz_grid)

  contexts = [
    SimpleNamespace(
      z0=jnp.array([0.0, 0.0], dtype=jnp.float64),
      s_to_z_fn=lambda s: s,
      z_to_s_fn=lambda z: z,
      log_pz_fn=_standard_normal_log_pz,
    ),
    SimpleNamespace(
      z0=jnp.array([0.5, -0.25], dtype=jnp.float64),
      s_to_z_fn=lambda s: s,
      z_to_s_fn=lambda z: z,
      log_pz_fn=_standard_normal_log_pz,
    ),
  ]

  with caplog.at_level(logging.INFO, logger=disentanglement_neighborhood.logger.name):
    fig = disentanglement_neighborhood.build_visualization_figure(
      contexts,
      epsilon=0.3,
      n_samples=5,
      key=jax.random.PRNGKey(0),
      ball_shortlist_size=6,
      n_grid=6,
    )

  messages = [record.getMessage() for record in caplog.records]
  assert any("Building visualization figure with num_z0=2, n_samples=5, n_grid=6" in message for message in messages)
  assert any("Sampled 5 prior points for graph neighborhoods" in message for message in messages)
  assert any("Processing neighborhood 1/2" in message for message in messages)
  assert any("Processing neighborhood 2/2" in message for message in messages)
  assert any("Computed density grid for neighborhood 2/2" in message for message in messages)
  assert any("Assembling figure and panels" in message for message in messages)
  plt.close(fig)


def test_main_emits_progress_logs_and_saves(monkeypatch, caplog, tmp_path):
  output_path = tmp_path / "disentanglement_neighborhood.pdf"
  saved_paths = []
  closed_figures = []

  class FakeParser:
    def parse_args(self):
      return SimpleNamespace(
        seed=0,
        num_z0=2,
        epsilon=0.6,
        n_samples=8,
        ball_shortlist_size=12,
        n_grid=10,
        method="ode",
        fixed_decoder=False,
        monge_decoder=None,
        decoder_offset=None,
        d=2,
        n=3,
        decoder_eps=1.0,
        z_plot_radius=2.0,
        n_contour_levels=8,
        figwidth=None,
        figheight=None,
        dpi=200,
        output=str(output_path),
      )

  class FakeFigure:
    def savefig(self, path, bbox_inches, dpi):
      saved_paths.append((path, bbox_inches, dpi))

  def fake_sample_z0s_from_prior(key, num_z0, d):
    del key
    return jnp.arange(num_z0 * d, dtype=jnp.float64).reshape(num_z0, d)

  def fake_build_context_spec_from_args(args, z0):
    del args
    return tuple(float(x) for x in z0)

  def fake_build_context_from_spec(spec):
    return SimpleNamespace(z0=jnp.asarray(spec, dtype=jnp.float64))

  def fake_build_visualization_figure(
    contexts,
    epsilon,
    n_samples,
    key,
    *,
    ball_shortlist_size=None,
    method=None,
    n_grid=80,
    z_plot_radius=2.0,
    n_contour_levels=8,
    figwidth=None,
    figheight=None,
    dpi=200,
  ):
    del epsilon, n_samples, key, ball_shortlist_size, method, n_grid, z_plot_radius, n_contour_levels, figwidth, figheight, dpi
    assert len(contexts) == 2
    return FakeFigure()

  monkeypatch.setattr(disentanglement_neighborhood, "build_parser", lambda: FakeParser())
  monkeypatch.setattr(disentanglement_neighborhood, "sample_z0s_from_prior", fake_sample_z0s_from_prior)
  monkeypatch.setattr(
    disentanglement_neighborhood,
    "build_context_spec_from_args",
    fake_build_context_spec_from_args,
  )
  monkeypatch.setattr(disentanglement_neighborhood, "build_context_from_spec", fake_build_context_from_spec)
  monkeypatch.setattr(
    disentanglement_neighborhood,
    "build_visualization_figure",
    fake_build_visualization_figure,
  )
  monkeypatch.setattr(plt, "close", lambda fig: closed_figures.append(fig))

  with caplog.at_level(logging.INFO, logger=disentanglement_neighborhood.logger.name):
    disentanglement_neighborhood.main()

  messages = [record.getMessage() for record in caplog.records]
  assert any(
    "Running disentanglement neighborhood with num_z0=2, n_samples=8, epsilon=0.600, ball_shortlist_size=12, method=ode" == message
    for message in messages
  )
  assert any("Sampled 2 base points from prior" == message for message in messages)
  assert any("Building decoder context 1/2" == message for message in messages)
  assert any("Building decoder context 2/2" == message for message in messages)
  assert any("Building visualization figure" == message for message in messages)
  assert any("Figure built, saving output" == message for message in messages)
  assert any(f"Saved figure to {output_path}" == message for message in messages)
  assert saved_paths == [(str(output_path), "tight", 200)]
  assert len(closed_figures) == 1
