"""The struck FPUT-beta chain, vectorized over an ensemble of strikes.

N unit masses between two fixed walls (u_0 = u_{N+1} = 0), joined by springs with force law
    f(Delta) = Delta + beta Delta^3        (Fermi-Pasta-Ulam-Tsingou, beta model)
and optionally held by on-site springs of stiffness mu^2 (which makes the spectrum inharmonic). The
energy is
    H = sum_j v_j^2 / 2 + sum_{j=0}^{N} (Delta_j^2 / 2 + beta Delta_j^4 / 4) + mu^2 sum_j u_j^2 / 2,
with Delta_j = u_{j+1} - u_j. The linear normal modes are sin(j k pi / (N + 1)), k = 1..N, with
    omega_k^2 = mu^2 + 4 sin^2(k pi / (2 (N + 1))).

Damping: friction c0 on every mass and dashpots c1 between neighbours,
    dv_j/dt += -c0 v_j + c1 (v_{j+1} - 2 v_j + v_{j-1}),
so that mode k decays (in amplitude) at (c0 + c1 (omega_k^2 - mu^2)) / 2. The dashpot term is the chain's version
of the plate's sigma1 gradient damping.

Integration: 4th-order symplectic composition (Yoshida) of leapfrog for the conservative part; after each
step the velocities are multiplied by the exact propagator exp(-dt D) of the damping (an N x N matrix,
precomputed). The tangent dynamics (largest Lyapunov exponent) uses the same splitting, linearized.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.linalg import expm

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


@dataclass(frozen=True)
class Chain:
    N: int = 32
    beta: float = 1.0
    mu: float = 0.0
    c0: float = 0.0
    c1: float = 0.0

    @property
    def k(self):
        return np.arange(1, self.N + 1)

    @property
    def omega(self):
        return np.sqrt(self.mu**2 + 4 * np.sin(self.k * np.pi / (2 * (self.N + 1))) ** 2)

    @property
    def modes(self):
        """Orthonormal mode matrix S[j, k] = sqrt(2 / (N + 1)) sin(j k pi / (N + 1)), j, k = 1..N."""
        j = np.arange(1, self.N + 1)
        return np.sqrt(2 / (self.N + 1)) * np.sin(np.outer(j, self.k) * np.pi / (self.N + 1))

    def damping_matrix(self):
        N = self.N
        L = -2 * np.eye(N) + np.eye(N, k=1) + np.eye(N, k=-1)
        return self.c0 * np.eye(N) - self.c1 * L

    def force(self, u):
        """u: (M, N) -> force (M, N)."""
        up = np.pad(u, ((0, 0), (1, 1)))
        d = np.diff(up, axis=1)                    # (M, N + 1) stretches
        f = d + self.beta * d**3
        return f[:, 1:] - f[:, :-1] - self.mu**2 * u

    def tangent_force(self, u, du):
        up = np.pad(u, ((0, 0), (1, 1)))
        dup = np.pad(du, ((0, 0), (1, 1)))
        d = np.diff(up, axis=1)
        dd = np.diff(dup, axis=1)
        g = (1 + 3 * self.beta * d**2) * dd
        return g[:, 1:] - g[:, :-1] - self.mu**2 * du

    def energy(self, u, v):
        up = np.pad(u, ((0, 0), (1, 1)))
        d = np.diff(up, axis=1)
        return 0.5 * (v**2).sum(1) + (0.5 * d**2 + 0.25 * self.beta * d**4).sum(1) + 0.5 * self.mu**2 * (u**2).sum(1)

    def mode_energies(self, u, v):
        S = self.modes
        q, p = u @ S, v @ S
        return 0.5 * (p**2 + (self.omega * q) ** 2)


def strike(chain: Chain, E, center, width, M=1, scale=None):
    """Initial state of a strike: at rest at u = 0, velocity profile exp(-((j - center) / width)^2 / 2)
    normalized to kinetic energy E. scale (M,) multiplies the velocity of each member (for ensembles)."""
    j = np.arange(1, chain.N + 1)
    prof = np.exp(-0.5 * ((j - center) / width) ** 2)
    prof *= np.sqrt(2 * E / (prof**2).sum())
    v = np.tile(prof, (M, 1))
    if scale is not None:
        v *= np.asarray(scale)[:, None]
    return np.zeros_like(v), v


def run(chain: Chain, u0, v0, t_end, dt, lyap=False, record_every=None, renorm_every=10):
    """Integrate an ensemble (M, N). Returns final state, the finite-time Lyapunov exponent per member
    (if lyap), and records (t, mode energies, total energy, Lyapunov log-sum) every record_every steps."""
    u, v = np.array(u0, float), np.array(v0, float)
    M = u.shape[0]
    D = chain.damping_matrix()
    damp = expm(-dt * D) if (chain.c0 or chain.c1) else None
    n = int(round(t_end / dt))
    if lyap:
        rng = np.random.default_rng(0)
        du = rng.standard_normal(u.shape)
        dv = rng.standard_normal(v.shape)
        nrm = np.sqrt((du**2).sum(1) + (dv**2).sum(1))[:, None]
        du /= nrm
        dv /= nrm
    logsum = np.zeros(M)
    rec = dict(t=[], Ek=[], E=[], logsum=[])
    for i in range(n):
        for c in YOSHIDA:
            h = c * dt
            v = v + 0.5 * h * chain.force(u)
            if lyap:
                dv = dv + 0.5 * h * chain.tangent_force(u, du)
                du = du + h * dv
            u = u + h * v
            v = v + 0.5 * h * chain.force(u)
            if lyap:
                dv = dv + 0.5 * h * chain.tangent_force(u, du)
        if damp is not None:
            v = v @ damp.T
            if lyap:
                dv = dv @ damp.T
        if lyap and (i + 1) % renorm_every == 0:
            nrm = np.sqrt((du**2).sum(1) + (dv**2).sum(1))
            logsum += np.log(nrm)
            du /= nrm[:, None]
            dv /= nrm[:, None]
        if record_every and (i + 1) % record_every == 0:
            rec["t"].append((i + 1) * dt)
            rec["Ek"].append(chain.mode_energies(u, v))
            rec["E"].append(chain.energy(u, v))
            rec["logsum"].append(logsum.copy())
    out = dict(u=u, v=v)
    if lyap:
        out["lyap"] = logsum / (n * dt)
    if record_every:
        out.update({k: np.array(x) for k, x in rec.items()})
    return out


def spectral_entropy(Ek):
    """Normalized effective number of modes exp(S) / N, S = -sum e ln e, e = E_k / sum E_k (last axis).
    1 at equipartition, 1/N when all energy is in one mode."""
    e = Ek / Ek.sum(-1, keepdims=True)
    S = -(e * np.log(np.where(e > 0, e, 1))).sum(-1)
    return np.exp(S) / Ek.shape[-1]