"""
Building a Custom System with the Low-Level Integrators
===========================================================

Every built-in continuous system in :mod:`physicskit.chaos.systems.continuous`
follows the same recipe: a module-level Numba ``@njit`` right-hand-side
function with signature ``rhs(state, t, params) -> ndarray``, wrapped by a
:class:`physicskit.chaos.core.base_system.DynamicalSystem` subclass whose
``trajectory()`` method calls :func:`physicskit.chaos.core.integrators.rk4_integrate`.
This example follows that same recipe to add a new system -- the Van der Pol
oscillator -- that is not one of the built-ins, and then shows the even more
low-level path of calling an integrator directly on a bare right-hand-side
function with no wrapper class at all.
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
from numba import njit
from numpy.typing import NDArray

from physicskit.chaos.core.base_system import DynamicalSystem
from physicskit.chaos.core.integrators import leapfrog_integrate, rk4_integrate


# %%
# Recipe 1: a full ``DynamicalSystem`` subclass
# ------------------------------------------------
# This is the same pattern used by :class:`physicskit.chaos.systems.continuous.Lorenz`
# and friends -- it plugs into every other ``physicskit.chaos`` tool (the Lyapunov
# metrics, the divergence visualizer, etc.) that expects a
# :class:`~physicskit.chaos.core.base_system.DynamicalSystem`. The system
# built here, the Van der Pol oscillator, is the classic model of a
# *relaxation oscillator* -- a self-sustaining electrical circuit oscillation
# (originally a vacuum-tube circuit) with amplitude-dependent damping:
#
# .. math::
#
#     \ddot{x} - \mu (1 - x^2) \dot{x} + x = 0,
#
# or, as the first-order system in :math:`(x, v)` actually integrated,
# :math:`\dot{x} = v`, :math:`\dot{v} = \mu(1-x^2)v - x`. The nonlinear
# damping term pumps energy in when :math:`|x| < 1` (negative damping) and
# removes it when :math:`|x| > 1` (positive damping), so trajectories from
# any nonzero initial condition spiral onto the same stable limit cycle
# regardless of amplitude.
@njit(cache=True)
def _van_der_pol_rhs(state: NDArray[np.float64], t: float, params: NDArray[np.float64]) -> NDArray[np.float64]:
    r"""Van der Pol vector field ``dx/dt = f(x, t; mu)``."""
    mu = params[0]
    x, v = state[0], state[1]
    out = np.empty(2)
    out[0] = v
    out[1] = mu * (1.0 - x * x) * v - x
    return out


class VanDerPol(DynamicalSystem):
    r"""The Van der Pol oscillator: a relaxation oscillator with a stable limit cycle."""

    dim = 2

    def __init__(self, mu: float = 3.0):
        self.mu = float(mu)

    @property
    def params(self) -> NDArray[np.float64]:
        return np.array([self.mu])

    def rhs(self, state: NDArray[np.float64], t: float) -> NDArray[np.float64]:
        return np.asarray(_van_der_pol_rhs(np.asarray(state, dtype=np.float64), t, self.params))

    def initial_state(self) -> NDArray[np.float64]:
        return np.array([0.5, 0.0])

    def trajectory(self, state0=None, t0: float = 0.0, dt: float = 0.01, n_steps: int = 5000):
        state0 = self.initial_state() if state0 is None else np.asarray(state0, dtype=np.float64)
        return rk4_integrate(_van_der_pol_rhs, state0, t0, dt, n_steps, self.params)


system = VanDerPol(mu=3.0)
t, states = system.trajectory(n_steps=5000, dt=0.01)

fig, ax = plt.subplots(figsize=(6, 6))
ax.plot(states[:, 0], states[:, 1], lw=0.6, color="darkorange")
ax.set_xlabel("x")
ax.set_ylabel("v")
ax.set_title(r"Van der Pol oscillator ($\mu=3$): trajectory spirals onto a limit cycle")

# %%
# Recipe 2: calling an integrator directly, with no wrapper class
# --------------------------------------------------------------------
# For quick experiments you don't need a full class at all: any Numba
# ``@njit`` function with signature ``f(state, t, params) -> ndarray`` can be
# handed straight to :func:`~physicskit.chaos.core.integrators.rk4_integrate`. Here we
# integrate a plain (undamped, undriven) simple harmonic oscillator,
#
# .. math::
#
#     \ddot{x} = -\omega^2 x,
#
# with the *symplectic* leapfrog integrator instead, which conserves the
# conserved energy :math:`E = \tfrac{1}{2}v^2 + \tfrac{1}{2}\omega^2 x^2` far
# better than RK4 over long integrations because it exactly preserves
# phase-space volume.


@njit(cache=True)
def _harmonic_force(pos: NDArray[np.float64], t: float, params: NDArray[np.float64]) -> NDArray[np.float64]:
    """Restoring force ``-omega^2 * x`` for a simple harmonic oscillator."""
    omega = params[0]
    return -(omega**2) * pos


omega = 2.0
pos0 = np.array([1.0])
vel0 = np.array([0.0])
t2, positions, velocities = leapfrog_integrate(_harmonic_force, pos0, vel0, 0.0, 0.05, 2000, np.array([omega]))

energy = 0.5 * velocities[:, 0] ** 2 + 0.5 * omega**2 * positions[:, 0] ** 2

fig2, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))
ax1.plot(t2, positions[:, 0])
ax1.set_xlabel("t")
ax1.set_ylabel("x")
ax1.set_title("Harmonic oscillator (leapfrog)")
ax2.plot(t2, (energy - energy[0]) / energy[0])
ax2.set_xlabel("t")
ax2.set_ylabel("relative energy drift")
ax2.set_title("Leapfrog conserves energy far better than RK4")
fig2.tight_layout()

plt.show()