import jax
import jax.numpy as jnp
import pytest

blackjax = pytest.importorskip("blackjax")

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

from local_geometry.paper.figures.dhsic_independence import compute_log_weights
from local_geometry.paper.figures.neighborhood_independence import sample_uniform_region
from local_geometry.rica_jax.rnc_nbhd_sample import rnc_nbhd_sample
from local_geometry.rica_jax.rnc_pushforward import ODEMethod, s_to_z


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 _rotation(theta):
  return jnp.array([
    [jnp.cos(theta), -jnp.sin(theta)],
    [jnp.sin(theta), jnp.cos(theta)],
  ])


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


def _build_coord_fns(metric_fn, z0, method=None, frame_rotation=None):
  def s_to_z_fn(s):
    return s_to_z(
      metric_fn,
      z0,
      s,
      method=method,
      frame_rotation=frame_rotation,
    )

  return s_to_z_fn, jax.jacobian(s_to_z_fn)


def _weighted_mean(samples, log_weights):
  weights = jax.nn.softmax(log_weights)
  return jnp.sum(samples * weights[:, None], axis=0)


def _weighted_second_moment(samples, log_weights):
  weights = jax.nn.softmax(log_weights)
  squared_norms = jnp.sum(samples ** 2, axis=-1)
  return jnp.sum(weights * squared_norms)


def test_flat_metric_samples_stay_in_ball_and_are_finite():
  z0 = jnp.array([0.2, -0.1], dtype=jnp.float64)
  epsilon = 0.75
  key = jax.random.PRNGKey(0)

  samples = rnc_nbhd_sample(
    _flat_metric,
    z0,
    epsilon,
    256,
    key,
    log_pz_fn=_standard_normal_log_pz,
    proposal_scale=0.18,
    n_warmup=200,
  )

  assert samples.shape == (256, 2)
  assert jnp.isfinite(samples).all()
  assert jnp.all(jnp.linalg.norm(samples, axis=-1) <= epsilon + 1e-12)


def test_flat_metric_nuts_samples_stay_in_ball_and_center_near_zero():
  z0 = jnp.array([0.2, -0.1], dtype=jnp.float64)
  epsilon = 0.75
  key = jax.random.PRNGKey(11)

  samples = rnc_nbhd_sample(
    _flat_metric,
    z0,
    epsilon,
    256,
    key,
    log_pz_fn=_standard_normal_log_pz,
    sampler="nuts",
    proposal_scale=0.1,
    n_warmup=200,
  )

  assert samples.shape == (256, 2)
  assert jnp.isfinite(samples).all()
  assert jnp.all(jnp.linalg.norm(samples, axis=-1) <= epsilon + 1e-12)
  assert jnp.allclose(jnp.mean(samples, axis=0), jnp.zeros(2), atol=0.15)


def test_flat_metric_rotation_leaves_samples_unchanged():
  z0 = jnp.array([0.0, 0.0], dtype=jnp.float64)
  epsilon = 0.8
  key = jax.random.PRNGKey(1)
  frame_rotation = _rotation(0.41)

  samples = rnc_nbhd_sample(
    _flat_metric,
    z0,
    epsilon,
    128,
    key,
    log_pz_fn=_standard_normal_log_pz,
    proposal_scale=0.2,
    n_warmup=150,
  )
  rotated_samples = rnc_nbhd_sample(
    _flat_metric,
    z0,
    epsilon,
    128,
    key,
    log_pz_fn=_standard_normal_log_pz,
    frame_rotation=frame_rotation,
    proposal_scale=0.2,
    n_warmup=150,
  )

  assert jnp.allclose(samples, rotated_samples)


def test_curved_metric_runs_with_taylor_and_ode_methods():
  z0 = jnp.array([0.17, -0.23], dtype=jnp.float64)
  epsilon = 0.3
  key_taylor, key_ode = jax.random.split(jax.random.PRNGKey(2))

  samples_taylor = rnc_nbhd_sample(
    _curved_metric,
    z0,
    epsilon,
    64,
    key_taylor,
    log_pz_fn=_standard_normal_log_pz,
    proposal_scale=0.08,
    n_warmup=100,
  )
  samples_ode = rnc_nbhd_sample(
    _curved_metric,
    z0,
    epsilon,
    64,
    key_ode,
    method=ODEMethod(max_iters=8, tol=1e-6, num_steps=64),
    log_pz_fn=_standard_normal_log_pz,
    proposal_scale=0.08,
    n_warmup=50,
  )

  assert samples_taylor.shape == (64, 2)
  assert samples_ode.shape == (64, 2)
  assert jnp.isfinite(samples_taylor).all()
  assert jnp.isfinite(samples_ode).all()
  assert jnp.all(jnp.linalg.norm(samples_taylor, axis=-1) <= epsilon + 1e-12)
  assert jnp.all(jnp.linalg.norm(samples_ode, axis=-1) <= epsilon + 1e-12)


def test_curved_metric_ode_nuts_sampler_runs_on_small_chain():
  z0 = jnp.array([0.17, -0.23], dtype=jnp.float64)
  epsilon = 0.3
  key = jax.random.PRNGKey(21)

  samples = rnc_nbhd_sample(
    _curved_metric,
    z0,
    epsilon,
    8,
    key,
    method=ODEMethod(max_iters=8, tol=1e-6, num_steps=64),
    log_pz_fn=_standard_normal_log_pz,
    sampler="nuts",
    proposal_scale=0.08,
    n_warmup=10,
  )

  assert samples.shape == (8, 2)
  assert jnp.isfinite(samples).all()
  assert jnp.all(jnp.linalg.norm(samples, axis=-1) <= epsilon + 1e-12)


def test_mcmc_matches_weighted_uniform_reference_on_small_problem():
  z0 = jnp.array([0.12, -0.08], dtype=jnp.float64)
  epsilon = 0.28
  key_chain, key_uniform = jax.random.split(jax.random.PRNGKey(3))

  samples = rnc_nbhd_sample(
    _curved_metric,
    z0,
    epsilon,
    1500,
    key_chain,
    log_pz_fn=_standard_normal_log_pz,
    proposal_scale=0.06,
    n_warmup=400,
  )

  s_uniform = sample_uniform_region(
    key_uniform,
    25_000,
    z0.shape[0],
    epsilon,
    use_sphere=True,
  )
  s_to_z_fn, dz_ds_fn = _build_coord_fns(_curved_metric, z0)
  log_weights = compute_log_weights(
    s_uniform,
    s_to_z_fn,
    dz_ds_fn,
  )

  reference_mean = _weighted_mean(s_uniform, log_weights)
  reference_second_moment = _weighted_second_moment(s_uniform, log_weights)
  chain_mean = jnp.mean(samples, axis=0)
  chain_second_moment = jnp.mean(jnp.sum(samples ** 2, axis=-1))

  assert jnp.allclose(chain_mean, reference_mean, atol=0.08)
  assert jnp.allclose(chain_second_moment, reference_second_moment, atol=0.05)
