"""One explicit time step (symplectic Euler) of the von Karman plate with a
thickness field, a stress-free profile, the losses, and the drumstick, as a
pure function of PyTorch tensors over the ``(ny, nx)`` grid.

The bending forces are the gradient of the discrete Kirchhoff energy with the
local bending stiffness D(h) at the nodes and at the cells (natural free-edge
conditions); the membrane forces come from bilinear cells with 2 x 2 Gauss
integration, the cell's mean thickness, and the Marguerre form of the von
Karman strains. Nodes outside the plate (weight 0) stay at zero: a second
difference exists only where the node and both neighbors belong to the plate,
and a cell contributes only when its four corners do. That gives the free edge
of the rectangle and a staircase free edge for any other outline, hole included.

Three places would otherwise turn single-precision rounding into audible
noise, and are written as in the reference implementation: the membrane
strains carry the deviation from the stress-free shape in every term (the
direct form 0.5 (wx^2 - w0x^2) leaves a rounding error proportional to w0x^2,
which does not decay); the update of w is compensated, with the rounding error
carried in a separate array; and the listener sum (in ``simulation.py``) uses
the first difference of the delayed velocity rather than the instantaneous
acceleration.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F

from ..constants import STICK_RELEASE
from ..linear.assemble import shape_derivatives
from ..params import Strike
from ..plate import Plate


def pad(t: torch.Tensor, left: int, right: int, top: int, bottom: int) -> torch.Tensor:
    return F.pad(t, (left, right, top, bottom))


@dataclass
class Static:
    """Per-node and per-cell data of a plate on the device, fixed for a run."""

    ins: torch.Tensor
    """Node in the plate, as 0/1."""
    wgt: torch.Tensor
    inv_mass: torch.Tensor
    """1 / lumped mass, zero outside."""
    Dn: torch.Tensor
    """Bending stiffness at the nodes."""
    Dc: torch.Tensor
    """Bending stiffness at the cells (zero outside), ``(ny - 1, nx - 1)``."""
    Cm: torch.Tensor
    """Membrane stiffness at the cells times the Gauss weight hx hy / 4 (zero outside)."""
    Gs: torch.Tensor
    has_a: torch.Tensor
    has_b: torch.Tensor
    nb: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
    """Left, right, lower, upper neighbor in the plate (0/1)."""
    w0x: torch.Tensor
    """Slopes of the stress-free shape at the Gauss points, ``(4, ny - 1, nx - 1)``."""
    w0y: torch.Tensor
    dNdx: torch.Tensor
    """Shape function derivatives, ``(4 gp, 4 corners)``."""
    dNdy: torch.Tensor
    relief: torch.Tensor
    """Inertia-relief acceleration field per unit stick force."""
    patch_w: torch.Tensor
    """m_k / M_patch on the contact patch, zero elsewhere: the contact coordinate is its dot with w."""
    patch_f: torch.Tensor
    """1 / M_patch on the contact patch: the force distribution per unit mass."""
    hx: float
    hy: float
    nu: float
    dt: float
    sigma0: float
    sigma1: float
    sigma_in: float
    K: float
    alpha: float
    stick_mass: float
    strike: bool


def static_data(pl: Plate, strike: Strike | None, dt: float, sigma0: float, sigma1: float, sigma_in: float, device, dtype) -> Static:
    g, m = pl.grid, pl.material
    t = lambda a: torch.as_tensor(np.asarray(a, dtype=np.float64), dtype=dtype, device=device)  # noqa: E731
    ins = g.mask.astype(float)
    mass = pl.node_mass
    inv_mass = np.where(mass > 0, 1 / np.where(mass > 0, mass, 1), 0.0)
    cm = g.cell_mask.astype(float)
    hc = pl.cell_thickness
    left, right, down, up = g.neighbor_masks()
    w0c = [np.asarray(c) for c in (pl.w0[:-1, :-1], pl.w0[:-1, 1:], pl.w0[1:, 1:], pl.w0[1:, :-1])]
    dNdx = np.zeros((4, 4))
    dNdy = np.zeros((4, 4))
    w0x = np.zeros((4,) + cm.shape)
    w0y = np.zeros((4,) + cm.shape)
    for gp in range(4):
        dNdx[gp], dNdy[gp] = shape_derivatives(gp, g.hx, g.hy)
        w0x[gp] = sum(dNdx[gp, b] * w0c[b] for b in range(4))
        w0y[gp] = sum(dNdy[gp, b] * w0c[b] for b in range(4))
    if strike is not None:
        patch = pl.strike_patch(strike)
        cx, cy = pl.center
        Ix, Iy = pl.inertia
        relief = ins * (1 / pl.mass + (g.Y - cy) * (patch.cy - cy) / Ix + (g.X - cx) * (patch.cx - cx) / Iy)
        patch_w = np.where(patch.mask, mass / patch.mass, 0.0)
        patch_f = np.where(patch.mask, 1 / patch.mass, 0.0)
        K, alpha, ms = strike.stiffness, strike.exponent, strike.mass
    else:
        relief = patch_w = patch_f = np.zeros(g.shape)
        K, alpha, ms = 0.0, 1.0, 1.0
    return Static(
        ins=t(ins),
        wgt=t(g.node_weight),
        inv_mass=t(inv_mass),
        Dn=t(m.bending_stiffness(pl.h)),
        Dc=t(m.bending_stiffness(hc) * cm),
        Cm=t(g.hx * g.hy / 4 * m.membrane_stiffness(hc) * cm),
        Gs=t(g.hx * g.hy / 4 * m.shear_stiffness(hc) * cm),
        has_a=t(left & right),
        has_b=t(down & up),
        nb=(t(left), t(right), t(down), t(up)),
        w0x=t(w0x),
        w0y=t(w0y),
        dNdx=t(dNdx),
        dNdy=t(dNdy),
        relief=t(relief),
        patch_w=t(patch_w),
        patch_f=t(patch_f),
        hx=g.hx,
        hy=g.hy,
        nu=m.nu,
        dt=dt,
        sigma0=sigma0,
        sigma1=sigma1,
        sigma_in=sigma_in,
        K=K,
        alpha=alpha,
        stick_mass=ms,
        strike=strike is not None,
    )


def accelerations(S: Static, u: torch.Tensor, v: torch.Tensor, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Undamped accelerations of every node for the displacements ``(u, v, w)`` from the stress-free shape."""
    hx, hy, nu = S.hx, S.hy, S.nu
    ihx2, ihy2, ihxy = 1 / hx**2, 1 / hy**2, 1 / (hx * hy)

    # Bending: weighted moments at the nodes, mixed differences at the cells.
    wp = pad(w, 1, 1, 1, 1)
    a = S.has_a * (wp[1:-1, 2:] - 2 * w + wp[1:-1, :-2]) * ihx2
    b = S.has_b * (wp[2:, 1:-1] - 2 * w + wp[:-2, 1:-1]) * ihy2
    wD = S.wgt * S.Dn
    Ma = S.has_a * wD * (a + nu * b)
    Mb = S.has_b * wD * (b + nu * a)
    Cc = S.Dc * (w[1:, 1:] - w[1:, :-1] - w[:-1, 1:] + w[:-1, :-1]) * ihxy
    Map = pad(Ma, 1, 1, 0, 0)
    Mbp = pad(Mb, 0, 0, 1, 1)
    Cp = pad(Cc, 1, 1, 1, 1)
    bend = (Map[:, 2:] - 2 * Ma + Map[:, :-2]) * ihx2 + (Mbp[2:, :] - 2 * Mb + Mbp[:-2, :]) * ihy2 + 2 * (1 - nu) * (Cp[1:, 1:] - Cp[1:, :-1] - Cp[:-1, 1:] + Cp[:-1, :-1]) * ihxy
    Fw = -hx * hy * bend

    # Membrane: bilinear cells, 2 x 2 Gauss points, Marguerre strains with the deviation in every term.
    U = (u[:-1, :-1], u[:-1, 1:], u[1:, 1:], u[1:, :-1])
    V = (v[:-1, :-1], v[:-1, 1:], v[1:, 1:], v[1:, :-1])
    W = (w[:-1, :-1], w[:-1, 1:], w[1:, 1:], w[1:, :-1])
    Fu_c = [None] * 4
    Fv_c = [None] * 4
    Fw_c = [None] * 4
    for gp in range(4):
        dNdx, dNdy = S.dNdx[gp], S.dNdy[gp]
        ux = dNdx[0] * U[0] + dNdx[1] * U[1] + dNdx[2] * U[2] + dNdx[3] * U[3]
        uy = dNdy[0] * U[0] + dNdy[1] * U[1] + dNdy[2] * U[2] + dNdy[3] * U[3]
        vx = dNdx[0] * V[0] + dNdx[1] * V[1] + dNdx[2] * V[2] + dNdx[3] * V[3]
        vy = dNdy[0] * V[0] + dNdy[1] * V[1] + dNdy[2] * V[2] + dNdy[3] * V[3]
        dwx = dNdx[0] * W[0] + dNdx[1] * W[1] + dNdx[2] * W[2] + dNdx[3] * W[3]
        dwy = dNdy[0] * W[0] + dNdy[1] * W[1] + dNdy[2] * W[2] + dNdy[3] * W[3]
        w0x, w0y = S.w0x[gp], S.w0y[gp]
        wx, wy = w0x + dwx, w0y + dwy
        exx = ux + dwx * (w0x + 0.5 * dwx)
        eyy = vy + dwy * (w0y + 0.5 * dwy)
        gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy
        Nxx = S.Cm * (exx + nu * eyy)
        Nyy = S.Cm * (eyy + nu * exx)
        Nxy = S.Gs * gxy
        fx = Nxx * wx + Nxy * wy
        fy = Nxy * wx + Nyy * wy
        for k in range(4):
            fu = -(Nxx * dNdx[k] + Nxy * dNdy[k])
            fv = -(Nyy * dNdy[k] + Nxy * dNdx[k])
            fw = -(fx * dNdx[k] + fy * dNdy[k])
            Fu_c[k] = fu if Fu_c[k] is None else Fu_c[k] + fu
            Fv_c[k] = fv if Fv_c[k] is None else Fv_c[k] + fv
            Fw_c[k] = fw if Fw_c[k] is None else Fw_c[k] + fw
    # Corner k of every cell goes to node offsets (0,0), (0,1), (1,1), (1,0) in (j, i).
    pads = ((0, 1, 0, 1), (1, 0, 0, 1), (1, 0, 1, 0), (0, 1, 1, 0))
    Fu = sum(pad(Fu_c[k], *pads[k]) for k in range(4))
    Fv = sum(pad(Fv_c[k], *pads[k]) for k in range(4))
    Fw = Fw + sum(pad(Fw_c[k], *pads[k]) for k in range(4))
    return Fu * S.inv_mass, Fv * S.inv_mass, Fw * S.inv_mass


def damping(S: Static, du: torch.Tensor, dv: torch.Tensor, dw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """The loss terms: -2 sigma0 v on every component, +2 sigma1 (graph Laplacian
    of w_t) on the transverse component, -2 sigma_in v on the in-plane components."""
    left, right, down, up = S.nb
    p = pad(dw, 1, 1, 1, 1)
    lap = (left * (p[1:-1, :-2] - dw) + right * (p[1:-1, 2:] - dw)) / S.hx**2 + (down * (p[:-2, 1:-1] - dw) + up * (p[2:, 1:-1] - dw)) / S.hy**2
    r = 2 * (S.sigma0 + S.sigma_in)
    return -r * du, -r * dv, -2 * S.sigma0 * dw + 2 * S.sigma1 * lap


def step(
    S: Static,
    u: torch.Tensor,
    v: torch.Tensor,
    w: torch.Tensor,
    wlo: torch.Tensor,
    du: torch.Tensor,
    dv: torch.Tensor,
    dw: torch.Tensor,
    stick: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """One step. ``stick`` holds (position, velocity, force, max force, time since release, withdrawn)."""
    dt = S.dt
    au, av, aw = accelerations(S, u, v, w)
    lu, lv, lw = damping(S, du, dv, dw)
    au = au + lu
    av = av + lv
    aw = aw + lw
    if S.strike:
        # Hertz contact between the free stick mass and the contact patch (the
        # plate nodes within the tip radius, plus the nearest one). The contact
        # coordinate is the mass-weighted mean displacement of the patch and the
        # force is spread over the patch in proportion to the node masses, so
        # the pair is conjugate. Once the contact has been released for
        # STICK_RELEASE the stick is withdrawn. The force of the previous step
        # is applied to the whole plate as the rigid acceleration field that
        # balances the patch load and its moments (inertia relief of the free
        # plate), so the plate gains no net momentum.
        sx, sv, f_prev, f_max, released, withdrawn = stick.unbind()
        aw = aw + f_prev * S.relief
        wc = (S.patch_w * w).sum()
        pen = wc - sx
        active = withdrawn == 0
        contact = active & (pen > 0)
        fc = torch.where(contact, S.K * torch.clamp(pen, min=0) ** S.alpha, torch.zeros_like(pen))
        released = torch.where(contact, torch.zeros_like(released), torch.where(active & (f_max > 0), released + dt, released))
        withdrawn = torch.where(active & (released > STICK_RELEASE), torch.ones_like(withdrawn), withdrawn)
        aw = aw - fc * S.patch_f
        svn = sv + dt * fc / S.stick_mass
        stick = torch.stack([sx + dt * svn, svn, fc, torch.maximum(f_max, fc), released, withdrawn])
    du = S.ins * (du + dt * au)
    dv = S.ins * (dv + dt * av)
    dw = S.ins * (dw + dt * aw)
    u = S.ins * (u + dt * du)
    v = S.ins * (v + dt * dv)
    # Compensated update of w: the increment dt v can be far below the float32
    # resolution of w, so the rounding error is carried separately.
    y = dt * dw + wlo
    t = w + y
    lo = (w - t) + y
    w = S.ins * t
    wlo = S.ins * lo
    return u, v, w, wlo, du, dv, dw, stick