import argparse
import json
import jax
import jax.numpy as jnp
import matplotlib
import numpy as np
from pathlib import Path

from local_coordinates.metric import get_euclidean_metric, pullback_metric
from local_geometry.experiments.figure1 import (
  build_context,
  build_context_spec_from_args,
  compute_figure_arrays,
  log_ps_via_inverse_map,
  make_fixed_decoder,
  save_archive,
)
from local_geometry.independent_components import riemannian_ica as legacy_riemannian_ica
from local_geometry.paper.figures.analytic_verification import (
  log_ps as legacy_log_ps,
  log_pz,
  make_coord_fns,
)

matplotlib.use("Agg")

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


def _build_fixed_context():
  z0 = jnp.array([0.3, -0.2])
  decoder, _, _, _ = make_fixed_decoder(d=2, n=3, eps=1.0)
  context = build_context(decoder, z0, method="taylor")
  return decoder, z0, context


def _build_legacy_coord_fns(decoder, z0):
  metric = pullback_metric(z0, decoder, get_euclidean_metric(decoder(z0)))
  j_z_to_s = legacy_riemannian_ica(metric, log_pz)
  j_s_to_z = j_z_to_s.get_inverse()
  return make_coord_fns("taylor", z0, j_s_to_z), j_z_to_s, j_s_to_z


def test_frame_rotation_is_orthogonal():
  _, _, context = _build_fixed_context()
  identity = context.frame_rotation.T @ context.frame_rotation
  assert jnp.allclose(identity, jnp.eye(2), atol=1e-10)


def test_disentangled_coord_fns_roundtrip_locally():
  _, _, context = _build_fixed_context()
  s = jnp.array([0.02, -0.015])
  z = context.s_to_z_fn(s)
  s_recovered = context.z_to_s_fn(z)
  assert jnp.allclose(s_recovered, s, atol=1e-6)


def test_log_ps_hessian_is_diagonal_at_origin():
  _, _, context = _build_fixed_context()

  def log_ps_fn(s):
    return log_ps_via_inverse_map(s, context.s_to_z_fn, context.z_to_s_fn)

  hessian = jax.hessian(log_ps_fn)(jnp.zeros(2))
  off_diag = hessian - jnp.diag(jnp.diag(hessian))
  assert jnp.allclose(off_diag, 0.0, atol=1e-8)


def test_compute_figure_arrays_are_finite():
  _, _, context = _build_fixed_context()
  arrays = compute_figure_arrays(context, R=0.6, n_grid=11)

  assert arrays["surface"].shape == (11, 11, 3)
  assert arrays["log_ps_grid"].shape == (11, 11)
  assert arrays["log_rho_grid"].shape == (11, 11)
  assert np.isfinite(arrays["surface"]).all()
  assert np.isfinite(arrays["log_ps_grid"]).all()
  assert np.isfinite(arrays["log_rho_grid"]).all()


def test_new_s_to_z_matches_legacy_taylor_map():
  decoder, z0, context = _build_fixed_context()
  (legacy_s_to_z_fn, _), _, _ = _build_legacy_coord_fns(decoder, z0)

  s = jnp.array([0.03, -0.02])
  z_new = context.s_to_z_fn(s)
  z_old = legacy_s_to_z_fn(s)
  assert jnp.allclose(z_new, z_old, atol=1e-6)


def test_new_log_ps_matches_legacy_change_of_variables():
  decoder, z0, context = _build_fixed_context()
  (legacy_s_to_z_fn, legacy_dz_ds_fn), _, _ = _build_legacy_coord_fns(decoder, z0)

  s = jnp.array([0.03, -0.02])
  log_ps_new = log_ps_via_inverse_map(s, context.s_to_z_fn, context.z_to_s_fn)
  log_ps_old = legacy_log_ps(s, legacy_s_to_z_fn, legacy_dz_ds_fn)
  assert jnp.allclose(log_ps_new, log_ps_old, atol=1e-6)


def test_save_archive_copies_figure_script_and_args(tmp_path):
  figure_path = tmp_path / "figure1.pdf"
  figure_path.write_bytes(b"fake-pdf")

  args = argparse.Namespace(
    seed=15,
    fixed_decoder=True,
    monge_decoder=None,
    decoder_offset=None,
    d=2,
    n=3,
    eps=1.0,
    z0=[0.3, -0.2],
    R=1.0,
    n_grid=80,
    n_curves=7,
    n_curve_pts=100,
    log_map="taylor",
    elev=10,
    azim=-145,
    figwidth=14,
    figheight=10,
    dpi=200,
    n_contour_levels=8,
    n_surface_contours=6,
    output=str(figure_path),
    archive=True,
  )

  archive_dir = Path(save_archive(str(figure_path), args))

  assert archive_dir.exists()
  assert (archive_dir / "figure1.pdf").exists()
  assert (archive_dir / "figure1.py").exists()
  assert (archive_dir / "args.json").exists()

  with open(archive_dir / "args.json") as fp:
    archived_args = json.load(fp)

  assert archived_args["archive"] is True
  assert archived_args["fixed_decoder"] is True


def test_build_context_spec_from_args_uses_shared_names():
  args = argparse.Namespace(
    seed=15,
    fixed_decoder=True,
    monge_decoder=None,
    decoder_offset=[0.1, -0.2],
    d=2,
    n=3,
    eps=1.0,
    z0=[0.3, -0.2],
    log_map="taylor",
  )

  spec = build_context_spec_from_args(args)

  assert spec.decoder_name == "fixed"
  assert spec.prior_name == "standard_normal"
  assert spec.z0 == (0.3, -0.2)
  assert spec.decoder_offset == (0.1, -0.2)
  assert spec.method == "taylor"
