"""Von Karman (Airy stress function) solution of cymbal-flat-disk-v1 by a modal method on a Cartesian grid.

The disk is approximated by the square cells of a uniform Cartesian grid whose centers lie inside
it (see bfs_disk.py). The transverse displacement is expanded in the linear free-plate modes Phi_k
of that staircase domain, computed with Bogner-Fox-Schmit elements, and the Airy stress function
in the clamped biharmonic eigenfunctions Psi_p on the same mesh. With mass-normalized Phi and
L2-normalized Psi the equations are

    q_k'' + 2 alpha_k q_k' + omega_k^2 q_k = f_k psi(t/T_s) + sum_{j,p} H^p_{kj} q_j eta_p,
    eta_p = -(E h / (2 zeta_p^4)) sum_{i,j} H^p_{ij} q_i q_j,      H^p_{ij} = int Psi_p L(Phi_i, Phi_j) dA,

where L(f, g) = f_xx g_yy + f_yy g_xx - 2 f_xy g_xy (Touze et al.). Time stepping uses a two-step
Gautschi-type scheme at a step of 1/(fs m): the damped linear part is propagated exactly and the
strike and the cubic coupling force are applied with a static-exact weight, which is second-order
accurate and free of numerical dispersion in the linear limit. The Rayleigh integral is evaluated
with 3 x 3 Gauss points per cell; each point's retarded acceleration is a cubic Lagrange
interpolation of the modal accelerations at neighbouring time steps, folded into a per-step
weight table. Damping acts on the bending motion only (the in-plane state is quasi-static),
a further approximation of the challenge model.

The convergence chain has three axes: the grid spacing (8 mm, 4 mm), the number of modes kept
(transverse and Airy), and the time step. The mode truncation caps the frequency the nonlinear
cascade can reach.

Because the strike is centered on the x axis, the solution is even in y, and by default the modal bases
(transverse and Airy) are restricted to that symmetric subspace, which is exact and halves the number of
modes needed for a given frequency reach (the coupling cost scales with the cube of the count). --full
uses the unrestricted bases instead.

If the challenge has a shell-of-revolution geometry (cymbal-bow-v1: a paraboloidal bow z = Z(r);
cymbal-bell-v1: the bow plus a cosine-dome bell; cymbal-hammered-v1: with raised-cosine dents, which break the
symmetry in y so the full modal basis is used; cymbal-tapered-v1: with a thickness that varies along the
meridian, so the mass, the bending stiffness and the membrane compliance 1/(E h) become variable coefficients
in the assembled operators), the same solver applies the Marguerre shallow-shell equations, which add the initial shape w0 = Z to the von Karman
system: L(w + w0, Phi) on the right of the bending equation and 2 L(w0, w) in the compatibility equation.
With G^p_j = int Psi_p L(w0, Phi_j) dA this gives a linear membrane stiffening 2 sum_p |c_p| G^p (G^p)^T,
which is diagonalized once in the flat-mode basis to give the shell modes (exact in the Gautschi step,
damped at sigma0 + eta omega^2 / 2 like every mode, as the Kelvin-Voigt law of the challenge implies), and
quadratic coupling terms F_k = sum_p c_p [ (r^T H^p r + 2 G^p . r) (H^p r)_k + (r^T H^p r) G^p_k ]. The strike
acts along -z on the curved top face with the normalization of the specification, the listener integral uses
the distance to the top face, and the field export is the normal displacement n_z w. Shallow-shell theory
neglects in-plane inertia and terms of the order of the slope squared (0.06 at the rim of the bow, 0.19 on the
flank of the bell). The mounting hole of the bell challenge is not represented: the solver closes it, which
neglects the mass and stiffness of a 12 mm disk at the apex; a stated approximation.

With --device cuda the coupling tensor is assembled and the time stepping runs on a GPU through PyTorch
(the modes and quadrature stay in numpy/scipy); the CPU path is the reference and the two agree to round-off.

Usage: python vk_modal_disk.py <challenge.json> <out-dir> [--fields tag,...] [--force F] [--duration T] [--seconds S] [--full] [--device cuda] [--strikes FILE] [tag ...]
  --strikes FILE replaces the challenge's strike by a schedule, a JSON list of {"t", "force", "duration", "x", "y"}, for
  experiments outside the challenge (the output is then not a valid submission).
  where a tag is g<grid spacing in mm>-w<transverse modes>-a<airy modes>-m<oversampling>, e.g. g8-w200-a200-m1
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import numpy as np
from scipy.io import wavfile
sys.path.insert(0, str(Path(__file__).resolve().parent))
from bfs_disk import BFSDisk, load_challenge  # noqa: E402


class Profile:
    """An axisymmetric mid-surface z = Z(r): the paraboloidal bow of cymbal-bow-v1, optionally with the cosine-dome
    bell of cymbal-bell-v1. Uniform thickness along the normal. The hole of the bell challenge is not represented
    (the solver closes it; see the module docstring)."""

    def __init__(self, R: float, H: float, h: float, bell_radius: float = 0.0, bell_height: float = 0.0, hole_radius: float = 0.0, dents=(), thickness_center=None, thickness_rim=None):
        self.R, self.H, self.h = R, H, h
        self.rb, self.hb, self.rh = bell_radius, bell_height, hole_radius
        self.dents = [(float(d["x"]), float(d["y"]), float(d["radius"]), float(d["depth"])) for d in dents]
        self.hc, self.he = thickness_center, thickness_rim  # tapered thickness (cymbal-tapered-v1), None for uniform

    @property
    def tapered(self) -> bool:
        return self.hc is not None

    def h_at(self, x, y):
        """Thickness at (x, y): uniform, or h_c over the bell and linear to h_e at the rim."""
        x = np.asarray(x, dtype=float); y = np.asarray(y, dtype=float)
        if not self.tapered:
            return np.full(np.shape(x), self.h)
        r = np.sqrt(x * x + y * y)
        return np.where(r <= self.rb, self.hc, self.hc + (self.he - self.hc) * (r - self.rb) / (self.R - self.rb))

    @property
    def symmetric_y(self) -> bool:
        """True when the surface is even in y (no dents): the axisymmetric bow and bell."""
        return len(self.dents) == 0

    @staticmethod
    def from_challenge(ch: dict):
        g = ch.get("geometry")
        if not g or g.get("type") != "shell-of-revolution":
            return None
        p = ch["parameters"]
        h = p.get("h", p.get("thickness_center"))
        return Profile(p["disk_radius"], p["bow_height"], h, p.get("bell_radius", 0.0), p.get("bell_height", 0.0), p.get("hole_radius", 0.0), g.get("dents", ()), p.get("thickness_center"), p.get("thickness_rim"))

    # ---- radial profile and its derivatives
    def _bell(self, r):
        if self.hb == 0.0:
            return np.zeros_like(r), np.zeros_like(r), np.zeros_like(r)
        k = np.pi / self.rb
        inside = r < self.rb
        z = np.where(inside, 0.5 * self.hb * (1.0 + np.cos(k * r)), 0.0)
        dz = np.where(inside, -0.5 * self.hb * k * np.sin(k * r), 0.0)
        d2z = np.where(inside, -0.5 * self.hb * k * k * np.cos(k * r), 0.0)
        return z, dz, d2z

    def radial(self, r):
        """Z, Z' and Z'' as functions of r."""
        r = np.asarray(r, dtype=float)
        c = -2.0 * self.H / self.R**2
        z = self.H + 0.5 * c * r * r
        dz = c * r
        d2z = c * np.ones_like(r)
        zb, dzb, d2zb = self._bell(r)
        return z + zb, dz + dzb, d2z + d2zb

    @staticmethod
    def _dent(rho, rd, depth):
        """A raised-cosine depression: value, first and second radial derivatives at distance rho from its center."""
        k = np.pi / rd
        inside = rho < rd
        z = np.where(inside, -0.5 * depth * (1.0 + np.cos(k * rho)), 0.0)
        dz = np.where(inside, 0.5 * depth * k * np.sin(k * rho), 0.0)
        d2z = np.where(inside, 0.5 * depth * k * k * np.cos(k * rho), 0.0)
        return z, dz, d2z

    @staticmethod
    def _cartesian(x, y, dz, d2z):
        """Gradient and second derivatives of a radial function about the origin of (x, y), from Z'(r) and Z''(r)."""
        r = np.sqrt(x * x + y * y)
        rr = np.where(r > 0, r, 1.0)
        s = np.where(r > 0, dz / rr, d2z)  # Z'/r -> Z'' on the axis
        cx, cy = x / rr, y / rr
        return s * x, s * y, d2z * cx * cx + s * cy * cy, d2z * cy * cy + s * cx * cx, (d2z - s) * cx * cy

    def z(self, x, y):
        x = np.asarray(x, dtype=float); y = np.asarray(y, dtype=float)
        z = self.radial(np.sqrt(x * x + y * y))[0]
        for xc, yc, rd, depth in self.dents:
            z = z + self._dent(np.hypot(x - xc, y - yc), rd, depth)[0]
        return z

    def grad(self, x, y):
        x = np.asarray(x, dtype=float); y = np.asarray(y, dtype=float)
        _, dz, d2z = self.radial(np.sqrt(x * x + y * y))
        gx, gy, _, _, _ = self._cartesian(x, y, dz, d2z)
        for xc, yc, rd, depth in self.dents:
            _, ddz, dd2z = self._dent(np.hypot(x - xc, y - yc), rd, depth)
            ex, ey, _, _, _ = self._cartesian(x - xc, y - yc, ddz, dd2z)
            gx, gy = gx + ex, gy + ey
        return gx, gy

    def curvatures(self, x, y):
        """Second derivatives (Z_xx, Z_yy, Z_xy) of the mid-surface height."""
        x = np.asarray(x, dtype=float); y = np.asarray(y, dtype=float)
        _, dz, d2z = self.radial(np.sqrt(x * x + y * y))
        _, _, zxx, zyy, zxy = self._cartesian(x, y, dz, d2z)
        for xc, yc, rd, depth in self.dents:
            _, ddz, dd2z = self._dent(np.hypot(x - xc, y - yc), rd, depth)
            _, _, exx, eyy, exy = self._cartesian(x - xc, y - yc, ddz, dd2z)
            zxx, zyy, zxy = zxx + exx, zyy + eyy, zxy + exy
        return zxx, zyy, zxy

    def normal(self, x, y):
        zx, zy = self.grad(x, y)
        nz = 1.0 / np.sqrt(1.0 + zx * zx + zy * zy)
        return -zx * nz, -zy * nz, nz

    def top(self, x, y):
        """The point of the top face on the normal through the mid-surface point (x, y), with the local thickness."""
        nx, ny, nz = self.normal(x, y)
        hh = 0.5 * self.h_at(x, y)
        return x + hh * nx, y + hh * ny, self.z(x, y) + hh * nz

    def max_slope(self):
        g = np.linspace(-self.R, self.R, 801)
        X, Y = np.meshgrid(g, g, indexing="ij")
        m = X * X + Y * Y <= self.R**2
        gx, gy = self.grad(X[m], Y[m])
        return float(np.sqrt(gx * gx + gy * gy).max())


Paraboloid = Profile  # former name


class VonKarmanDisk:
    def __init__(self, ch: dict, spacing: float, n_w: int, n_airy: int, force_override: float | None = None, log=None, symmetric_y: bool | None = None, device: str = "cpu", duration_override: float | None = None, strikes=None):
        """strikes: optional list of {"t", "force", "duration", "x", "y"} replacing the challenge's single strike at
        t = 0 (for experiments outside the challenge; the output is then not a valid submission)."""
        self.ch = ch
        self.strikes = strikes
        self.device = device
        self.torch = None
        if device != "cpu":
            import torch

            self.torch = torch
        p = dict(ch["parameters"])
        if force_override is not None:
            p["strike_peak_force"] = force_override
        if duration_override is not None:
            p["strike_duration"] = duration_override
        ch2 = dict(ch)
        ch2["parameters"] = p
        self.p = p
        self.log = log or (lambda s: None)
        self.profile = Profile.from_challenge(ch)
        if symmetric_y is None:
            symmetric_y = p["strike_y"] == 0.0 and (self.profile is None or self.profile.symmetric_y)
        if symmetric_y and self.profile is not None and not self.profile.symmetric_y:
            raise ValueError("the dented surface is not symmetric in y; the symmetric basis cannot be used")
        self.symmetric_y = symmetric_y
        t0 = time.perf_counter()
        thick = self.profile.h_at if (self.profile is not None and self.profile.tapered) else None
        self.plate = BFSDisk(ch2, spacing, n_w, symmetric_y=symmetric_y, thickness=thick)  # linear modes, modal forces, alpha, omega
        pl = self.plate
        self.log(f"  linear modes: {pl.n_modes}{' (even in y)' if symmetric_y else ''} on {pl.n_cells} cells ({pl.ndof} dof) in {time.perf_counter() - t0:.1f}s, f_max={pl.omega[-1] / (2 * np.pi):.0f} Hz, rigid={pl.n_rigid}")
        self.n_w = pl.n_modes
        self.omega = pl.omega
        self.alpha = pl.alpha
        self.f = pl.f
        V = pl.V
        # Airy modes: eigenfunctions of the compatibility operator int (1/(E h)) [Psi_xx v_xx + Psi_yy v_yy
        # - nu (Psi_xx v_yy + Psi_yy v_xx) + 2 (1 + nu) Psi_xy v_xy] with clamped conditions, L2-normalized:
        # K_b psi = zeta^4 M_1 psi on interior dofs (for uniform h this is the biharmonic operator over E h)
        t0 = time.perf_counter()
        nu = p["nu"]
        Cnu = np.array([[1.0, -nu, 0.0], [-nu, 1.0, 0.0], [0.0, 0.0, 2.0 * (1.0 + nu)]])
        gx, gy = pl._gauss_xy
        Kb = pl.assemble_coeff(Cnu, 1.0 / (p["E"] * pl.h_at(gx, gy)))
        z4, Psi = pl.eigenmodes(Kb, pl.M1, n_airy, sigma=0.0, dofs=pl.interior_dofs)
        self.n_airy = len(z4)
        self.zeta4 = z4
        self.log(f"  airy modes: {self.n_airy} in {time.perf_counter() - t0:.1f}s")
        # coupling tensor H[p, i, j] by 3 x 3 Gauss quadrature per cell
        t0 = time.perf_counter()
        xq, yq, WQ = pl.quadrature(3)
        Bxx, Byy, Bxy = pl.basis_derivs_at(xq, yq)
        Pxx, Pyy, Pxy = Bxx @ V, Byy @ V, Bxy @ V  # (Q, n_w)
        Bn = pl.basis_at(xq, yq)
        PsiQ = Bn @ Psi  # (Q, n_airy)
        self.coef = -1.0 / (2.0 * self.zeta4)  # E h is inside the compatibility operator
        if self.torch is None:
            H = np.zeros((self.n_airy, self.n_w, self.n_w))
            for pi in range(self.n_airy):
                wp = (WQ * PsiQ[:, pi])[:, None]
                A = (Pxx * wp).T @ Pyy
                H[pi] = A + A.T - 2.0 * (Pxy * wp).T @ Pxy
            self.H = H
            self.Hflat = H.reshape(self.n_airy * self.n_w, self.n_w)
            shape = H.shape
        else:
            # the same assembly on the device; the tensor stays there for the time stepping
            t = self.torch
            dev = t.device(self.device)
            Pxx_t, Pyy_t, Pxy_t = (t.as_tensor(a, device=dev) for a in (Pxx, Pyy, Pxy))
            WPsi = t.as_tensor(WQ[:, None] * PsiQ, device=dev)  # (Q, n_airy)
            H = t.empty((self.n_airy, self.n_w, self.n_w), dtype=t.float64, device=dev)
            for pi in range(self.n_airy):
                wp = WPsi[:, pi : pi + 1]
                A = (Pxx_t * wp).T @ Pyy_t
                H[pi] = A + A.T - 2.0 * (Pxy_t * wp).T @ Pxy_t
            t.cuda.synchronize() if self.device.startswith("cuda") else None
            self.H = None
            self.Hflat_t = H.reshape(self.n_airy * self.n_w, self.n_w)
            shape = tuple(H.shape)
            del Pxx_t, Pyy_t, Pxy_t, WPsi
        self.log(f"  coupling tensor {shape} on {self.device} in {time.perf_counter() - t0:.1f}s")
        self.G = None  # shell coupling (n_airy, n_w), None for the flat disk
        self.V = V
        if self.profile is not None:
            t0 = time.perf_counter()
            self._apply_profile(xq, yq, WQ, PsiQ, Pxx, Pyy, Pxy)
            V = self.V
            self.log(f"  shell modes (max slope {self.profile.max_slope():.2f}): f_low={np.round(self.omega[self.omega > 0][:4] / (2 * np.pi), 2)} Hz, f_max={self.omega.max() / (2 * np.pi):.0f} Hz in {time.perf_counter() - t0:.1f}s")
        # Rayleigh integral: quadrature weights, delays, modal shapes at the points
        _, _, self.tau, self.cq = pl.listener_quadrature(3)
        self.PhiQ = Bn @ V  # (Q, n_w)
        # the strike schedule: (t0, duration, modal force vector) per hit; default is the challenge's single strike at t = 0
        if self.strikes:
            self.schedule = [(float(st["t"]), float(st["duration"]), self.strike_force(float(st["x"]), float(st["y"]), float(st["force"]))) for st in self.strikes]
        else:
            self.schedule = [(0.0, p["strike_duration"], self.f)]

    def strike_force(self, xs: float, ys: float, F: float) -> np.ndarray:
        """Modal force vector of a raised-cosine strike of peak total force F centered at (xs, ys) (horizontal
        coordinates of the top face), for the flat disk or the shell, normalized by the surface integral of the bump."""
        from numpy.polynomial import legendre as leg

        p, pl, pr = self.p, self.plate, self.profile
        Rs = p["strike_radius"]
        gr, gw = leg.leggauss(32)
        rq = (Rs * 1.1) * (gr + 1.0) / 2.0
        wr = gw * (Rs * 1.1) / 2.0
        nth = 96
        th = 2 * np.pi * np.arange(nth) / nth
        rr, tt = np.meshgrid(rq, th, indexing="ij")
        wa = (wr[:, None] * rr * (2 * np.pi / nth)).ravel()
        px = (xs + rr * np.cos(tt)).ravel()
        py = (ys + rr * np.sin(tt)).ravel()
        if pr is None:
            tx, ty, nz = px, py, np.ones_like(px)
        else:
            tx, ty, _ = pr.top(px, py)
            _, _, nz = pr.normal(px, py)
        rprime = np.sqrt((tx - xs) ** 2 + (ty - ys) ** 2)
        phi = np.where(rprime <= Rs, 0.5 * (1.0 + np.cos(np.pi * np.minimum(rprime / Rs, 1.0))), 0.0)
        C = np.sum(wa * phi / nz)  # surface integral of the bump over the top face
        Phi_s = pl.basis_at(px, py) @ self.V
        return -(F / C) * ((wa * phi / nz) @ Phi_s)

    def _apply_profile(self, xq, yq, WQ, PsiQ, Pxx, Pyy, Pxy):
        """Marguerre shallow-shell terms for the initial shape: the G tensor, the shell modes (rotation T of the
        flat modes), the rotated coupling tensors, the strike on the curved top face and the listener distances."""
        pl, pr, p = self.plate, self.profile, self.p
        t = self.torch
        zxx, zyy, zxy = pr.curvatures(xq, yq)
        # G[p, j] = int Psi_p L(w0, Phi_j) dA,  L(w0, Phi) = Z_xx Phi_yy + Z_yy Phi_xx - 2 Z_xy Phi_xy
        Lw0 = zxx[:, None] * Pyy + zyy[:, None] * Pxx - 2.0 * zxy[:, None] * Pxy  # (Q, n_w)
        G = (WQ[:, None] * PsiQ).T @ Lw0  # (n_airy, n_w)
        # linear membrane stiffening in the flat-mode basis: K = diag(omega^2) - 2 G^T diag(c) G  (c < 0)
        K = np.diag(self.omega**2) - 2.0 * (G.T * self.coef) @ G
        K = 0.5 * (K + K.T)
        w2, T = np.linalg.eigh(K)
        omega = np.sqrt(np.maximum(w2, 0.0))
        omega[omega < 1.0] = 0.0  # rigid modes stay rigid (L(w0, harmonic) = 0)
        # rotate: H'^p = T^T H^p T, G' = G T, f' = T^T f, V' = V T
        if t is None:
            Hr = np.einsum("pij,ik,jl->pkl", self.H, T, T, optimize=True)
            self.H = Hr
            self.Hflat = Hr.reshape(self.n_airy * self.n_w, self.n_w)
        else:
            dev = t.device(self.device)
            Tt = t.as_tensor(T, device=dev)
            Ht = self.Hflat_t.reshape(self.n_airy, self.n_w, self.n_w)
            Ht = t.einsum("pij,ik,jl->pkl", Ht, Tt, Tt)
            self.Hflat_t = Ht.reshape(self.n_airy * self.n_w, self.n_w)
        self.G = G @ T
        self.V = self.V @ T
        self.omega = omega
        self.alpha = p["sigma0"] + 0.5 * p["eta"] * omega**2
        # strike: -int_{S+} p Phi_k dS with p = F/C phi(r'/R_s) psi, r' the horizontal distance of the top-face point,
        # dS = dA / n_z; quadrature over the projected footprint in polar coordinates about the strike center
        from numpy.polynomial import legendre as leg

        Rs, xs, ys = p["strike_radius"], p["strike_x"], p["strike_y"]
        C = p["strike_normalization_m2"]
        gr, gw = leg.leggauss(32)
        rq = (Rs * 1.1) * (gr + 1.0) / 2.0
        wr = gw * (Rs * 1.1) / 2.0
        nth = 96
        th = 2 * np.pi * np.arange(nth) / nth
        rr, tt = np.meshgrid(rq, th, indexing="ij")
        wa = (wr[:, None] * rr * (2 * np.pi / nth)).ravel()
        px = (xs + rr * np.cos(tt)).ravel()
        py = (ys + rr * np.sin(tt)).ravel()
        tx, ty, _ = pr.top(px, py)
        rprime = np.sqrt((tx - xs) ** 2 + (ty - ys) ** 2)
        phi = np.where(rprime <= Rs, 0.5 * (1.0 + np.cos(np.pi * np.minimum(rprime / Rs, 1.0))), 0.0)
        _, _, nz = pr.normal(px, py)
        Phi_s = pl.basis_at(px, py) @ self.V
        self.f = -(p["strike_peak_force"] / C) * ((wa * phi / nz) @ Phi_s)
        self.total_force_check = float((p["strike_peak_force"] / C) * np.sum(wa * phi / nz))  # should be ~F_max
        # listener: n . a dS = w_tt dA in the shallow-shell model; distances to the top-face points
        xq_, yq_, wq_ = pl.quadrature(3)
        tx, ty, tz = pr.top(xq_, yq_)
        r = np.sqrt((p["listener_x"] - tx) ** 2 + (p["listener_y"] - ty) ** 2 + (p["listener_z"] - tz) ** 2)
        pl._listener_q = (xq_, yq_, r / p["c_air"], p["rho_air"] / (2 * np.pi) * wq_ / r)
        # the plate object doubles as the closed-form linear solver of the shell (for checks): give it the shell modes
        pl.omega, pl.alpha, pl.f, pl.V = self.omega, self.alpha, self.f, self.V
        pl.rigid = self.omega == 0.0
        pl.n_rigid = int(pl.rigid.sum())
        pl._setup_linear_response()

    def nonlinear_force(self, q: np.ndarray) -> np.ndarray:
        v = (self.Hflat @ q).reshape(self.n_airy, self.n_w)  # v[p, k] = sum_j H[p,k,j] q_j
        sq = v @ q  # q^T H^p q
        if self.G is None:
            return v.T @ (self.coef * sq)  # F_k = sum_p c_p (q^T H^p q) v[p,k]
        g = self.G @ q  # G^p . q
        return v.T @ (self.coef * (sq + 2.0 * g)) + self.G.T @ (self.coef * sq)

    def _prepare(self, fs: int, nsamp: int, m: int, field_grid, field_rate):
        p = self.p
        dt = 1.0 / (fs * m)
        w2, al = self.omega**2, self.alpha
        # exact two-step recurrence of the damped linear oscillator: q+ = a1 q + a2 q- + b F
        rigid = self.omega == 0.0
        disc = np.sqrt(al**2 - w2 + 0j)
        E1, E2 = np.exp((-al + disc) * dt), np.exp((-al - disc) * dt)
        a1 = (E1 + E2).real
        a2 = -(E1 * E2).real
        with np.errstate(divide="ignore", invalid="ignore"):
            b = np.where(rigid, 0.0, (1.0 - a1 - a2) / np.where(rigid, 1.0, w2))
        if rigid.any():
            e2 = np.exp(-2.0 * al[rigid] * dt)
            a1[rigid] = 1.0 + e2
            a2[rigid] = -e2
            b[rigid] = dt * (1.0 - e2) / (2.0 * al[rigid])
        # fractional delays folded into a per-step weight table G[step, mode]
        tau_steps = self.tau / dt
        d0 = np.floor(tau_steps).astype(int)
        fr = tau_steps - d0
        lag = np.stack([
            -fr * (fr - 1.0) * (fr - 2.0) / 6.0,
            (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,
            -(fr + 1.0) * fr * (fr - 2.0) / 2.0,
            (fr + 1.0) * fr * (fr - 1.0) / 6.0,
        ], axis=1)
        dmax = int(d0.max()) + 2
        G = np.zeros((dmax + 1, self.n_w))
        for j, off in enumerate([-1, 0, 1, 2]):
            dd = d0 + off
            for u in np.unique(dd):
                sel = dd == u
                G[u] += (self.cq[sel] * lag[sel, j]) @ self.PhiQ[sel]
        fld = None
        if field_grid is not None:
            fx, fy, ft = field_grid
            XX, YY = np.meshgrid(fx, fy, indexing="xy")
            inside = self.plate.inside(XX.ravel(), YY.ravel())
            PhiG = self.plate.basis_at(XX.ravel(), YY.ravel()) @ self.V
            if self.profile is not None:
                PhiG = PhiG * self.profile.normal(XX.ravel(), YY.ravel())[2][:, None]  # u . n = n_z w
            fld = dict(PhiG=PhiG, inside=inside, every=int(round(fs * m / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx))
        return dict(dt=dt, T=p["strike_duration"], Om=2 * np.pi / p["strike_duration"], nsteps=nsamp * m, a1=a1, a2=a2, b=b, w2=w2, al=al, G=G, dmax=dmax, fld=fld)

    def run(self, fs: int, nsamp: int, m: int, field_grid=None, field_rate: float | None = None):
        pr = self._prepare(fs, nsamp, m, field_grid, field_rate)
        if self.torch is not None:
            return self._run_device(pr, m, nsamp)
        dt, T, Om, nsteps, a1, a2, b, w2, al, G, dmax, fld = (pr[k] for k in ("dt", "T", "Om", "nsteps", "a1", "a2", "b", "w2", "al", "G", "dmax", "fld"))
        f = self.f
        ring = np.zeros((dmax + 1, self.n_w))
        out = np.zeros(nsamp)
        q_prev = np.zeros(self.n_w)
        q = np.zeros(self.n_w)
        frames = [] if fld else None
        t_start = time.perf_counter()
        report = max(1, nsteps // 10)
        for nstep in range(nsteps + 1):
            t = nstep * dt
            Fext = self.nonlinear_force(q)
            for t0, Ts, fvec in self.schedule:
                if t0 <= t <= t0 + Ts:
                    Fext = Fext + fvec * (0.5 * (1.0 - np.cos(2 * np.pi * (t - t0) / Ts)))
            q_next = a1 * q + a2 * q_prev + b * Fext
            qdd = Fext - w2 * q - 2.0 * al * (q_next - q_prev) / (2.0 * dt)
            ring[nstep % (dmax + 1)] = qdd
            if nstep % m == 0:
                k = nstep // m
                if k < nsamp:
                    idx = (nstep - np.arange(dmax + 1)) % (dmax + 1)
                    out[k] = np.sum(G * ring[idx])
            if frames is not None and nstep % fld["every"] == 0 and len(frames) < fld["nt"]:
                fr_ = (fld["PhiG"] @ q).astype(np.float32)
                fr_[~fld["inside"]] = np.nan
                frames.append(fr_)
            q_prev, q = q, q_next
            if nstep % report == 0 and nstep:
                if not np.all(np.isfinite(q)):
                    raise RuntimeError(f"solution diverged at step {nstep} (t = {t:.4f} s); increase the oversampling m")
                el = time.perf_counter() - t_start
                self.log(f"  step {nstep}/{nsteps} t={t:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(q).max():.3e}")
        if not np.all(np.isfinite(out)):
            raise RuntimeError("output is not finite; increase the oversampling m")
        fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
        return out, fields

    def _run_device(self, pr, m: int, nsamp: int):
        """The same loop with the tensor contraction and the recurrence on the device (PyTorch)."""
        t = self.torch
        dev = t.device(self.device)
        T64 = lambda a: t.as_tensor(np.asarray(a, dtype=np.float64), device=dev)
        dt, T, Om, nsteps, dmax, fld = pr["dt"], pr["T"], pr["Om"], pr["nsteps"], pr["dmax"], pr["fld"]
        a1, a2, b, w2, al, G, coef = (T64(x) for x in (pr["a1"], pr["a2"], pr["b"], pr["w2"], pr["al"], pr["G"], self.coef))
        schedule = [(t0, Ts, T64(fvec)) for t0, Ts, fvec in self.schedule]
        Hf = self.Hflat_t
        Gt = T64(self.G) if self.G is not None else None
        L = dmax + 1
        ar = t.arange(L, device=dev)
        ring = t.zeros((L, self.n_w), dtype=t.float64, device=dev)
        out = t.zeros(nsamp, dtype=t.float64, device=dev)
        q = t.zeros(self.n_w, dtype=t.float64, device=dev)
        q_prev = t.zeros_like(q)
        frames = [] if fld else None
        PhiG = T64(fld["PhiG"]) if fld else None
        invdt = 1.0 / dt  # 2 alpha (q+ - q-) / (2 dt)
        t_start = time.perf_counter()
        report = max(1, nsteps // 10)
        for nstep in range(nsteps + 1):
            tt = nstep * dt
            v = (Hf @ q).reshape(self.n_airy, self.n_w)
            sq = v @ q
            if Gt is None:
                Fext = v.T @ (coef * sq)
            else:
                Fext = v.T @ (coef * (sq + 2.0 * (Gt @ q))) + Gt.T @ (coef * sq)
            for t0, Ts, fvec in schedule:
                if t0 <= tt <= t0 + Ts:
                    Fext = Fext + fvec * (0.5 * (1.0 - np.cos(2 * np.pi * (tt - t0) / Ts)))
            q_next = a1 * q + a2 * q_prev + b * Fext
            qdd = Fext - w2 * q - al * ((q_next - q_prev) * invdt)
            ring[nstep % L] = qdd
            if nstep % m == 0:
                k = nstep // m
                if k < nsamp:
                    idx = (nstep - ar) % L
                    out[k] = t.sum(G * ring[idx])
            if frames is not None and nstep % fld["every"] == 0 and len(frames) < fld["nt"]:
                fr_ = (PhiG @ q).cpu().numpy().astype(np.float32)
                fr_[~fld["inside"]] = np.nan
                frames.append(fr_)
            q_prev, q = q, q_next
            if nstep % report == 0 and nstep:
                qc = q.cpu().numpy()
                if not np.all(np.isfinite(qc)):
                    raise RuntimeError(f"solution diverged at step {nstep} (t = {tt:.4f} s); increase the oversampling m")
                el = time.perf_counter() - t_start
                self.log(f"  step {nstep}/{nsteps} t={tt:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(qc).max():.3e}")
        out = out.cpu().numpy()
        if not np.all(np.isfinite(out)):
            raise RuntimeError("output is not finite; increase the oversampling m")
        fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
        return out, fields


def hardware_string(device: str) -> str:
    cpu = "unknown CPU"
    try:
        for line in open("/proc/cpuinfo"):
            if line.startswith("model name"):
                cpu = " ".join(line.split(":", 1)[1].split())
                break
    except OSError:
        pass
    s = f"{cpu} ({os.cpu_count()} threads), numpy/scipy for the setup"
    if device != "cpu":
        import torch

        s += f"; coupling tensor and time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
    else:
        s += " and the time stepping"
    return s


def parse_tag(tag: str):
    parts = dict((kv[0], int(kv[1:])) for kv in tag.split("-"))
    return parts["g"] * 1e-3, parts["w"], parts["a"], parts["m"]


def main():
    ch = load_challenge(Path(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]
    fields_for = set()
    force = None
    seconds = None

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

    v = take("--fields", lambda s: set(s.split(",")))
    fields_for = v or set()
    force = take("--force", float)
    duration = take("--duration", float)
    strikes_file = take("--strikes", str)
    strikes = json.load(open(strikes_file)) if strikes_file else None
    seconds = take("--seconds", float)
    device = take("--device", str) or "cpu"
    full = "--full" in args
    if full:
        args.remove("--full")
    tags = args or ["g8-w150-a150-m1", "g8-w300-a300-m1", "g4-w300-a300-m2", "g4-w400-a400-m2"]
    out = ch["output"]
    fs = out["sample_rate_hz"]
    nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
    log = lambda s: print(s, file=sys.stderr, flush=True)
    meta = []
    for tag in tags:
        spacing, n_w, n_a, m = parse_tag(tag)
        log(f"{tag}: force={force if force is not None else ch['parameters']['strike_peak_force']} N, duration={duration if duration is not None else ch['parameters']['strike_duration']} s")
        t0 = time.perf_counter()
        model = VonKarmanDisk(ch, spacing, n_w, n_a, force, log=log, symmetric_y=False if full else None, device=device, duration_override=duration, strikes=strikes)
        t1 = time.perf_counter()
        grid = None
        if tag in fields_for:
            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"])
        s, fields = model.run(fs, nsamp, m, grid, ch["fields"]["frame_rate_hz"] if grid else None)
        dt_total = time.perf_counter() - t0
        sid = f"vk-{tag}" + ("-" + Path(strikes_file).stem if strikes_file else "")
        out_dir.mkdir(parents=True, exist_ok=True)
        scale = out["signal_scale"] if out.get("signal_scale") else 1.0  # a draft challenge may have no scale yet
        wavfile.write(out_dir / f"{sid}.wav", fs, (s / scale).astype(np.float32))
        entry = {
            "solution_id": sid,
            "output": f"{out_dir.name}/{sid}.wav",
            "method": (("Marguerre shallow shell (Airy form) on a Cartesian grid masked to the disk, modal method on shell modes built from Bogner-Fox-Schmit plate modes" + (" with variable thickness" if model.profile.tapered else "") + ", Gautschi-type integrator exact in the linear part, retarded Rayleigh-type integral with cubic fractional delays") if model.profile is not None else
                       "von Karman plate (Airy form) on a Cartesian grid masked to the disk, modal method on Bogner-Fox-Schmit modes, Gautschi-type integrator exact in the linear part, retarded Rayleigh integral with cubic fractional delays"),
            "parameters": {
                "grid_spacing_m": spacing,
                "cells": model.plate.n_cells,
                "n_dof": model.plate.ndof,
                "transverse_modes": model.n_w,
                "basis_even_in_y": model.symmetric_y,
                "airy_modes": model.n_airy,
                "oversampling": m,
                "dt_s": 1.0 / (fs * m),
                "highest_transverse_mode_hz": round(float(model.omega.max() / (2 * np.pi)), 1),
            },
            "timing_seconds": round(dt_total, 3),
            "hardware": hardware_string(device),
            "peak_si": float(np.max(np.abs(s))),
            "setup_seconds": round(t1 - t0, 3),
            **({"strikes": strikes, "note": "experiment with a strike schedule, not a solution of the challenge"} if strikes else {}),
        }
        if fields is not None:
            import h5py

            fg = ch["fields"]
            with h5py.File(out_dir / f"{sid}.h5", "w") as fh:
                fh.attrs["challenge_id"] = ch["challenge_id"]
                fh.attrs["solution_id"] = sid
                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][: fields.shape[0]])
                fh.create_dataset("u", data=fields, compression="gzip", compression_opts=4, chunks=(8, fields.shape[1], fields.shape[2]))
            entry["fields"] = {"path": f"{out_dir.name}/{sid}.h5", "format": "hdf5"}
        meta.append(entry)
        log(f"{tag}: done, setup={t1 - t0:.1f}s total={dt_total:.1f}s peak={np.max(np.abs(s)):.4e} Pa")
        with open(out_dir / f"{sid}.meta.json", "w") as fh:
            json.dump(entry, fh, indent=1)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()