"""
Numerical verification of the geodesic Taylor expansion derived in
research_journal/feb_15/notes_feb_15.md.

The main result is that for f: R^n -> R with Monge metric g = I + grad_f grad_f^T,
the Taylor expansion of f along a geodesic gamma(t) with gamma(0) = x and
dot_gamma(0) = V is:

  f(gamma_t) = f(x) + t * grad_f^T V + (t^2/2) * V^T H_f V / (1 + ||grad_f||^2) + O(t^3)

where H_f is the Euclidean Hessian, and the second-order coefficient is the
covariant Hessian of f with respect to the Monge metric.
"""
import jax
import jax.numpy as jnp
import pytest
from local_coordinates.monge import get_monge_metric
from local_coordinates.connection import get_levi_civita_connection
from local_coordinates.exponential_map import exponential_map_ode
from local_coordinates.basis import get_standard_basis
from local_coordinates.jet import Jet
from local_coordinates.tangent import TangentVector


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_tangent_vector(x, v_components):
  """Build a TangentVector from raw arrays."""
  dim = x.shape[0]
  basis = get_standard_basis(x)
  v_jet = Jet(
    value=v_components,
    gradient=jnp.zeros((dim, dim)),
    hessian=jnp.zeros((dim, dim, dim)),
  )
  return TangentVector(p=x, components=v_jet, basis=basis)


def _analytic_first_deriv(f, x, V):
  """grad_f(x)^T V"""
  grad_f = jax.grad(f)(x)
  return jnp.dot(grad_f, V)


def _analytic_second_deriv(f, x, V):
  """V^T H_f V / (1 + ||grad_f||^2)"""
  grad_f = jax.grad(f)(x)
  hess_f = jax.hessian(f)(x)
  norm_sq = jnp.dot(grad_f, grad_f)
  return jnp.dot(V, hess_f @ V) / (1.0 + norm_sq)


def _covariant_hessian_from_christoffel(f, x):
  """Compute (nabla^2 f)_{ij} = d^2_{ij} f - Gamma^k_{ij} d_k f using the library."""
  metric = get_monge_metric(f, x)
  connection = get_levi_civita_connection(metric)
  Gamma = connection.christoffel_symbols.value  # [i, j, k] = Gamma^k_{ij}
  grad_f = jax.grad(f)(x)
  hess_f = jax.hessian(f)(x)
  correction = jnp.einsum("ijk,k->ij", Gamma, grad_f)
  return hess_f - correction


def _covariant_hessian_closed_form(f, x):
  """H_f / (1 + ||grad_f||^2)"""
  grad_f = jax.grad(f)(x)
  hess_f = jax.hessian(f)(x)
  norm_sq = jnp.dot(grad_f, grad_f)
  return hess_f / (1.0 + norm_sq)


# ---------------------------------------------------------------------------
# Test functions and points
# ---------------------------------------------------------------------------

def _quadratic_f(x):
  A = jnp.array([[2.0, 0.7], [0.7, 5.0]])
  return 0.5 * x @ A @ x


def _nonlinear_f(x):
  return jnp.sin(x[0]) * jnp.cos(x[1]) + 0.3 * x[0] ** 2


QUADRATIC_POINT = jnp.array([1.0, -0.5])
NONLINEAR_POINT = jnp.array([0.7, -0.4])


# ---------------------------------------------------------------------------
# Tests: Covariant Hessian closed-form vs Christoffel computation
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("f, x", [
  (_quadratic_f, QUADRATIC_POINT),
  (_nonlinear_f, NONLINEAR_POINT),
])
def test_covariant_hessian_closed_form(f, x):
  """The closed-form H_f/(1+||grad_f||^2) matches the Christoffel-based computation."""
  cov_hess_christoffel = _covariant_hessian_from_christoffel(f, x)
  cov_hess_closed = _covariant_hessian_closed_form(f, x)
  assert jnp.allclose(cov_hess_christoffel, cov_hess_closed, atol=1e-10), (
    f"max diff = {jnp.max(jnp.abs(cov_hess_christoffel - cov_hess_closed))}"
  )


# ---------------------------------------------------------------------------
# Tests: First derivative along geodesic
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("f, x", [
  (_quadratic_f, QUADRATIC_POINT),
  (_nonlinear_f, NONLINEAR_POINT),
])
def test_first_derivative_along_geodesic(f, x):
  """
  Finite-difference d/dt f(gamma_t) at t=0 matches the analytic grad_f^T V.

  We integrate a geodesic and use a forward finite difference on the first
  few trajectory points to estimate the derivative.
  """
  dim = x.shape[0]
  V = jnp.array([0.6, 0.8])
  V = V / jnp.linalg.norm(V)

  metric_fn = lambda z: get_monge_metric(f, z)
  scale = 0.01
  v_tan = _make_tangent_vector(x, scale * V)

  ts, trajectory = exponential_map_ode(
    v_tan, metric_fn=metric_fn, num_steps=500, return_trajectory=True,
  )

  f_vals = jax.vmap(f)(trajectory)

  # Forward finite difference at t=0 (the ts are multiplied by scale since
  # we passed scale*V as the velocity)
  dt = float(ts[1] - ts[0]) * scale
  fd_first = float((f_vals[1] - f_vals[0]) / dt)

  analytic = float(_analytic_first_deriv(f, x, V))
  assert abs(fd_first - analytic) < 1e-3, (
    f"FD first deriv = {fd_first}, analytic = {analytic}, diff = {abs(fd_first - analytic)}"
  )


# ---------------------------------------------------------------------------
# Tests: Second derivative along geodesic
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("f, x", [
  (_quadratic_f, QUADRATIC_POINT),
  (_nonlinear_f, NONLINEAR_POINT),
])
def test_second_derivative_along_geodesic(f, x):
  """
  Finite-difference d^2/dt^2 f(gamma_t) at t=0 matches V^T H_f V / (1+||grad_f||^2).

  We use f(gamma_t) = f(x) + t*a1 + t^2/2*a2 + O(t^3) and fit a2 from the
  trajectory data rather than relying on a central difference, which is
  more stable for second derivatives.
  """
  dim = x.shape[0]
  V = jnp.array([0.6, 0.8])
  V = V / jnp.linalg.norm(V)

  metric_fn = lambda z: get_monge_metric(f, z)
  scale = 0.01
  v_tan = _make_tangent_vector(x, scale * V)

  ts, trajectory = exponential_map_ode(
    v_tan, metric_fn=metric_fn, num_steps=500, return_trajectory=True,
  )

  f_vals = jax.vmap(f)(trajectory)
  f0 = f_vals[0]
  a1_analytic = _analytic_first_deriv(f, x, V)

  # The actual time along the geodesic is t_actual = ts * scale
  # since we passed scale*V. So f(gamma(t)) = f(x) + t*a1 + t^2/2*a2 + ...
  # where t = ts * scale.
  # Rearranging: (f_vals - f0 - t*a1) / (t^2/2) = a2 + O(t)
  t_actual = ts * scale
  residual = f_vals - f0 - t_actual * a1_analytic
  # Use points near t=0 (but not t=0 itself) to estimate a2.
  # Keep the window small to avoid third-order contamination.
  mask = (t_actual > 1e-6) & (t_actual < 0.003)
  t_masked = t_actual[mask]
  residual_masked = residual[mask]
  a2_estimates = 2.0 * residual_masked / (t_masked ** 2)
  a2_fd = float(jnp.mean(a2_estimates))

  a2_analytic = float(_analytic_second_deriv(f, x, V))
  assert abs(a2_fd - a2_analytic) / (abs(a2_analytic) + 1e-12) < 0.02, (
    f"FD second deriv = {a2_fd}, analytic = {a2_analytic}, "
    f"rel diff = {abs(a2_fd - a2_analytic) / (abs(a2_analytic) + 1e-12)}"
  )


# ---------------------------------------------------------------------------
# Test: Full Taylor expansion accuracy
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("f, x", [
  (_quadratic_f, QUADRATIC_POINT),
  (_nonlinear_f, NONLINEAR_POINT),
])
def test_taylor_expansion_accuracy(f, x):
  """
  The second-order Taylor expansion should match f(gamma_t) to O(t^3) for small t.
  """
  dim = x.shape[0]
  V = jnp.array([0.6, 0.8])
  V = V / jnp.linalg.norm(V)

  metric_fn = lambda z: get_monge_metric(f, z)
  scale = 0.05
  v_tan = _make_tangent_vector(x, scale * V)

  ts, trajectory = exponential_map_ode(
    v_tan, metric_fn=metric_fn, num_steps=500, return_trajectory=True,
  )

  f_vals = jax.vmap(f)(trajectory)
  f0 = f(x)
  a1 = _analytic_first_deriv(f, x, V)
  a2 = _analytic_second_deriv(f, x, V)

  t_actual = ts * scale
  taylor_approx = f0 + t_actual * a1 + 0.5 * t_actual ** 2 * a2

  # At t ~ 0.01, the O(t^3) error should be very small
  mask = (t_actual > 1e-6) & (t_actual < 0.015)
  errors = jnp.abs(f_vals[mask] - taylor_approx[mask])
  max_error = float(jnp.max(errors))
  assert max_error < 1e-5, f"Max Taylor error = {max_error} (should be O(t^3))"


# ---------------------------------------------------------------------------
# Helpers for principal direction tests
# ---------------------------------------------------------------------------

def _get_principal_directions(f, x):
  """
  Compute the principal directions via the sequential greedy procedure.

  Returns (V, mus) where V is (n, n) with V[k] = k-th principal direction
  and mus is (n,) with the associated eigenvalues of the projected Hessian
  (the first entry is ||grad_f|| for the gradient direction).
  """
  n = x.shape[0]
  grad_f = jax.grad(f)(x)
  hess_f = jax.hessian(f)(x)
  norm_g = jnp.linalg.norm(grad_f)

  g_hat = grad_f / norm_g
  P = jnp.eye(n) - jnp.outer(g_hat, g_hat)

  projected_hess = P @ hess_f @ P

  evals, evecs = jnp.linalg.eigh(projected_hess)
  # eigh returns ascending eigenvalues. The projected Hessian has one
  # zero eigenvalue (the gradient direction). Sort the rest by |eval|
  # descending to get the principal directions.
  abs_evals = jnp.abs(evals)
  order = jnp.argsort(-abs_evals)
  evals = evals[order]
  evecs = evecs[:, order]

  # The first column of evecs after sorting by |eval| might be the
  # gradient direction (zero eigenvalue) or a non-trivial direction.
  # We build the full frame: V_1 = g_hat, then the (n-1) eigenvectors
  # with non-negligible eigenvalue.
  # Since P has rank (n-1), exactly one eigenvalue is ~0.
  # We take the (n-1) eigenvectors with |eval| > threshold.
  V = jnp.zeros((n, n))
  V = V.at[0].set(g_hat)
  mus = jnp.zeros(n)
  mus = mus.at[0].set(norm_g)

  idx = 1
  for k in range(n):
    if jnp.abs(evals[k]) > 1e-10:
      V = V.at[idx].set(evecs[:, k])
      mus = mus.at[idx].set(evals[k])
      idx += 1

  return V, mus


# ---------------------------------------------------------------------------
# 3D test function where grad_f is NOT an eigenvector of H_f
# ---------------------------------------------------------------------------

def _3d_f(x):
  return jnp.sin(x[0]) * jnp.cos(x[1]) + 0.3 * x[0] ** 2 + x[2] ** 2 + 0.5 * x[0] * x[2]


THREE_D_POINT = jnp.array([0.7, -0.4, 0.3])


# ---------------------------------------------------------------------------
# Tests: Principal directions from projected Hessian
# ---------------------------------------------------------------------------

def test_principal_directions_orthonormal():
  """The principal directions form an orthonormal set."""
  V, _ = _get_principal_directions(_3d_f, THREE_D_POINT)
  gram = V @ V.T
  assert jnp.allclose(gram, jnp.eye(3), atol=1e-10), (
    f"Gram matrix not identity, max diff = {jnp.max(jnp.abs(gram - jnp.eye(3)))}"
  )


def test_principal_directions_first_is_gradient():
  """The first principal direction is the normalized gradient."""
  V, _ = _get_principal_directions(_3d_f, THREE_D_POINT)
  grad_f = jax.grad(_3d_f)(THREE_D_POINT)
  g_hat = grad_f / jnp.linalg.norm(grad_f)
  assert jnp.allclose(jnp.abs(jnp.dot(V[0], g_hat)), 1.0, atol=1e-10)


def test_principal_directions_orthogonal_to_gradient():
  """The remaining principal directions are orthogonal to the gradient."""
  V, _ = _get_principal_directions(_3d_f, THREE_D_POINT)
  grad_f = jax.grad(_3d_f)(THREE_D_POINT)
  for k in range(1, 3):
    dot = jnp.dot(V[k], grad_f)
    assert jnp.abs(dot) < 1e-10, f"V_{k} not orthogonal to grad_f, dot = {dot}"


def test_principal_directions_are_projected_hessian_eigenvectors():
  """
  V_2, ..., V_n satisfy P H_f P @ V_k = mu_k V_k, confirming they are
  eigenvectors of the projected Hessian.
  """
  x = THREE_D_POINT
  V, mus = _get_principal_directions(_3d_f, x)
  n = x.shape[0]
  grad_f = jax.grad(_3d_f)(x)
  hess_f = jax.hessian(_3d_f)(x)
  g_hat = grad_f / jnp.linalg.norm(grad_f)
  P = jnp.eye(n) - jnp.outer(g_hat, g_hat)
  projected_hess = P @ hess_f @ P

  for k in range(1, n):
    lhs = projected_hess @ V[k]
    rhs = mus[k] * V[k]
    assert jnp.allclose(lhs, rhs, atol=1e-10), (
      f"V_{k} not an eigenvector of P H_f P, max diff = {jnp.max(jnp.abs(lhs - rhs))}"
    )


def test_principal_directions_differ_from_hessian_eigenvectors():
  """
  At a non-critical point where grad_f is not an eigenvector of H_f, the
  projected Hessian eigenvectors should differ from the eigenvectors of H_f.
  """
  x = THREE_D_POINT
  grad_f = jax.grad(_3d_f)(x)
  hess_f = jax.hessian(_3d_f)(x)

  # Verify that grad_f is NOT an eigenvector of H_f (precondition for this test)
  Hg = hess_f @ grad_f
  # If grad_f were an eigenvector, Hg would be parallel to grad_f
  cross = Hg / jnp.linalg.norm(Hg) - grad_f / jnp.linalg.norm(grad_f)
  assert jnp.linalg.norm(cross) > 0.01, "grad_f is an eigenvector of H_f at this point"

  # Now check that the projected eigenvectors differ from H_f eigenvectors
  V_proj, _ = _get_principal_directions(_3d_f, x)
  evals_h, evecs_h = jnp.linalg.eigh(hess_f)

  # The projected directions V_proj[1:] should NOT be aligned with any
  # single eigenvector of H_f. Check that the maximum alignment is < 1.
  for k in range(1, 3):
    alignments = jnp.abs(evecs_h.T @ V_proj[k])
    max_alignment = jnp.max(alignments)
    assert max_alignment < 1.0 - 1e-6, (
      f"V_{k} is aligned with an eigenvector of H_f (max alignment = {max_alignment})"
    )


def test_principal_directions_at_critical_point():
  """
  At a critical point (grad_f = 0), P = I and the principal directions
  reduce to the eigenvectors of H_f sorted by absolute eigenvalue.
  """
  def f_quad(x):
    A = jnp.array([[2.0, 0.7, 0.0], [0.7, 5.0, 0.0], [0.0, 0.0, 3.0]])
    return 0.5 * x @ A @ x

  x = jnp.zeros(3)
  hess_f = jax.hessian(f_quad)(x)
  evals_h, evecs_h = jnp.linalg.eigh(hess_f)

  # At the origin, grad_f = 0 and the Monge metric is the identity.
  # The principal directions should be the eigenvectors of H_f.
  # Since grad_f = 0, the greedy procedure's first step is degenerate
  # (all directions have zero first-order contribution), so ALL directions
  # come from the Hessian. We verify this by checking that the covariant
  # Hessian equals H_f and its eigenvectors span the same space.
  cov_hess = _covariant_hessian_closed_form(f_quad, x)
  assert jnp.allclose(cov_hess, hess_f, atol=1e-12), (
    "Covariant Hessian should equal H_f at critical point"
  )

  evals_c, evecs_c = jnp.linalg.eigh(cov_hess)
  # Eigenvectors should span the same eigenspaces
  for k in range(3):
    alignment = jnp.abs(jnp.dot(evecs_c[:, k], evecs_h[:, k]))
    assert jnp.allclose(alignment, 1.0, atol=1e-10), (
      f"Eigenvector {k} mismatch at critical point, alignment = {alignment}"
    )
