"""Two oscillators coupled by a quartic potential, vectorized over many orbits.

    H = (px^2 + py^2) / 2 + (w1^2 x^2 + w2^2 y^2) / 2 + beta x^2 y^2

With w1 = w2 this is the system studied by Pullen and Edmonds (1981). Scaling x, y, px, py by sqrt(E)
shows that for fixed w1, w2 the dynamics depends only on beta E, the analogue of the strike strength
of the plate. The forces are

    Fx = -w1^2 x - 2 beta x y^2,    Fy = -w2^2 y - 2 beta x^2 y.

Integration: 4th-order symplectic composition (Yoshida 1990) of the kick-drift-kick leapfrog, which
conserves energy to high accuracy over long times. Optional linear damping p -> p exp(-gamma dt) after
each step (dp/dt = -gamma p: the amplitude decays at rate gamma / 2 and the energy at rate gamma,
the analogue of the plate's sigma0). The tangent
dynamics (for Lyapunov exponents) is integrated with the same splitting, linearized.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

_C = 2 ** (1 / 3)
YOSHIDA = (1 / (2 - _C), -_C / (2 - _C), 1 / (2 - _C))


@dataclass(frozen=True)
class System:
    w1: float = 1.0
    w2: float = 1.0
    beta: float = 1.0

    def force(self, x, y):
        return -self.w1**2 * x - 2 * self.beta * x * y * y, -self.w2**2 * y - 2 * self.beta * x * x * y

    def tangent_force(self, x, y, dx, dy):
        """Linearized force applied to a tangent vector (dx, dy) at (x, y)."""
        b = self.beta
        fxx = -self.w1**2 - 2 * b * y * y
        fyy = -self.w2**2 - 2 * b * x * x
        fxy = -4 * b * x * y
        return fxx * dx + fxy * dy, fxy * dx + fyy * dy

    def energy(self, s):
        x, y, px, py = s
        return 0.5 * (px**2 + py**2) + 0.5 * (self.w1**2 * x**2 + self.w2**2 * y**2) + self.beta * x**2 * y**2

    def mode_energies(self, s):
        """Linear (uncoupled) energies of the two oscillators."""
        x, y, px, py = s
        return 0.5 * (px**2 + self.w1**2 * x**2), 0.5 * (py**2 + self.w2**2 * y**2)


def _leapfrog(sys, s, d, h):
    """One kick-drift-kick step of size h on state s (4, M) and tangent d (4, M) or None."""
    x, y, px, py = s
    fx, fy = sys.force(x, y)
    px = px + 0.5 * h * fx
    py = py + 0.5 * h * fy
    if d is not None:
        dx, dy, dpx, dpy = d
        gx, gy = sys.tangent_force(x, y, dx, dy)
        dpx = dpx + 0.5 * h * gx
        dpy = dpy + 0.5 * h * gy
        dx = dx + h * dpx
        dy = dy + h * dpy
    x = x + h * px
    y = y + h * py
    fx, fy = sys.force(x, y)
    px = px + 0.5 * h * fx
    py = py + 0.5 * h * fy
    if d is not None:
        gx, gy = sys.tangent_force(x, y, dx, dy)
        dpx = dpx + 0.5 * h * gx
        dpy = dpy + 0.5 * h * gy
        d = np.array([dx, dy, dpx, dpy])
    return np.array([x, y, px, py]), d


def step(sys, s, d, dt, gamma=0.0):
    for c in YOSHIDA:
        s, d = _leapfrog(sys, s, d, c * dt)
    if gamma:
        f = np.exp(-gamma * dt)
        s[2:] *= f
        if d is not None:
            d[2:] *= f
    return s, d


def run(sys, s0, t_end, dt=0.01, gamma=0.0, lyap=False, renorm_every=10, record_every=None,
        section=False, d0=None):
    """Integrate orbits s0 (4, M) to t_end.

    lyap: integrate a tangent vector and return the finite-time largest Lyapunov exponent of each orbit
          (Benettin: renormalize every renorm_every steps and average the log growth), plus the running
          log-growth sum at each record time if record_every is given.
    section: record crossings of y = 0 with py > 0 (Poincare section), as (x, px) per orbit.
    record_every: record the state every that many steps.
    """
    s = np.array(s0, dtype=float)
    M = s.shape[1]
    n = int(round(t_end / dt))
    d = None
    logsum = np.zeros(M)
    if lyap:
        if d0 is None:
            rng = np.random.default_rng(0)
            d = rng.standard_normal((4, M))
        else:
            d = np.array(d0, dtype=float)
        d /= np.linalg.norm(d, axis=0)
    rec_t, rec_s, rec_log = [], [], []
    sec = [[] for _ in range(M)] if section else None
    for i in range(n):
        y_old = s[1].copy()
        s_old = s.copy()
        s, d = step(sys, s, d, dt, gamma)
        if section:
            cross = np.nonzero((y_old < 0) & (s[1] >= 0) & (s[3] > 0))[0]
            for m in cross:
                a = -y_old[m] / (s[1, m] - y_old[m])
                sec[m].append(((1 - a) * s_old[0, m] + a * s[0, m], (1 - a) * s_old[2, m] + a * s[2, m]))
        if lyap and (i + 1) % renorm_every == 0:
            nrm = np.linalg.norm(d, axis=0)
            logsum += np.log(nrm)
            d /= nrm
        if record_every and (i + 1) % record_every == 0:
            rec_t.append((i + 1) * dt)
            rec_s.append(s.copy())
            if lyap:
                rec_log.append(logsum.copy())
    out = dict(s=s)
    if lyap:
        out["lyap"] = logsum / (n * dt)
    if record_every:
        out["t"] = np.array(rec_t)
        out["states"] = np.array(rec_s)          # (T, 4, M)
        if lyap:
            out["logsum"] = np.array(rec_log)    # (T, M)
    if section:
        out["section"] = [np.array(p) for p in sec]
    return out


def section_initial_conditions(sys, E, n_grid):
    """Initial conditions on the section y = 0, py > 0 at energy E: a grid over the allowed (x, px) ellipse."""
    xm = np.sqrt(2 * E) / sys.w1
    pm = np.sqrt(2 * E)
    xs = np.linspace(-xm, xm, n_grid + 2)[1:-1]
    ps = np.linspace(-pm, pm, n_grid + 2)[1:-1]
    X, P = np.meshgrid(xs, ps)
    X, P = X.ravel(), P.ravel()
    rest = E - 0.5 * P**2 - 0.5 * sys.w1**2 * X**2
    ok = rest > 1e-9 * E
    X, P, rest = X[ok], P[ok], rest[ok]
    return np.array([X, np.zeros_like(X), P, np.sqrt(2 * rest)])