r"""Validates ``DegenerateVariationalIntegratorFlow`` (DVI) as a plain
numerical integrator -- no trainable network, no gradient, no Projector
(see ``implicit_vs_explicit_euler.py``/``adams_bashforth_multistep.py``).

The DVI is built directly from the discrete Euler-Lagrange equations of a
properly-degenerate Lagrangian L(q, p, qdot) = theta(q, p) . qdot - H(q, p)
(eq. 10 in Franck, *Learning non-canonical Hamiltonian ODEs*), NOT by
discretizing the continuous vector field F = W^-1 grad H (see
``LagrangianDegenerateVectorFieldSpace`` in ``vector_fields.py`` for that
continuous F, used below only as an independent reference/starter).

Test system -- a non-canonical rewriting of the harmonic oscillator:
    theta(q, p) = p + eps*q     (eps != 0: "properly degenerate")
    H(q, p)     = 0.5*(p^2 + q^2)
which drives qdot = p/eps, pdot = -q/eps (still exactly energy-conserving,
just at angular frequency 1/eps instead of 1).

This script checks the two properties that matter for a variational
integrator (NOT raw trajectory accuracy, which a low-order symplectic
method trades away on purpose):
  1. Local truncation order: comparing one DVI step (from an
     RK4-bootstrapped window) against the known analytic solution as dt
     shrinks. The scheme mixes theta/grad H evaluated at n in the first
     discrete Euler-Lagrange equation with theta/grad H at n+1 in the
     second (see ``DegenerateVariationalIntegratorFlow``'s docstring) --
     the same asymmetric-time structure as symplectic Euler, hence the
     same first order (NOT the second order of, say, ``AdamsBashforth2Flow``).
  2. Long-time energy conservation: bounded oscillation around H0 over
     thousands of steps, unlike a naive (non-symplectic) explicit Euler
     applied to the same continuous F, whose energy grows without bound.

Caveat (Remark 3 in the paper): unlike the continuous F, this DISCRETE
scheme is sensitive to the gauge of theta -- replacing theta(q, p) with
theta(q, p) + g(q) for a g with symmetric Jacobian leaves the continuous
dynamics unchanged but can noticeably change (10)'s solution/stability.
The theta above is a simple, well-behaved gauge; not every valid theta for
a given system is (e.g. the ``theta = -log(q)/p`` gauge used for the
continuous flow in ``lotka_volterra_noncanonical.py`` is a poor gauge for
this discrete scheme on that system).
"""

# %%
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (  # noqa: E501
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.approximation_spaces.flow_approximation_spaces import (  # noqa: E501
    FlowsApproximationSpace,
    MultistepPhaseSpaceApproxSpace,
)
from scimba_jax.ode_approx.basic_discrete_ode_nets import (
    ExplicitEulerFlow,
    Rk4Flow,
    bootstrap_multistep_window,
)
from scimba_jax.ode_approx.symplec_discrete_ode_nets import (
    DegenerateVariationalIntegratorFlow,
)
from scimba_jax.ode_approx.vector_fields import (
    LagrangianDegenerateVectorFieldSpace,
)
from scimba_jax.utils.scimba_pytree import ScimbaPytree

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

# %% Non-canonical harmonic oscillator: theta(q,p) = p + eps*q, H = 0.5*(p^2+q^2)

EPS = 0.3


class AnalyticNet(ScimbaPytree):
    """Plain analytic callable wrapped as a (parameter-free) ScimbaPytree model."""

    fn: object

    def __init__(self, fn):
        self.fn = fn

    def __call__(self, inputs):
        return self.fn(inputs)

    def ndof(self) -> int:
        return 0


def theta_fn(u):
    # AnalyticNet always receives ONE pre-concatenated array (see
    # default_pre_processing), regardless of how many named variables the
    # wrapping ApproximationSpace declares -- so this same function works
    # both for the merged "u_mu" convention
    # (LagrangianDegenerateVectorFieldSpace) and the separate "q_p_mu"
    # convention (DegenerateVariationalIntegratorFlow, where u = [q, p]
    # concatenated).
    q, p = u[0:1], u[1:2]
    return p + EPS * q


def hamiltonian_fn(u):
    q, p = u[0:1], u[1:2]
    return 0.5 * (p**2 + q**2)


# LagrangianDegenerateVectorFieldSpace needs the merged "u_mu" convention,
# one single-model space per potential.
theta_space = ApproximationSpace(
    dims={"u": 2, "mu": 0},
    list_models=[(AnalyticNet(theta_fn), "vec", 1)],
    model_type="u_mu",
)
hamiltonian_space = ApproximationSpace(
    dims={"u": 2, "mu": 0},
    list_models=[(AnalyticNet(hamiltonian_fn), "scalar", None)],
    model_type="u_mu",
)

# DegenerateVariationalIntegratorFlow needs q, p separate ("q_p_mu"), same
# convention as the canonical-Hamiltonian schemes (VerletFlowNonSep, ...).
# Same underlying analytic functions (theta_fn/hamiltonian_fn concatenate
# their args either way), just wrapped in a differently-typed space.
potentials_qp = ApproximationSpace(
    dims={"q": 1, "p": 1, "mu": 0},
    list_models=[
        (AnalyticNet(theta_fn), "vec", 1),
        (AnalyticNet(hamiltonian_fn), "scalar", None),
    ],
    model_type="q_p_mu",
)


def hamiltonian(u):
    return float(0.5 * (u[1] ** 2 + u[0] ** 2))


def exact(u0, t):
    q0, p0 = u0
    omega = 1.0 / EPS
    q = q0 * jnp.cos(omega * t) + p0 * jnp.sin(omega * t)
    p = -q0 * jnp.sin(omega * t) + p0 * jnp.cos(omega * t)
    return jnp.array([q, p])


u0 = jnp.array([1.0, 0.0])
mu = jnp.array([])

# %% Local truncation order: one DVI step (after a 1-step RK4 bootstrap) vs
# the exact analytic solution, as dt shrinks.

F_space = LagrangianDegenerateVectorFieldSpace(
    dim=2, theta=theta_space, hamiltonian=hamiltonian_space
)

dts = [0.1, 0.05, 0.025, 0.0125, 0.00625]
errors = []
for dt in dts:
    dvi_model = DegenerateVariationalIntegratorFlow(
        dim=2, potentials_hamiltonian_space=potentials_qp, dt=dt, newton_iter=25
    )
    dvi_space = MultistepPhaseSpaceApproxSpace(
        state_dim=1, params_dim=0, model=dvi_model, model_type="q_p_qm_pm_mu"
    )
    starter = Rk4Flow(dim=2, flownet=F_space, dt=dt, params_dim=0)
    x1, x0 = bootstrap_multistep_window(starter, u0, mu, 2)
    window0 = (x1[:1], x1[1:], x0[:1], x0[1:])
    qs, ps = dvi_space.rollout_trajectory(dvi_space, window0, mu, 1)
    u_dvi = jnp.concatenate([qs[1], ps[1]])
    u_exact = exact(u0, 2 * dt)
    errors.append(float(jnp.max(jnp.abs(u_dvi - u_exact))))

print(f"{'dt':>8}{'error':>14}")
for dt, e in zip(dts, errors):
    print(f"{dt:8.5f}{e:14.4e}")
orders = [
    jnp.log(errors[i] / errors[i + 1]) / jnp.log(dts[i] / dts[i + 1])
    for i in range(len(dts) - 1)
]
print(
    "Observed order (consecutive dt pairs, expect ~1):",
    [f"{float(o):.2f}" for o in orders],
)

# %% Long-time energy conservation: DVI (bounded) vs explicit Euler on the
# same continuous F (unbounded growth) -- the property that matters for a
# variational/symplectic integrator, over a horizon many periods long.

dt_long = 0.05
n_periods = 30
n_steps = int(n_periods * 2 * jnp.pi * EPS / dt_long)

dvi_model = DegenerateVariationalIntegratorFlow(
    dim=2, potentials_hamiltonian_space=potentials_qp, dt=dt_long, newton_iter=25
)
dvi_space = MultistepPhaseSpaceApproxSpace(
    state_dim=1, params_dim=0, model=dvi_model, model_type="q_p_qm_pm_mu"
)
starter = Rk4Flow(dim=2, flownet=F_space, dt=dt_long, params_dim=0)
x1, x0 = bootstrap_multistep_window(starter, u0, mu, 2)
window0 = (x1[:1], x1[1:], x0[:1], x0[1:])
qs_dvi, ps_dvi = dvi_space.rollout_trajectory(dvi_space, window0, mu, n_steps)
traj_dvi = jnp.concatenate([qs_dvi, ps_dvi], axis=-1)

euler_model = ExplicitEulerFlow(dim=2, flownet=F_space, dt=dt_long, params_dim=0)
euler_space = FlowsApproximationSpace(
    state_dim=2, params_dim=0, model=euler_model, model_type="u_mu", rollout=1
)
traj_euler = euler_space.rollout_trajectory(euler_space, u0, mu, n_steps)

H0 = hamiltonian(u0)
H_dvi = jnp.array([hamiltonian(u) for u in traj_dvi])
H_euler = jnp.array([hamiltonian(u) for u in traj_euler])

print(f"\nH0 = {H0:.4f}  ({n_periods} periods, {n_steps} steps, dt={dt_long})")
print(
    f"DVI   energy: min={float(H_dvi.min()):.4f} max={float(H_dvi.max()):.4f}  "
    f"rel. range={(float(H_dvi.max()) - float(H_dvi.min())) / H0:.3e}"
)
print(
    f"Euler energy: min={float(H_euler.min()):.4f} max={float(H_euler.max()):.4e}  "
    f"final/H0={float(H_euler[-1]) / H0:.3e}"
)

# %% Plots

t_axis_dvi = jnp.arange(n_steps + 1) * dt_long
t_axis_euler = jnp.arange(n_steps + 1) * dt_long

fig, axes = plt.subplots(1, 4, figsize=(20, 5))

axes[0].loglog(dts, errors, "o-")
axes[0].set_xlabel("dt")
axes[0].set_ylabel("error after 1 DVI step")
axes[0].set_title("Local truncation order (~1)")

ax = axes[1]
ax.plot(traj_dvi[:, 0], traj_dvi[:, 1], "C0-", linewidth=0.7, label="DVI")
ax.plot(
    traj_euler[:, 0], traj_euler[:, 1], "C1-", linewidth=0.7, label="Explicit Euler"
)
ax.set_xlabel("q")
ax.set_ylabel("p")
ax.set_title("Phase portrait")
ax.legend()

# DVI's own energy oscillation (~5e-2 rel. range) is invisible next to
# Euler's unbounded growth (~1e13x) on a shared axis -- a dedicated panel
# at DVI's own scale, plus a shared log-scale panel for the qualitative
# bounded-vs-unbounded contrast.
ax = axes[2]
ax.plot(t_axis_dvi, H_dvi - H0, "C0-")
ax.set_xlabel("t")
ax.set_ylabel("H(t) - H0")
ax.set_title("Energy conservation (DVI only)")

ax = axes[3]
ax.semilogy(t_axis_dvi, jnp.abs(H_dvi - H0), "C0-", label="DVI")
ax.semilogy(t_axis_euler, jnp.abs(H_euler - H0), "C1-", label="Explicit Euler")
ax.set_xlabel("t")
ax.set_ylabel("|H(t) - H0|")
ax.set_title("Energy conservation (log scale, both)")
ax.legend()

plt.tight_layout()
plt.show()

# %%