"""Displacement-based finite element solution of the cymbal2 shallow-shell model on a Cartesian grid.

The annulus is approximated by the square cells of a uniform grid whose centers lie inside it (a staircase
outer rim and a staircase hole). The three displacement fields u, v, w are each expanded in bicubic Hermite
(Bogner-Fox-Schmit) elements, C1 across cells, with the four nodal unknowns f, f_x, f_y, f_xy per field. The
weak form of MODEL.md is integrated with 3 x 3 Gauss points per cell: the membrane strains of Marguerre theory
and the bending strains are evaluated at the points, the stress resultants include the Kelvin-Voigt rates, and
the internal force is the assembled virtual work. Nothing is linearized and no modes are computed: the mass is
lumped (diagonal scaling of the consistent element mass), the equations are stepped explicitly with the
leapfrog scheme at a step below the stability limit of the in-plane waves, and the three sticks are point
masses integrated with the same step through the Hertz contact law. The Rayleigh integral is accumulated at
every step from the transverse acceleration at the Gauss points with cubic Lagrange interpolation in time at
the retarded times, so nothing about the acceleration history is stored. Free edges need no treatment: they
are the natural boundary conditions of the weak form.

The whole time step runs on the GPU through PyTorch (one dense matmul from element unknowns to Gauss-point
derivatives, elementwise constitutive evaluation, one matmul back to element forces, a scatter-add), captured
in a CUDA graph. The CPU path is the same code and serves for tests.

Usage:
  python fem_shell.py <challenge.json> <out-dir> [--spacing 0.008] [--seconds S] [--device cuda] [--fields]
                      [--report 60] [--linear] [--speed V] [--tag NAME] [--no-graph]
Writes <out>/<tag>.wav, <tag>.meta.json, and while running <tag>.progress.json and <tag>.partial.wav (the samples
computed so far, the rest zero), rewritten every --report seconds of wall time.
"""
from __future__ import annotations

import json
import math
import sys
import time
from pathlib import Path

import numpy as np
from numpy.polynomial import legendre as leg
from scipy.io import wavfile

# ----------------------------------------------------------------------------------------------------------------
# geometry of a challenge

class Profile:
    """Mid-surface height Z(r), its slope, and the thickness h(r) of a cymbal2 challenge."""

    def __init__(self, p: dict):
        self.R = p["disk_radius"]
        self.rh = p.get("mount_radius", p.get("hole_radius", 0.0))  # inner radius of the shell (the clamped mount)
        self.H = p.get("bow_height", 0.0)
        self.rb = p.get("bell_radius", 0.0)
        self.hb = p.get("bell_height", 0.0)
        self.h0 = p.get("thickness")
        self.hc, self.he = p.get("thickness_center"), p.get("thickness_rim")

    def Z(self, r):
        z = self.H * (1.0 - r * r / self.R**2)
        if self.rb > 0:
            z = z + np.where(r < self.rb, 0.5 * self.hb * (1.0 + np.cos(np.pi * np.minimum(r, self.rb) / self.rb)), 0.0)
        return z

    def dZ(self, r):
        d = -2.0 * self.H * r / self.R**2
        if self.rb > 0:
            d = d + np.where(r < self.rb, -0.5 * self.hb * (np.pi / self.rb) * np.sin(np.pi * np.minimum(r, self.rb) / self.rb), 0.0)
        return d

    def slope(self, x, y):
        r = np.hypot(x, y)
        d = self.dZ(r)
        with np.errstate(invalid="ignore", divide="ignore"):
            zx = np.where(r > 0, d * x / r, 0.0)
            zy = np.where(r > 0, d * y / r, 0.0)
        return zx, zy

    def h(self, r):
        if self.h0 is not None:
            return np.full_like(np.asarray(r, dtype=float), self.h0)
        return np.where(r <= self.rb, self.hc, self.hc + (self.he - self.hc) * (r - self.rb) / (self.R - self.rb))


# ----------------------------------------------------------------------------------------------------------------
# bicubic Hermite element on a square cell of side a

def hermite_1d(s, a):
    """Cubic Hermite functions on [0, a] at s = x/a: values and x-derivatives up to second order.
    Returns arrays (4, n) for H, H', H'' with the ordering [f(0), f'(0), f(a), f'(a)]."""
    s = np.asarray(s, dtype=float)
    H = np.stack([1 - 3 * s**2 + 2 * s**3, a * (s - 2 * s**2 + s**3), 3 * s**2 - 2 * s**3, a * (-(s**2) + s**3)])
    dH = np.stack([-6 * s + 6 * s**2, a * (1 - 4 * s + 3 * s**2), 6 * s - 6 * s**2, a * (-2 * s + 3 * s**2)]) / a
    ddH = np.stack([-6 + 12 * s, a * (-4 + 6 * s), 6 - 12 * s, a * (-2 + 6 * s)]) / a**2
    return H, dH, ddH


def element_basis(sx, sy, a):
    """Shape functions of the 16 element unknowns at points (sx, sy) in [0,1]^2: returns a dict of arrays (n, 16)
    for N, Nx, Ny, Nxx, Nyy, Nxy. Local unknown k = 4 * node + d with nodes (0,0), (1,0), (0,1), (1,1) and
    d in [f, f_x, f_y, f_xy]."""
    Hx, dHx, ddHx = hermite_1d(sx, a)
    Hy, dHy, ddHy = hermite_1d(sy, a)
    out = {k: np.zeros((len(np.atleast_1d(sx)), 16)) for k in ("N", "Nx", "Ny", "Nxx", "Nyy", "Nxy")}
    for node, (i, j) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
        for d, (p, q) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):  # (x-derivative flag, y-derivative flag)
            ix, iy = 2 * i + p, 2 * j + q
            k = 4 * node + d
            out["N"][:, k] = Hx[ix] * Hy[iy]
            out["Nx"][:, k] = dHx[ix] * Hy[iy]
            out["Ny"][:, k] = Hx[ix] * dHy[iy]
            out["Nxx"][:, k] = ddHx[ix] * Hy[iy]
            out["Nyy"][:, k] = Hx[ix] * ddHy[iy]
            out["Nxy"][:, k] = dHx[ix] * dHy[iy]
    return out


# ----------------------------------------------------------------------------------------------------------------

class ShellFEM:
    NQ = 3  # Gauss points per direction
    QUANT = ["w_x", "w_y", "w_xx", "w_yy", "w_xy", "u_x", "u_y", "v_x", "v_y"]

    def __init__(self, ch: dict, spacing: float, log=print, linear: bool = False, speed: float | None = None, c_air: float | None = None, no_drift: bool = False, inplane_scale: float = 1.0):
        self.ch = ch
        self.no_drift = no_drift
        self.inplane_scale = inplane_scale
        p = dict(ch["parameters"])
        if c_air is not None:
            p["c_air"] = c_air
        self.p = p
        self.a = a = spacing
        self.log = log
        self.linear = linear
        self.v0 = p["stick_speed"] if speed is None else speed
        pr = self.profile = Profile(p)
        R, rh = pr.R, pr.rh
        # ---- cells: centers inside the annulus; nodes: their corners; the origin is a node
        n = int(math.ceil(R / a))
        ii, jj = np.meshgrid(np.arange(-n, n), np.arange(-n, n), indexing="ij")
        cx, cy = (ii.ravel() + 0.5) * a, (jj.ravel() + 0.5) * a
        rc = np.hypot(cx, cy)
        keep = (rc < R) & (rc > rh)
        self.cells = np.stack([ii.ravel()[keep], jj.ravel()[keep]], axis=1)  # lower-left node index (i, j)
        self.n_cells = len(self.cells)
        inner_cells = set(map(tuple, np.stack([ii.ravel()[rc <= rh], jj.ravel()[rc <= rh]], axis=1).tolist()))  # cells inside the mount
        # node map
        node_ids = {}
        cell_nodes = np.zeros((self.n_cells, 4), dtype=np.int64)
        for c, (i, j) in enumerate(self.cells):
            for k, (di, dj) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
                key = (int(i + di), int(j + dj))
                if key not in node_ids:
                    node_ids[key] = len(node_ids)
                cell_nodes[c, k] = node_ids[key]
        self.n_nodes = len(node_ids)
        self.node_xy = np.zeros((self.n_nodes, 2))
        for (i, j), k in node_ids.items():
            self.node_xy[k] = (i * a, j * a)
        self.node_ids = node_ids
        self.n_dof = 12 * self.n_nodes  # blocks w, u, v of 4 n_nodes each
        # the clamped inner edge: every node that touches a cell inside the mount has all its unknowns fixed at zero
        # (displacements and the slopes of w), which is the clamp of MODEL.md on the staircase inner boundary
        fixed_nodes = np.array([k for (i, j), k in node_ids.items() if any((i + di, j + dj) in inner_cells for di in (-1, 0) for dj in (-1, 0))], dtype=np.int64)
        self.fixed_dofs = np.concatenate([b * 4 * self.n_nodes + 4 * fixed_nodes[:, None] + np.arange(4)[None, :] for b in range(3)]).ravel() if len(fixed_nodes) else np.zeros(0, dtype=np.int64)
        self.free_dofs = np.setdiff1d(np.arange(self.n_dof), self.fixed_dofs)
        self.clamped = len(fixed_nodes) > 0
        if self.clamped and no_drift:
            no_drift = False  # the clamp removes the rigid motions; no reaction or projection is needed
            self.no_drift = False
        loc = (4 * cell_nodes[:, :, None] + np.arange(4)[None, None, :]).reshape(self.n_cells, 16)
        self.cell_dofs = np.concatenate([loc, loc + 4 * self.n_nodes, loc + 8 * self.n_nodes], axis=1)  # (C, 48): w | u | v
        self.cell_dofs_w = loc
        self.area = self.n_cells * a * a
        # ---- quadrature and the derivative operator B: (C-independent) (NQ*NQ*9, 48)
        g, wg = leg.leggauss(self.NQ)
        s = (g + 1) / 2
        ws = wg / 2
        SX, SY = np.meshgrid(s, s, indexing="ij")
        self.sx, self.sy = SX.ravel(), SY.ravel()
        self.wq = (ws[:, None] * ws[None, :]).ravel() * a * a  # (G,)
        G = len(self.wq)
        self.G = G
        eb = element_basis(self.sx, self.sy, a)
        B = np.zeros((G, 9, 48))
        B[:, 0, :16], B[:, 1, :16], B[:, 2, :16], B[:, 3, :16], B[:, 4, :16] = eb["Nx"], eb["Ny"], eb["Nxx"], eb["Nyy"], eb["Nxy"]
        B[:, 5, 16:32], B[:, 6, 16:32] = eb["Nx"], eb["Ny"]
        B[:, 7, 32:48], B[:, 8, 32:48] = eb["Nx"], eb["Ny"]
        self.B = B.reshape(G * 9, 48)
        self.Nw = eb["N"]  # (G, 16) transverse displacement at the Gauss points
        # Gauss point coordinates, slopes, thickness, stiffnesses
        gx = self.cells[:, 0][:, None] * a + self.sx[None, :] * a  # (C, G)
        gy = self.cells[:, 1][:, None] * a + self.sy[None, :] * a
        self.gx, self.gy = gx, gy
        gr = np.hypot(gx, gy)
        self.Zx, self.Zy = pr.slope(gx, gy)
        self.hq = pr.h(gr)
        E, nu = p["E"], p["nu"]
        self.Cm = E * self.hq / (1 - nu**2)  # membrane stiffness per Gauss point (C, G)
        self.Db = E * self.hq**3 / (12 * (1 - nu**2))
        self.mass_total = float(np.sum(p["rho"] * self.hq * self.wq[None, :]))
        # ---- lumped mass (HRZ diagonal scaling per element, per field)
        Mdiag = np.zeros(self.n_dof)
        N2 = eb["N"] ** 2  # (G, 16)
        for blk in range(3):
            me = (p["rho"] * self.hq * self.wq[None, :]) @ N2  # (C, 16) diagonal of the consistent element mass
            mcell = np.sum(p["rho"] * self.hq * self.wq[None, :], axis=1)  # (C,)
            trans = me[:, 0::4].sum(axis=1)  # translational entries
            me = me * (mcell / trans)[:, None]
            np.add.at(Mdiag, self.cell_dofs[:, 16 * blk : 16 * (blk + 1)], me)
        # optional scaling of the lumped masses of the in-plane derivative unknowns (--inplane-scale). They are tiny and
        # set the explicit time step through the in-plane block, and scaling them by 16 halves the step count; but it
        # also pulls the spurious derivative-dominated in-plane modes down from 40 to 60 kHz into the audible band, where
        # the nonlinear cascade pumps them and they ring as steady lines (seen at 10.7 kHz on the 8 mm grid). The default
        # is therefore 1 (no scaling).
        nw = 4 * self.n_nodes
        for b in (1, 2):
            for d in (1, 2, 3):
                Mdiag[nw * b + d::4][: self.n_nodes] *= inplane_scale
        self.Mdiag = Mdiag
        # ---- strikes: footprint vectors phi_j = int phi N_j dA (fine quadrature on the cells they touch)
        self.strikes = []
        Rs = p["strike_radius"]
        c_phi = 0.5 - 2.0 / np.pi**2
        gf, wf = leg.leggauss(12)
        sf = (gf + 1) / 2
        wsf = wf / 2
        FX, FY = np.meshgrid(sf, sf, indexing="ij")
        ebf = element_basis(FX.ravel(), FY.ravel(), a)["N"]  # (144, 16)
        wff = (wsf[:, None] * wsf[None, :]).ravel() * a * a
        for st in ch["strikes"]:
            xs, ys = st["x"], st["y"]
            phi = np.zeros(self.n_dof)
            near = np.flatnonzero(np.hypot(cx[keep] - xs, cy[keep] - ys) < Rs + a)
            for c in near:
                px = self.cells[c, 0] * a + FX.ravel() * a
                py = self.cells[c, 1] * a + FY.ravel() * a
                rr = np.hypot(px - xs, py - ys)
                bump = np.where(rr <= Rs, 0.5 * (1 + np.cos(np.pi * np.minimum(rr / Rs, 1.0))), 0.0) / (c_phi * np.pi * Rs**2)
                np.add.at(phi, self.cell_dofs_w[c], (wff * bump) @ ebf)
            self.strikes.append(dict(t=st["t"], phi=phi, integral=float(phi[0 : 4 * self.n_nodes : 4].sum())))
        # ---- rigid reaction: for a unit downward force at each strike point, the rigid acceleration field of the shell
        # (translation along z plus the two tilts) that balances its net force and moments, as a nodal force vector M g.
        # With --no-drift this reaction is applied together with the stick force, so the shell gains no momentum or
        # angular momentum and does not drift (the constraint of zero rigid motion; the elastic motion is untouched).
        nw = 4 * self.n_nodes
        mw = self.Mdiag[0:nw:4]  # translational masses of the w nodes
        xn, yn = self.node_xy[:, 0], self.node_xy[:, 1]
        Mt = mw.sum()
        xc, yc = (mw * xn).sum() / Mt, (mw * yn).sum() / Mt
        Ixx, Iyy, Ixy = (mw * (yn - yc) ** 2).sum(), (mw * (xn - xc) ** 2).sum(), (mw * (xn - xc) * (yn - yc)).sum()
        for st in self.strikes:
            xs, ys = [s_ for s_ in ch["strikes"] if abs(s_["t"] - st["t"]) < 1e-12][0]["x"], [s_ for s_ in ch["strikes"] if abs(s_["t"] - st["t"]) < 1e-12][0]["y"]
            # rigid field a = c0 + cx (x - xc) + cy (y - yc) with net force F and moments F xs, F ys about the center (per unit F)
            Minv = np.linalg.inv(np.array([[Iyy, Ixy], [Ixy, Ixx]]))
            cx, cy = Minv @ np.array([xs - xc, ys - yc])
            g = np.zeros(self.n_dof)
            g[0:nw:4] = 1.0 / Mt + cx * (xn - xc) + cy * (yn - yc)
            g[1:nw:4] = cx  # f_x of a linear field
            g[2:nw:4] = cy
            st["reaction"] = self.Mdiag * g  # nodal force of the rigid acceleration field for unit stick force (upward)
        # the six rigid fields (M-orthonormalized): used with --no-drift to project the rigid components out of the state
        rig = np.zeros((6, self.n_dof))
        rig[0, 0:nw:4] = 1.0  # w translation
        rig[1, 0:nw:4], rig[1, 1:nw:4] = xn - xc, 1.0  # tilt about y (w = x)
        rig[2, 0:nw:4], rig[2, 2:nw:4] = yn - yc, 1.0  # tilt about x (w = y)
        rig[3, nw:2 * nw:4] = 1.0  # u translation
        rig[4, 2 * nw:3 * nw:4] = 1.0  # v translation
        rig[5, nw:2 * nw:4], rig[5, nw + 2:2 * nw:4] = -(yn - yc), -1.0  # in-plane rotation: u = -y, v = x
        rig[5, 2 * nw:3 * nw:4], rig[5, 2 * nw + 1:3 * nw:4] = xn - xc, 1.0
        for i in range(6):  # Gram-Schmidt in the M inner product
            for j in range(i):
                rig[i] -= (rig[j] @ (self.Mdiag * rig[i])) * rig[j]
            rig[i] /= np.sqrt(rig[i] @ (self.Mdiag * rig[i]))
        self.rigid = rig
        # ---- Rayleigh integral: weights and delays per Gauss point
        L = np.array([p["listener_x"], p["listener_y"], p["listener_z"]])
        gz = pr.Z(gr)
        d = np.sqrt((L[0] - gx) ** 2 + (L[1] - gy) ** 2 + (L[2] - gz) ** 2)  # (C, G)
        self.ray_c = (p["rho_air"] / (2 * np.pi)) * self.wq[None, :] / d
        self.ray_tau = d / p["c_air"]
        self.log(f"  mesh: {self.n_cells} cells, {self.n_nodes} nodes, {self.n_dof} unknowns ({len(self.fixed_dofs)} fixed at the mount), spacing {a * 1e3:g} mm, "
                 f"mass {self.mass_total:.4f} kg, footprint integrals {[round(s['integral'], 4) for s in self.strikes]}")

    # ------------------------------------------------------------------------------------------------------------
    def assemble_linear(self):
        """Sparse stiffness matrix of the linearized shell (the Hessian of the strain energy at rest), from the same
        Gauss-point operators the time stepping uses; symmetric, with the six rigid fields in its null space."""
        import scipy.sparse as sp

        p = self.p
        C, G, nu = self.n_cells, self.G, p["nu"]
        B = self.B.reshape(G, 9, 48)
        L = np.zeros((C, G, 6, 9))  # 9 derivative quantities -> [exx, eyy, gxy, kxx, kyy, 2kxy] at the reference state
        L[:, :, 0, 0] = self.Zx; L[:, :, 0, 5] = 1.0
        L[:, :, 1, 1] = self.Zy; L[:, :, 1, 8] = 1.0
        L[:, :, 2, 0] = self.Zy; L[:, :, 2, 1] = self.Zx; L[:, :, 2, 6] = 1.0; L[:, :, 2, 7] = 1.0
        L[:, :, 3, 2] = 1.0; L[:, :, 4, 3] = 1.0; L[:, :, 5, 4] = 2.0
        Q = np.array([[1, nu, 0], [nu, 1, 0], [0, 0, (1 - nu) / 2]])
        Dm = np.zeros((C, G, 6, 6))
        Dm[:, :, :3, :3] = self.Cm[:, :, None, None] * Q
        Dm[:, :, 3:, 3:] = self.Db[:, :, None, None] * Q
        LB = np.einsum("cgij,gjk->cgik", L, B)
        Ke = np.einsum("g,cgik,cgij,cgjl->ckl", self.wq, LB, Dm, LB, optimize=True)
        rows = np.repeat(self.cell_dofs, 48, axis=1).ravel()
        cols = np.tile(self.cell_dofs, (1, 48)).ravel()
        K = sp.coo_matrix((Ke.ravel(), (rows, cols)), shape=(self.n_dof, self.n_dof)).tocsr()
        return 0.5 * (K + K.T)

    def build_force(self, torch, dev, compile_: bool):
        """The internal force f(q, v) of the assembled shell (Kelvin-Voigt rates included) as a torch function on the
        device, and the tensors it shares with the time stepping."""
        T64 = lambda x: torch.as_tensor(np.asarray(x, dtype=np.float64), device=dev)
        T64i = lambda x: torch.as_tensor(np.asarray(x, dtype=np.int64), device=dev)
        p = self.p
        C, G = self.n_cells, self.G
        cell_dofs = T64i(self.cell_dofs)
        BT = T64(self.B.T.copy())  # (48, 81)
        Bw = T64(self.B * np.repeat(self.wq, 9)[:, None])  # (81, 48) with quadrature weights folded in
        Zx, Zy, Cm, Db = T64(self.Zx), T64(self.Zy), T64(self.Cm), T64(self.Db)
        nu, eta = p["nu"], p["eta"]
        sig1rh = T64(2.0 * p.get("sigma1", 0.0) * p["rho"] * self.hq)  # 2 sigma1 rho h at the Gauss points (Laplacian damping of w)
        lin = 0.0 if self.linear else 1.0
        n_dof = self.n_dof

        def constitutive(D):
            Dq, Dv = D[0], D[1]
            wx, wy = Dq[..., 0], Dq[..., 1]
            ax, ay = Zx + lin * wx, Zy + lin * wy  # Z_x + w_x
            exx = Dq[..., 5] + Zx * wx + 0.5 * lin * wx * wx
            eyy = Dq[..., 8] + Zy * wy + 0.5 * lin * wy * wy
            exy = 0.5 * (Dq[..., 6] + Dq[..., 7] + Zx * wy + Zy * wx + lin * wx * wy)
            rxx = Dv[..., 5] + ax * Dv[..., 0]
            ryy = Dv[..., 8] + ay * Dv[..., 1]
            rxy = 0.5 * (Dv[..., 6] + Dv[..., 7] + ax * Dv[..., 1] + ay * Dv[..., 0])
            exx = exx + eta * rxx
            eyy = eyy + eta * ryy
            exy = exy + eta * rxy
            Nxx = Cm * (exx + nu * eyy)
            Nyy = Cm * (eyy + nu * exx)
            Nxy = Cm * (1 - nu) * exy
            kxx = Dq[..., 2] + eta * Dv[..., 2]
            kyy = Dq[..., 3] + eta * Dv[..., 3]
            kxy = Dq[..., 4] + eta * Dv[..., 4]
            Mxx = Db * (kxx + nu * kyy)
            Myy = Db * (kyy + nu * kxx)
            Mxy = Db * (1 - nu) * kxy
            # generalized stresses conjugate to [w_x, w_y, w_xx, w_yy, w_xy, u_x, u_y, v_x, v_y]; the sigma1 term is the
            # virtual work 2 sigma1 rho h grad(w_t) . grad(delta w) of the Laplacian loss
            return torch.stack([Nxx * ax + Nxy * ay + sig1rh * Dv[..., 0], Nyy * ay + Nxy * ax + sig1rh * Dv[..., 1], Mxx, Myy, 2 * Mxy, Nxx, Nxy, Nxy, Nyy], dim=-1)  # (C, G, 9)

        constitutive_c = torch.compile(constitutive, dynamic=False) if (compile_ and str(dev).startswith("cuda")) else constitutive

        def force(q, v):
            qe = q[cell_dofs]  # (C, 48)
            ve = v[cell_dofs]
            D = (torch.stack([qe, ve]) @ BT).view(2, C, G, 9)
            S = constitutive_c(D)
            fe = S.reshape(C, G * 9) @ Bw  # (C, 48)
            f = torch.zeros(n_dof, dtype=torch.float64, device=dev)
            f.index_add_(0, cell_dofs.view(-1), fe.view(-1))
            return f

        invM_np = 1.0 / self.Mdiag
        invM_np[self.fixed_dofs] = 0.0  # clamped unknowns never accelerate
        ops = dict(cell_dofs=cell_dofs, cell_dofs_w=T64i(self.cell_dofs_w), NwT=T64(self.Nw.T.copy()), invM=T64(invM_np), T64=T64, T64i=T64i)
        return force, ops

    def field_interpolator(self, xs, ys):
        """Indices (P, 16) and weights (P, 16) giving w at grid points from the w unknowns; mask of points inside."""
        XX, YY = np.meshgrid(xs, ys, indexing="xy")
        px, py = XX.ravel(), YY.ravel()
        a = self.a
        ci = np.floor(px / a + 1e-9).astype(int)
        cj = np.floor(py / a + 1e-9).astype(int)
        lut = {(int(i), int(j)): c for c, (i, j) in enumerate(self.cells)}
        idx = np.zeros((len(px), 16), dtype=np.int64)
        wts = np.zeros((len(px), 16))
        inside = np.zeros(len(px), dtype=bool)
        for k in range(len(px)):
            c = lut.get((ci[k], cj[k]))
            if c is None:
                # a point on the outer edge of the last cell row belongs to the cell below/left
                for di, dj in [(-1, 0), (0, -1), (-1, -1)]:
                    c = lut.get((ci[k] + di, cj[k] + dj))
                    if c is not None:
                        ci[k] += di
                        cj[k] += dj
                        break
            if c is None:
                continue
            inside[k] = True
            sx, sy = px[k] / a - ci[k], py[k] / a - cj[k]
            idx[k] = self.cell_dofs_w[c]
            wts[k] = element_basis([sx], [sy], a)["N"][0]
        return idx, wts, inside

    # ------------------------------------------------------------------------------------------------------------
    def run(self, fs: int, nsamp: int, device: str = "cpu", m: int | None = None, field_grid=None, field_rate=None,
            report: float = 60.0, out_dir: Path | None = None, tag: str = "run", use_graph: bool = True, scale: float = 1.0, check: bool = False, compile_: bool = True):
        import torch

        t_setup = time.perf_counter()
        dev = torch.device(device)
        force, ops = self.build_force(torch, dev, compile_)
        T64, T64i = ops["T64"], ops["T64i"]
        p = self.p
        C, G = self.n_cells, self.G
        cell_dofs, cell_dofs_w, NwT, invM = ops["cell_dofs"], ops["cell_dofs_w"], ops["NwT"], ops["invM"]
        sigma0 = p["sigma0"]
        ms, kc = p["stick_mass"], p["contact_stiffness"]
        v0 = self.v0

        zero = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
        torch.manual_seed(0)
        x = torch.randn(self.n_dof, dtype=torch.float64, device=dev)
        lam = 0.0
        saved = self.linear
        self.linear = True
        force_lin, _ = self.build_force(torch, dev, False)
        self.linear = saved
        for _ in range(60):
            y = invM * force_lin(x, zero)
            lam = float(torch.dot(x, y) / torch.dot(x, x))
            x = y / torch.linalg.norm(y)
        del force_lin
        omega_max = math.sqrt(max(lam, 1.0)) * 1.05  # power iteration converges from below; margin
        dt_stab = 2.0 / omega_max
        if m is None:
            m = int(math.ceil(1.0 / (fs * 0.85 * dt_stab)))
        dt = 1.0 / (fs * m)
        self.log(f"  omega_max ~ {omega_max:.3e} rad/s ({omega_max / 2 / math.pi:.0f} Hz), stable dt < {dt_stab:.3e} s; using m = {m}, dt = {dt:.3e} s")
        nsteps = nsamp * m

        # ---- Rayleigh integral bookkeeping
        dsteps = T64(self.ray_tau / dt).view(-1)  # (C*G,) delay in steps
        rc = T64(self.ray_c).view(-1)
        wmean = T64(np.tile(self.wq, C) / self.area)  # (C*G,) quadrature weights over the area: the area mean of a Gauss-point field
        out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
        # ---- strikes
        nst = len(self.strikes)
        phi = T64(np.stack([s["phi"] for s in self.strikes]))  # (nst, n_dof)
        phi_load = T64(np.stack([s["phi"] - (s["reaction"] if self.no_drift else 0.0) for s in self.strikes]))  # force on the shell per unit F, sign as phi
        rigM = T64(self.rigid * self.Mdiag[None, :])  # (6, n_dof): <rigid_i, .>_M
        rig = T64(self.rigid)

        def project_rigid():  # remove the rigid components of the state (the constraint of MODEL.md, enforced per chunk)
            for x in (q, v):
                x.sub_(rig.T @ (rigM @ x))
        t_on = T64([s["t"] for s in self.strikes])
        n_on = T64i([int(round(s["t"] / dt)) for s in self.strikes])
        z = torch.zeros(nst, dtype=torch.float64, device=dev)
        zd = torch.zeros(nst, dtype=torch.float64, device=dev)
        active = torch.zeros(nst, dtype=torch.float64, device=dev)
        Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
        imp = torch.zeros(nst, dtype=torch.float64, device=dev)  # contact impulse per stick
        dur = torch.zeros(nst, dtype=torch.float64, device=dev)  # time in contact
        fmax = torch.zeros(nst, dtype=torch.float64, device=dev)
        out_direct = torch.zeros(nsamp, dtype=torch.float64, device=dev) if check else None
        # ---- state
        q = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
        v = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
        acc = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
        nstep = torch.zeros((), dtype=torch.int64, device=dev)
        lag_c = 1.0 / 6.0

        assert m >= 4, "the retarded-time interpolation assumes at least 4 steps per sample"

        def rayleigh(n, ag):
            nf = n.to(torch.float64)
            kk = torch.ceil((nf - 2 + dsteps) / m)
            pos = kk * m - dsteps
            fl = torch.floor(pos)
            j = nf - fl + 1
            fr = pos - fl
            wgt = torch.where(j == 0, -fr * (fr - 1) * (fr - 2) * lag_c,
                  torch.where(j == 1, (fr + 1) * (fr - 1) * (fr - 2) * 0.5,
                  torch.where(j == 2, -(fr + 1) * fr * (fr - 2) * 0.5,
                  torch.where(j == 3, (fr + 1) * fr * (fr - 1) * lag_c, torch.zeros_like(fr)))))
            valid = (j >= 0) & (j <= 3) & (kk >= 0) & (kk < nsamp)
            idx = torch.where(valid, kk, torch.zeros_like(kk)).to(torch.int64)
            return idx, torch.where(valid, rc * wgt * ag, torch.zeros_like(ag))

        ray_c = torch.compile(rayleigh, dynamic=False) if (compile_ and device.startswith("cuda")) else rayleigh

        def step():
            n = nstep
            # sticks: start the ones whose time has come (at the current surface, with relative speed v0)
            wbar = phi @ q
            wbar_dot = phi @ v
            starting = (n == n_on).to(torch.float64)
            z.copy_(z * (1 - starting) + wbar * starting)
            zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
            active.copy_(torch.maximum(active, starting))
            pen = torch.clamp(wbar - z, min=0.0)
            F = active * kc * pen * torch.sqrt(pen)
            Fst.copy_(F)
            imp.add_(dt * F)
            dur.add_(dt * (F > 0).to(torch.float64))
            fmax.copy_(torch.maximum(fmax, F))
            f = force(q, v)
            f = f + F @ phi_load  # the shell is pushed down (and, with --no-drift, held rigidly): M a = -f
            acc.copy_(-f * invM - 2 * sigma0 * v)
            zd.add_(dt * F / ms)
            z.add_(dt * zd)
            v.add_(dt * acc)
            q.add_(dt * v)
            # Rayleigh integral: scatter the Gauss-point accelerations onto the retarded output samples
            ag = (acc[cell_dofs_w] @ NwT).reshape(-1)  # (C*G,)
            ag = ag - torch.dot(wmean, ag)  # the output uses the acceleration with its area mean removed (MODEL.md)
            idx, val = ray_c(n, ag)
            out.index_add_(0, idx, val)
            if out_direct is not None:  # the same integral without retardation, for checking the scatter above
                ks = torch.div(n, m, rounding_mode="floor")
                on_sample = ((n % m) == 0) & (ks < nsamp)
                out_direct.index_add_(0, torch.where(on_sample, ks, torch.zeros_like(ks)).view(1), (torch.sum(rc * ag) * on_sample.to(torch.float64)).view(1))
            nstep.add_(1)

        # fields
        fld = None
        if field_grid is not None:
            fx, fy, ft = field_grid
            idx, wts, inside = self.field_interpolator(fx, fy)
            fld = dict(idx=T64i(idx), wts=T64(wts), inside=inside, every=int(round(fs * m / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])

        def take_frame():
            fr = (q[fld["idx"]] * fld["wts"]).sum(1).cpu().numpy().astype(np.float32)
            fr[~fld["inside"]] = np.nan
            fld["frames"].append(fr)

        def reset_state():
            for tsr in (q, v, acc, out, z, zd, active, Fst, imp, dur, fmax, nstep):
                tsr.zero_()
            if out_direct is not None:
                out_direct.zero_()

        graph = None
        if use_graph and device.startswith("cuda"):
            s = torch.cuda.Stream()
            s.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(s):
                for _ in range(2):
                    step()
            torch.cuda.current_stream().wait_stream(s)
            # undo the warm-up steps
            reset_state()
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                step()
            reset_state()
        self.log(f"  setup {time.perf_counter() - t_setup:.1f}s; stepping {nsteps} steps ({nsamp / fs:.2f} s) on {device}{' with a CUDA graph' if graph else ''}")

        t_start = time.perf_counter()
        last_report = t_start
        chunk = fld["every"] if fld is not None else 2000
        if fld is not None:
            take_frame()
        stick_stats = lambda: dict(impulse_Ns=[round(float(x), 5) for x in imp.cpu().numpy()], contact_ms=[round(1e3 * float(x), 4) for x in dur.cpu().numpy()],
                                   peak_force_N=[round(float(x), 2) for x in fmax.cpu().numpy()], rebound_mps=[round(float(x), 4) for x in zd.cpu().numpy()])
        contact_log = []
        Fprev = np.zeros(nst)
        n_done = 0
        while n_done < nsteps:
            k = min(chunk, nsteps - n_done)
            for _ in range(k):
                if graph is not None:
                    graph.replay()
                else:
                    step()
            n_done += k
            if self.no_drift:
                project_rigid()
            if fld is not None and len(fld["frames"]) < fld["nt"]:
                take_frame()
            Fc = Fst.cpu().numpy()
            for i in range(nst):
                if (Fc[i] > 0) != (Fprev[i] > 0):
                    contact_log.append((round(n_done * dt, 6), i + 1, "on" if Fc[i] > 0 else "off"))
            Fprev = Fc
            now = time.perf_counter()
            if now - last_report >= report or n_done == nsteps:
                last_report = now
                el = now - t_start
                qn = q.cpu().numpy()
                if not np.all(np.isfinite(qn)):
                    raise RuntimeError(f"solution diverged at step {n_done} (t = {n_done * dt:.4f} s)")
                w = qn[0 : 4 * self.n_nodes : 4]
                sig = out.cpu().numpy()
                prog = dict(tag=tag, step=n_done, steps=nsteps, t=round(n_done * dt, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
                            eta_s=round(el * (nsteps - n_done) / max(n_done, 1), 1), steps_per_s=round(n_done / el, 1),
                            max_abs_w_m=float(np.abs(w).max()), rms_w_m=float(np.sqrt(np.mean(w**2))),
                            stick_z_m=[round(float(x), 6) for x in z.cpu().numpy()], stick_v_mps=[round(float(x), 4) for x in zd.cpu().numpy()],
                            contact_force_N=[round(float(x), 2) for x in Fc], contacts=contact_log[-40:], **stick_stats(),
                            peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device)
                self.log(f"  t={prog['t']:.3f}s step {n_done}/{nsteps} elapsed {el:.0f}s eta {prog['eta_s']:.0f}s ({prog['steps_per_s']:.0f} steps/s) "
                         f"max|w|={prog['max_abs_w_m']:.3e} peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
                if out_dir is not None:
                    with open(out_dir / f"{tag}.progress.json", "w") as fh:
                        json.dump(prog, fh, indent=1)
                    wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
        elapsed = time.perf_counter() - t_start
        sig = out.cpu().numpy()
        info = dict(m=m, dt=dt, omega_max=omega_max, steps=nsteps, elapsed=elapsed, contacts=contact_log, sticks=stick_stats(),
                    out_direct=out_direct.cpu().numpy() if out_direct is not None else None)
        fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
        return sig, fields, info


# ----------------------------------------------------------------------------------------------------------------

def hardware_string(device: str) -> str:
    cpu = "unknown CPU"
    try:
        for line in open("/proc/cpuinfo"):
            if line.startswith("model name"):
                cpu = line.split(":", 1)[1].strip()
                break
    except OSError:
        pass
    import os

    s = f"{cpu} ({os.cpu_count()} threads)"
    if device.startswith("cuda"):
        import torch

        s += f"; time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
    return s


def main():
    ch = json.load(open(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]

    def take(flag, conv, default=None):
        nonlocal args
        if flag in args:
            i = args.index(flag)
            val = conv(args[i + 1])
            args = args[:i] + args[i + 2 :]
            return val
        return default

    spacing = take("--spacing", float, 0.008)
    seconds = take("--seconds", float)
    device = take("--device", str, "cpu")
    report = take("--report", float, 60.0)
    speed = take("--speed", float)
    m_over = take("--m", int)
    tag = take("--tag", str)
    c_air = take("--c-air", float)
    check = "--check" in args
    no_drift = "--free" not in args  # the rigid constraint of MODEL.md is the default; --free leaves the shell free (diagnostics only)
    if "--no-drift" in args:
        args.remove("--no-drift")
    inplane_scale = take("--inplane-scale", float, 1.0)
    compile_ = "--no-compile" not in args
    linear = "--linear" in args
    fields = "--fields" in args
    use_graph = "--no-graph" not in args
    out = ch["output"]
    fs = out["sample_rate_hz"]
    nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
    tag = tag or f"fem-g{spacing * 1e3:g}"
    log = lambda s: print(s, file=sys.stderr, flush=True)
    out_dir.mkdir(parents=True, exist_ok=True)
    t0 = time.perf_counter()
    log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, {nsamp / fs:.2f} s, device {device}{' (linear)' if linear else ''}")
    model = ShellFEM(ch, spacing, log=log, linear=linear, speed=speed, c_air=c_air, no_drift=no_drift, inplane_scale=inplane_scale)
    scale = out.get("signal_scale") or 1.0
    grid = None
    if fields:
        fg = ch["fields"]
        nt, ny, nx = fg["shape"]
        R = ch["parameters"]["disk_radius"]
        grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
    sig, fld, info = model.run(fs, nsamp, device=device, m=m_over, report=report, out_dir=out_dir, tag=tag, use_graph=use_graph, scale=scale, check=check, compile_=compile_,
                               field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None)
    total = time.perf_counter() - t0
    wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
    meta = {
        "solution_id": tag,
        "output": f"{out_dir.name}/{tag}.wav",
        "method": "bicubic Hermite (Bogner-Fox-Schmit) finite elements for u, v, w on a Cartesian grid masked to the annulus, lumped mass, explicit leapfrog, Hertz stick, Rayleigh integral accumulated at the retarded times",
        "parameters": {"grid_spacing_m": spacing, "cells": model.n_cells, "n_dof": model.n_dof, "gauss_points_per_cell": model.G,
                       "oversampling": info["m"], "dt_s": info["dt"], "omega_max_rad_s": info["omega_max"], "inplane_derivative_mass_scale": inplane_scale},
        "timing_seconds": round(info["elapsed"], 3),
        "hardware": hardware_string(device),
        "peak_si": float(np.max(np.abs(sig))),
        "sticks": info["sticks"],
        **({"stick_speed_override": speed} if speed is not None else {}),
        **({"linear": True} if linear else {}),
        **({"no_drift": True} if no_drift else {}),
    }
    if fld is not None:
        import h5py

        fg = ch["fields"]
        with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
            fh.attrs["challenge_id"] = ch["challenge_id"]
            fh.attrs["solution_id"] = tag
            fh.attrs["quantity"] = fg["quantity"]
            fh.attrs["units"] = fg["units"]
            fh.create_dataset("x", data=grid[0])
            fh.create_dataset("y", data=grid[1])
            fh.create_dataset("t", data=grid[2][: fld.shape[0]])
            fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
        meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
    with open(out_dir / f"{tag}.meta.json", "w") as fh:
        json.dump(meta, fh, indent=1)
    if info["out_direct"] is not None:
        d = info["out_direct"]
        log(f"  check: max|out - direct| = {np.abs(sig - d).max():.3e} Pa against max|direct| = {np.abs(d).max():.3e} Pa (equal only when c_air is huge)")
        np.save(out_dir / f"{tag}.direct.npy", d)
    log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")


if __name__ == "__main__":
    main()