"""A struck sheet of beads: the FPUT chain made two-dimensional, as a toy cymbal or gong.

N x N unit beads on a square grid move up and down, w_ij; the boundary ring is pinned (w = 0 on it, and the
bending is simply supported). The sheet may be built as a shallow dome of height Hd, rest shape
z = 16 Hd x (1 - x) y (1 - y). Energy:
    links (to the 4 neighbours, including the pinned ring):  sum [ T D^2 / 2 + beta (z' D + D^2 / 2)^2 ]
    bending:                                                  Kb / 2 sum (Lap w)^2   (5-point Laplacian)
with D the height difference across a link and z' its rest slope. The links are the 2D version of the FPUT-beta
springs (a height difference D stretches a link by about D^2 / 2); bending gives plate-like dispersion; the dome adds
the quadratic (three-wave) coupling of curved shells. Without a dome the normal modes are products of sines,
    sin(k pi i / (N + 1)) sin(l pi j / (N + 1)),   omega^2 = T lam + Kb lam^2,
    lam = 4 sin^2(k pi / (2 (N + 1))) + 4 sin^2(l pi / (2 (N + 1))),
and the mode energies below use this basis (approximate when there is a dome). Damping: friction c0 and
neighbour dashpots c1 (Laplacian of the velocity), applied explicitly each step. Time stepping: leapfrog.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class Plate:
    N: int = 24
    T: float = 1.0
    beta: float = 1.0
    Kb: float = 4.0
    Hd: float = 0.0
    c0: float = 0.0
    c1: float = 0.0

    def _z(self):
        g = np.arange(self.N + 2) / (self.N + 1)
        return 16 * self.Hd * np.outer(g * (1 - g), g * (1 - g))

    @property
    def zh(self):
        """Rest slopes of the horizontal links, (N, N + 1)."""
        return np.diff(self._z()[1:-1, :], axis=1)

    @property
    def zv(self):
        """Rest slopes of the vertical links, (N + 1, N)."""
        return np.diff(self._z()[:, 1:-1], axis=0)

    def lam(self):
        s = 4 * np.sin(np.arange(1, self.N + 1) * np.pi / (2 * (self.N + 1))) ** 2
        return s[:, None] + s[None, :]

    def omega(self):
        lam = self.lam()
        return np.sqrt(self.T * lam + self.Kb * lam**2)

    def sine(self):
        N = self.N
        j = np.arange(1, N + 1)
        return np.sqrt(2 / (N + 1)) * np.sin(np.outer(j, j) * np.pi / (N + 1))

    @staticmethod
    def lap(u):
        """5-point Laplacian of (M, N, N) with zero boundary."""
        p = np.pad(u, ((0, 0), (1, 1), (1, 1)))
        return p[:, 2:, 1:-1] + p[:, :-2, 1:-1] + p[:, 1:-1, 2:] + p[:, 1:-1, :-2] - 4 * u

    def force(self, w):
        p = np.pad(w, ((0, 0), (1, 1), (1, 1)))
        Dh = np.diff(p[:, 1:-1, :], axis=2)                 # (M, N, N + 1)
        Dv = np.diff(p[:, :, 1:-1], axis=1)                 # (M, N + 1, N)
        zh, zv, T, b = self.zh, self.zv, self.T, self.beta
        eh = zh * Dh + 0.5 * Dh * Dh
        ev = zv * Dv + 0.5 * Dv * Dv
        fh = T * Dh + 2 * b * eh * (zh + Dh)
        fv = T * Dv + 2 * b * ev * (zv + Dv)
        f = (fh[:, :, 1:] - fh[:, :, :-1]) + (fv[:, 1:, :] - fv[:, :-1, :])
        if self.Kb:
            f -= self.Kb * self.lap(self.lap(w))
        return f

    def energy(self, w, v):
        p = np.pad(w, ((0, 0), (1, 1), (1, 1)))
        Dh = np.diff(p[:, 1:-1, :], axis=2)
        Dv = np.diff(p[:, :, 1:-1], axis=1)
        eh = self.zh * Dh + 0.5 * Dh * Dh
        ev = self.zv * Dv + 0.5 * Dv * Dv
        U = (0.5 * self.T * Dh**2 + self.beta * eh**2).sum((1, 2)) + (0.5 * self.T * Dv**2 + self.beta * ev**2).sum((1, 2))
        U += 0.5 * self.Kb * (self.lap(w) ** 2).sum((1, 2))
        return 0.5 * (v * v).sum((1, 2)) + U

    def mode_energies(self, w, v):
        S = self.sine()
        q = np.einsum("ai,mij,jb->mab", S, w, S)
        p = np.einsum("ai,mij,jb->mab", S, v, S)
        return 0.5 * (p * p + (self.omega() * q) ** 2)


def strike(pl: Plate, E, ci, cj, radius, M=1, scale=None):
    i = np.arange(1, pl.N + 1)
    prof = np.exp(-0.5 * (((i[:, None] - ci) ** 2 + (i[None, :] - cj) ** 2) / radius**2))
    prof *= np.sqrt(2 * E / (prof**2).sum())
    v = np.tile(prof, (M, 1, 1))
    if scale is not None:
        v *= np.asarray(scale)[:, None, None]
    return np.zeros_like(v), v


def run(pl: Plate, w0, v0, t_end, dt, record_every=100):
    w, v = np.array(w0, float), np.array(v0, float)
    n = int(round(t_end / dt))
    rec = dict(t=[], Ek=[], E=[])
    f = pl.force(w)
    for i in range(n):
        v = v + 0.5 * dt * f
        w = w + dt * v
        f = pl.force(w)
        v = v + 0.5 * dt * f
        if pl.c0 or pl.c1:
            v = v - dt * (pl.c0 * v - pl.c1 * pl.lap(v))
        if (i + 1) % record_every == 0:
            rec["t"].append((i + 1) * dt)
            rec["Ek"].append(pl.mode_energies(w, v))
            rec["E"].append(pl.energy(w, v))
    out = {k: np.array(x) for k, x in rec.items()}
    out.update(w=w, v=v)
    return out