"""A struck beaded wire: the FPUT chain generalized toward a cymbal or gong.

N unit masses on a wire between two pinned ends move sideways, with displacements y_j (y_0 = y_{N+1} = 0).
The wire may be built as a shallow arch (a "dome") of height Hd: rest shape z_j = 4 Hd x (1 - x), x = j / (N + 1).
Energy:
    links:   sum_j [ T D_j^2 / 2 + beta (z'_j D_j + D_j^2 / 2)^2 ],   D_j = y_{j+1} - y_j, z'_j = z_{j+1} - z_j
    bending: sum_j Kb (y_{j+1} - 2 y_j + y_{j-1})^2 / 2              (pinned ends: y_{-1} = -y_1, y_{N+2} = -y_N)
    on-site: mu^2 sum_j y_j^2 / 2
With Hd = 0 and Kb = 0 this is the FPUT-beta chain (sideways motion of a pre-tensioned beaded string; the link
stretch is about D^2 / 2, so the stretching energy is quartic in D). The bending term makes the spectrum
dispersive, omega ~ k^2 at large Kb, like a plate. The dome makes the link stretch linear in D as well
(z' D), which adds stiffness (as curvature does in a shell) and a quadratic nonlinearity (z' D^3), the
three-wave coupling of gongs and cymbals.

The linear part is K = T L + 2 beta diag-ish(z'^2) + Kb B + mu^2 I; with a dome its normal modes are not sines, so
they are computed numerically (eigh). Damping: friction c0 and neighbour dashpots c1 (Laplacian of the
velocity), applied per step with Crank-Nicolson (tridiagonal). Time stepping: Stormer-Verlet (leapfrog).
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.linalg import eigh, solve_banded


@dataclass(frozen=True)
class Wire:
    N: int = 64
    T: float = 1.0
    beta: float = 1.0
    Kb: float = 0.0
    Hd: float = 0.0
    mu: float = 0.0
    c0: float = 0.0
    c1: float = 0.0

    @property
    def zp(self):
        x = np.arange(self.N + 2) / (self.N + 1)
        z = 4 * self.Hd * x * (1 - x)
        return np.diff(z)                                   # (N + 1,) slopes of the links

    def force(self, y):
        """y: (M, N) -> force (M, N)."""
        M, N = y.shape
        yp = np.pad(y, ((0, 0), (1, 1)))
        D = np.diff(yp, axis=1)                             # (M, N + 1)
        zp = self.zp
        e = zp * D + 0.5 * D * D                            # link stretch
        fl = self.T * D + 2 * self.beta * e * (zp + D)      # dV/dD
        f = fl[:, 1:] - fl[:, :-1]
        if self.Kb:
            # bending: curvature c_j = y_{j+1} - 2 y_j + y_{j-1}, j = 0..N+1 with pinned ghosts
            yg = np.concatenate([-y[:, :1], np.zeros((M, 1)), y, np.zeros((M, 1)), -y[:, -1:]], axis=1)  # j = -1..N+2
            c = yg[:, 2:] - 2 * yg[:, 1:-1] + yg[:, :-2]    # j = 0..N+1
            # dU/dy_j = Kb (c_{j-1} - 2 c_j + c_{j+1}), j = 1..N
            f -= self.Kb * (c[:, :-2] - 2 * c[:, 1:-1] + c[:, 2:])
        f -= self.mu**2 * y
        return f

    def stiffness(self):
        """Linear stiffness matrix K (N x N) at y = 0."""
        N = self.N
        kl = self.T + 2 * self.beta * self.zp**2            # link stiffness at rest
        K = np.zeros((N, N))
        for j in range(N + 1):                              # link j joins y_j and y_{j+1} (1-based; 0 and N+1 fixed)
            a, b = j - 1, j
            if a >= 0:
                K[a, a] += kl[j]
            if b < N:
                K[b, b] += kl[j]
            if a >= 0 and b < N:
                K[a, b] -= kl[j]
                K[b, a] -= kl[j]
        if self.Kb:
            L = -2 * np.eye(N) + np.eye(N, k=1) + np.eye(N, k=-1)   # Dirichlet Laplacian; pinned bending = Kb L^2
            K += self.Kb * L @ L
        K += self.mu**2 * np.eye(N)
        return K

    def modes(self):
        w2, V = eigh(self.stiffness())
        return np.sqrt(np.maximum(w2, 0)), V

    def energy(self, y, v):
        M, N = y.shape
        yp = np.pad(y, ((0, 0), (1, 1)))
        D = np.diff(yp, axis=1)
        e = self.zp * D + 0.5 * D * D
        U = (0.5 * self.T * D * D + self.beta * e * e).sum(1)
        if self.Kb:
            yg = np.concatenate([-y[:, :1], np.zeros((M, 1)), y, np.zeros((M, 1)), -y[:, -1:]], axis=1)
            c = yg[:, 2:] - 2 * yg[:, 1:-1] + yg[:, :-2]
            U += 0.5 * self.Kb * (c[:, 1:-1] ** 2).sum(1)
        U += 0.5 * self.mu**2 * (y * y).sum(1)
        return 0.5 * (v * v).sum(1) + U


def strike(w: Wire, E, center, width, M=1, scale=None):
    j = np.arange(1, w.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(w: Wire, y0, v0, t_end, dt, record_every=100, lyap=False):
    """Stormer-Verlet (kick-drift-kick), damping by Crank-Nicolson on v after each step."""
    y, v = np.array(y0, float), np.array(v0, float)
    M, N = y.shape
    om, V = w.modes()
    damp = bool(w.c0 or w.c1)
    if damp:
        # (I + dt/2 Dm) v+ = (I - dt/2 Dm) v,  Dm = c0 I - c1 L
        a = 0.5 * dt
        ab = np.zeros((3, N))
        ab[0, 1:] = -a * w.c1
        ab[1, :] = 1 + a * (w.c0 + 2 * w.c1)
        ab[2, :-1] = -a * w.c1

        def apply_damp(x):
            rhs = (1 - a * (w.c0 + 2 * w.c1)) * x
            rhs[:, 1:] += a * w.c1 * x[:, :-1]
            rhs[:, :-1] += a * w.c1 * x[:, 1:]
            return solve_banded((1, 1), ab, rhs.T).T
    if lyap:
        rng = np.random.default_rng(0)
        dy, dv = rng.standard_normal(y.shape), rng.standard_normal(v.shape)
        nrm = np.sqrt((dy**2).sum(1) + (dv**2).sum(1))[:, None]
        dy /= nrm
        dv /= nrm
        eps = 1e-7
    logsum = np.zeros(M)
    n = int(round(t_end / dt))
    rec = dict(t=[], Ek=[], E=[], logsum=[])
    f = w.force(y)
    for i in range(n):
        v = v + 0.5 * dt * f
        y = y + dt * v
        f = w.force(y)
        v = v + 0.5 * dt * f
        if lyap:
            # tangent by finite differences of the force (keeps the model in one place)
            g = (w.force(y + eps * dy) - f) / eps
            dv = dv + dt * g
            dy = dy + dt * dv
        if damp:
            v = apply_damp(v)
            if lyap:
                dv = apply_damp(dv)
        if lyap and (i + 1) % 10 == 0:
            nrm = np.sqrt((dy**2).sum(1) + (dv**2).sum(1))
            logsum += np.log(nrm)
            dy /= nrm[:, None]
            dv /= nrm[:, None]
        if (i + 1) % record_every == 0:
            q, p = y @ V, v @ V
            rec["t"].append((i + 1) * dt)
            rec["Ek"].append(0.5 * (p * p + (om * q) ** 2))
            rec["E"].append(w.energy(y, v))
            rec["logsum"].append(logsum.copy())
    out = {k: np.array(x) for k, x in rec.items()}
    out.update(y=y, v=v, omega=om)
    return out