"""Tests for the interactive analytic verification state and computation logic.

Verifies that:
  - InteractiveState initializes for every decoder type and method
  - Tier 1 (rebuild_decoder) and Tier 2 (recompute_grids) produce valid data
  - JIT-compiled batch functions agree with their non-JIT counterparts
  - Grid-only updates preserve decoder identity
  - draw_panels renders without error
"""
import jax
import jax.numpy as jnp
import numpy as np
import pytest
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from local_geometry.paper.figures.interactive_analytic_verification import (
  InteractiveState,
  draw_panels,
  ParamControl,
  DECODER_CHOICES,
  METHOD_CHOICES,
  N_GRID,
  N_CURVES,
)
from local_geometry.paper.figures.analytic_verification import (
  map_grid_to_z,
  map_grid_to_surface,
  compute_density_grid,
  compute_rho_grid,
  build_rnc_grid,
  compute_orientation_arrows,
  plot_orientation_arrows,
)

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


@pytest.fixture
def default_state():
  return InteractiveState()


class TestStateInit:
  """InteractiveState should initialize for all decoder/method combos."""

  @pytest.mark.parametrize("decoder_type", DECODER_CHOICES)
  def test_init_decoder_types(self, decoder_type):
    state = InteractiveState(decoder_type=decoder_type)
    assert state.f is not None
    assert state.s_to_z_fn is not None
    assert "surf" in state.data
    assert state.data["surf"].shape == (N_GRID, N_GRID, 3)

  @pytest.mark.parametrize("method", ["taylor", "improved"])
  def test_init_methods(self, method):
    state = InteractiveState(method=method)
    assert state.data["surf"].shape[0] == N_GRID


class TestTierOneRebuild:
  """Tier 1 rebuilds should produce valid functions and JIT wrappers."""

  def test_rebuild_changes_decoder(self, default_state):
    f_z0_before = default_state.data["f_z0"].copy()
    default_state.decoder_type = "fixed"
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    assert not np.allclose(f_z0_before, default_state.data["f_z0"])

  def test_rebuild_with_z0_change(self, default_state):
    f_z0_before = default_state.data["f_z0"].copy()
    default_state.z0 = jnp.array([0.5, 0.5])
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    assert not np.allclose(f_z0_before, default_state.data["f_z0"])

  def test_rebuild_with_eps_change(self, default_state):
    f_z0_before = default_state.data["f_z0"].copy()
    default_state.eps = 0.1
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    assert not np.allclose(f_z0_before, default_state.data["f_z0"])


class TestTierTwoGrids:
  """Tier 2 grid recomputation should produce correct shapes."""

  def test_grid_uses_fixed_constants(self, default_state):
    assert default_state.data["surf"].shape == (N_GRID, N_GRID, 3)
    assert default_state.data["Z1"].shape == (N_GRID, N_GRID)
    n_total = N_CURVES + N_CURVES
    assert len(default_state.data["mapped_curves_z"]) == n_total

  def test_R_change(self, default_state):
    default_state.R = 1.5
    default_state.recompute_grids()
    assert default_state.data["S1"].max() == pytest.approx(1.5)

  def test_grid_preserves_decoder(self, default_state):
    f_z0_before = default_state.data["f_z0"].copy()
    default_state.R = 1.5
    default_state.recompute_grids()
    np.testing.assert_allclose(f_z0_before, default_state.data["f_z0"])


class TestJitConsistency:
  """JIT-compiled batch functions should agree with non-JIT versions."""

  def test_grid_to_z_matches(self, default_state):
    S1, S2 = build_rnc_grid(default_state.R, N_GRID)
    Z1_ref, Z2_ref = map_grid_to_z(S1, S2, default_state.s_to_z_fn)
    np.testing.assert_allclose(default_state.data["Z1"], Z1_ref, atol=1e-12)
    np.testing.assert_allclose(default_state.data["Z2"], Z2_ref, atol=1e-12)

  def test_density_matches(self, default_state):
    S1, S2 = build_rnc_grid(default_state.R, N_GRID)
    log_ps_ref = compute_density_grid(
      S1, S2, default_state.s_to_z_fn, default_state.dz_ds_fn,
    )
    np.testing.assert_allclose(
      default_state.data["log_ps_grid"], log_ps_ref, atol=1e-10,
    )

  def test_rho_matches(self, default_state):
    S1, S2 = build_rnc_grid(default_state.R, N_GRID)
    Z1_ref, Z2_ref = map_grid_to_z(S1, S2, default_state.s_to_z_fn)
    log_rho_ref = compute_rho_grid(Z1_ref, Z2_ref, default_state.log_det_g_fn)
    np.testing.assert_allclose(
      default_state.data["log_rho_grid"], log_rho_ref, atol=1e-10,
    )


class TestParamControl:
  """ParamControl should clamp, format, and fire callbacks correctly."""

  def test_increase_fires_callback(self):
    fig = plt.figure()
    received = []
    pc = ParamControl(fig, 0.5, "test", 1.0, 0.1, 0.0, 2.0, fmt=".1f")
    pc.on_changed(lambda v: received.append(v))
    pc._on_increase(None)
    assert len(received) == 1
    assert received[0] == pytest.approx(1.1)
    plt.close(fig)

  def test_decrease_fires_callback(self):
    fig = plt.figure()
    received = []
    pc = ParamControl(fig, 0.5, "test", 1.0, 0.1, 0.0, 2.0, fmt=".1f")
    pc.on_changed(lambda v: received.append(v))
    pc._on_decrease(None)
    assert len(received) == 1
    assert received[0] == pytest.approx(0.9)
    plt.close(fig)

  def test_clamp_at_min(self):
    fig = plt.figure()
    pc = ParamControl(fig, 0.5, "test", 0.0, 0.1, 0.0, 2.0, fmt=".1f")
    pc._on_decrease(None)
    assert pc.value == 0.0
    plt.close(fig)

  def test_clamp_at_max(self):
    fig = plt.figure()
    pc = ParamControl(fig, 0.5, "test", 2.0, 0.1, 0.0, 2.0, fmt=".1f")
    pc._on_increase(None)
    assert pc.value == 2.0
    plt.close(fig)

  def test_int_step(self):
    fig = plt.figure()
    pc = ParamControl(fig, 0.5, "test", 5, 1, 0, 10)
    pc._on_increase(None)
    assert pc.value == 6
    assert isinstance(pc.value, int)
    plt.close(fig)

  def test_reset(self):
    fig = plt.figure()
    pc = ParamControl(fig, 0.5, "test", 1.0, 0.1, 0.0, 2.0, fmt=".1f")
    pc._on_increase(None)
    pc._on_increase(None)
    assert pc.value == pytest.approx(1.2)
    pc.reset(1.0)
    assert pc.value == pytest.approx(1.0)
    plt.close(fig)


class TestSeedUpdatePath:
  """Simulates the exact callback path triggered by a seed +/- click."""

  def test_seed_increment_rebuild_redraw(self, default_state):
    f_z0_before = default_state.data["f_z0"].copy()
    default_state.seed = 16
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    assert not np.allclose(f_z0_before, default_state.data["f_z0"])
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    draw_panels(fig, ax_surf, ax_z, ax_s, default_state)
    plt.close(fig)

  def test_seed_change_with_contours_off(self, default_state):
    default_state.seed = 16
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    draw_panels(fig, ax_surf, ax_z, ax_s, default_state,
                show_surface_contours=False)
    plt.close(fig)


class TestDrawPanels:
  """draw_panels should render without errors."""

  def test_draw_default(self, default_state):
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    draw_panels(fig, ax_surf, ax_z, ax_s, default_state)
    plt.close(fig)

  def test_draw_with_contours(self, default_state):
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    draw_panels(fig, ax_surf, ax_z, ax_s, default_state,
                show_surface_contours=True)
    plt.close(fig)

  def test_draw_after_rebuild(self, default_state):
    default_state.decoder_type = "monkey"
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    draw_panels(fig, ax_surf, ax_z, ax_s, default_state)
    plt.close(fig)


class TestOrientationArrows:
  """Orientation arrows should be computed, stored, and drawn correctly."""

  def test_arrows_in_data(self, default_state):
    assert "arrows" in default_state.data
    arrows = default_state.data["arrows"]
    for key in ("s_e1", "s_e2", "z_e1", "z_e2", "x_e1", "x_e2"):
      assert key in arrows

  def test_arrow_shapes(self, default_state):
    arrows = default_state.data["arrows"]
    for key in ("s_e1", "s_e2", "z_e1", "z_e2"):
      assert arrows[key].shape == (2,)
    for key in ("x_e1", "x_e2"):
      assert arrows[key].shape == (3,)

  def test_rnc_arrows_axis_aligned(self, default_state):
    arrows = default_state.data["arrows"]
    np.testing.assert_allclose(arrows["s_e1"][1], 0.0, atol=1e-12)
    np.testing.assert_allclose(arrows["s_e2"][0], 0.0, atol=1e-12)

  def test_arrow_lengths_are_unit(self, default_state):
    arrows = default_state.data["arrows"]
    for key in ("s_e1", "s_e2", "z_e1", "z_e2", "x_e1", "x_e2"):
      length = np.linalg.norm(arrows[key])
      np.testing.assert_allclose(length, 1.0, rtol=1e-6)

  def test_arrows_change_with_decoder(self, default_state):
    arrows_before = {k: v.copy() for k, v in default_state.data["arrows"].items()}
    default_state.decoder_type = "saddle"
    default_state.rebuild_decoder()
    default_state.recompute_grids()
    arrows_after = default_state.data["arrows"]
    changed = any(
      not np.allclose(arrows_before[k], arrows_after[k])
      for k in ("z_e1", "z_e2", "x_e1", "x_e2")
    )
    assert changed

  def test_compute_orientation_arrows_directly(self):
    state = InteractiveState(decoder_type="fixed")
    arrows = compute_orientation_arrows(state.z0, state.f, state.dz_ds_fn)
    for key in ("s_e1", "s_e2", "z_e1", "z_e2", "x_e1", "x_e2"):
      length = np.linalg.norm(arrows[key])
      np.testing.assert_allclose(length, 1.0, rtol=1e-6)

  def test_plot_orientation_arrows_runs(self, default_state):
    fig = plt.figure(figsize=(16, 14))
    ax_surf = fig.add_axes([0.1, 0.5, 0.8, 0.45], projection="3d")
    ax_z = fig.add_axes([0.05, 0.25, 0.2, 0.2])
    ax_s = fig.add_axes([0.75, 0.25, 0.2, 0.2])
    d = default_state.data
    plot_orientation_arrows(ax_s, ax_z, ax_surf, d["z0_np"], d["f_z0"],
                            d["arrows"])
    plt.close(fig)
