/**
 * WebGPU port of plate_fd_loop.m: the explicit time stepping of the plate,
 * run as compute shaders. Everything around the loop stays in numbl: the host
 * runs the unmodified plate_fd_setup.m, hands its struct in here as FdLoopSetup,
 * and passes the record this returns to the unmodified plate_fd_output.m, which
 * computes the stresses, the audio and the spectra in double precision.
 *
 * One time step is five compute dispatches, the dispatch boundaries being the
 * global synchronization points the stencils need:
 *   strain   e = G phi: every difference stencil of the energy method (second
 *            differences at the nodes, twists at the cells, first differences
 *            along the edges) as an exact integer stencil applied to phi
 *   force    at every node, G' applied to the weighted strains: the bending
 *            force, the prestress and spring force, and the Laplacian, each
 *            from its own weights. This is the assembled operator applied in
 *            its factored form. In single precision the assembled 13-point
 *            stencil loses the smooth modes to rounding (the biharmonic of a
 *            smooth field is a small residual of large terms, and the rounding
 *            of the field alone is percent of it at low wavenumbers); the
 *            factored form keeps them to about 1e-5, because a second
 *            difference of nearby values is nearly exact in f32 and the
 *            scales live in the weights, where rounding is harmless
 *   dots     one workgroup: the nodal forces f (strike, contacts, dashpots) from
 *            the total field w = R c + phi and its velocity, the reductions
 *            R' f, (AR)' phi and the support sums, the rigid-body update
 *            c -> cn (thread 0), and the series of this step (probe, rigid
 *            coordinates, support forces) averaged into their samples
 *   update   phin = c_new (2 phi - c_old phim + dt^2 (f - B - AR c)/M - dt^2 R cacc
 *            + 2 sigma1 dt (Lphi - Lphi_prev)), zeroed at nodes that are not
 *            unknowns; frames recorded
 *   reduce   the acceleration of the total field summed per listener delay
 *            bin into the audio samples; max |phin| plus the rigid extent for
 *            the blow-up test
 *
 * Steps are recorded in chunks of a few thousand into one command buffer each,
 * with the step index a chunk base (ctl[0], written by the host once per chunk)
 * plus a within-chunk offset from a dynamic-offset uniform; every recorded
 * dispatch is identical. Two chunks are in flight at a time and the blow-up
 * flag is read back after each, so a diverging run stops a chunk late rather
 * than at the end.
 *
 * Precision. The GPU works in single precision. The flexural field is carried
 * as (phi, d = phi - phim) and the rigid coordinates as (c, dc): the update is
 * algebraically the .m's,
 *   dn = (d + rest) - cs (2 d + rest),   phin = phi + dn,   cs = sigma0 dt / (1 + sigma0 dt),
 * with rest = dt^2 acc (the sigma1 loss being part of the nodal force), so that no factor of
 * the form 1 - epsilon is stored in f32 and the rounding of the field does not
 * feed back into the velocity. The split into rigid coordinates and flexural
 * remainder is the .m's own (plate_rigid_basis); bending is never applied to
 * the rigid part, whose forces come from the precomputed AR and Krr.
 */

export interface Triplets {
  /** 1-based row and column indices, as MATLAB's find returns them */
  i: Float64Array;
  j: Float64Array;
  v: Float64Array;
}

/** The fields of plate_fd_setup's struct that the loop reads; indices are 1-based as in the .m. */
export interface FdLoopSetup {
  nx: number;
  ny: number;
  nsteps: number;
  dx: number;
  dy: number;
  dt: number;
  fs: number;
  /** Number of output samples */
  na: number;
  keep: Float64Array;
  Mvec: Float64Array;
  /**
   * The operators in factored form (plate_fd_setup): Gs stacks the difference
   * stencils (ng rows over the N nodes), Gscale the scale of each row, and the
   * symmetric weights give KB = Gs' WB Gs, KG - diag(ks) = Gs' WG Gs and
   * Lm = diag(inva) Gs' WL Gs on the unknowns.
   */
  Gs: Triplets;
  Gscale: Float64Array;
  ng: number;
  WB: Triplets;
  WG: Triplets;
  WL: Triplets;
  inva: Float64Array;
  /** Linear springs, the diagonal part of KG */
  ks: Float64Array;
  nr: number;
  /** N x nr column-major */
  R: Float64Array;
  AR: Float64Array;
  /** nr x nr column-major */
  Krr: Float64Array;
  Ginv: Float64Array;
  kb: Float64Array;
  cb: Float64Array;
  gb: Float64Array;
  kt: Float64Array;
  ct: Float64Array;
  gt: Float64Array;
  cs: Float64Array;
  /** Static force at rest of the springs, the bottom contacts and the top contacts */
  fstatS: Float64Array;
  fstatB: Float64Array;
  fstatT: Float64Array;
  /** Supports as a CSR list of 1-based node indices, each reporting one kind of force (0 spring, 1 bottom, 2 top) */
  supPtr: Float64Array;
  supNodes: Float64Array;
  supKind: Float64Array;
  nsup: number;
  /** Force per node per unit F */
  sa: Float64Array;
  /** Force at the nsteps + 1 time levels */
  F: Float64Array;
  ip: number;
  jp: number;
  frameSteps: Float64Array;
  /** Listener entries: node (1-based), delay bin (1-based), weight */
  lisNode: Float64Array;
  lisBin: Float64Array;
  lisWgt: Float64Array;
  nbins: number;
  c0: Float64Array;
  phi0: Float64Array;
  sigma0: number;
  /** The sigma1 loss as a nodal force per unit Lm(w - wm): 2 sigma1 M / dt, per node */
  s1f: Float64Array;
  /** Lm R, N x nr column-major: the Laplacian per node of the rigid basis */
  LR: Float64Array;
  dt2: number;
  /** Weight of a step in its output sample (1 over the steps the sample contains), na entries */
  wsamp: Float64Array;
  blowup: number;
}

/** What plate_fd_loop.m returns, in single precision, before conversion for plate_fd_output.m. */
export interface FdLoopRecord {
  w_p: Float32Array;
  wxx_p: Float32Array;
  wyy_p: Float32Array;
  wxy_p: Float32Array;
  /** nr x na, row k at C[k * na ..] */
  C: Float32Array;
  /** nsup x na */
  Fsup: Float32Array;
  /** ny x nx x nf */
  Wf: Float32Array;
  /** nr x nf, row k at Cf[k * nf ..] */
  Cf: Float32Array;
  /** nbins x na */
  Pbins: Float32Array;
  divergedStep: number;
  gpuMs: number;
}

export interface GpuHandle {
  device: GPUDevice;
  adapter: string;
}

export function describeAdapter(adapter: GPUAdapter): string {
  const info = (adapter as { info?: Partial<GPUAdapterInfo> }).info;
  if (!info) return "unknown adapter";
  const name = [info.vendor, info.architecture].filter(Boolean).join(" ");
  return name || info.description || "unknown adapter";
}

/** The device request the browser worker and the Node test both make. */
export async function requestGpuDevice(): Promise<GpuHandle> {
  const gpu = (globalThis as { navigator?: { gpu?: GPU } }).navigator?.gpu;
  if (!gpu) throw new Error("WebGPU is not available");
  const adapter = await gpu.requestAdapter();
  if (!adapter) throw new Error("no WebGPU adapter");
  const device = await adapter.requestDevice({
    requiredLimits: {
      maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
      maxBufferSize: adapter.limits.maxBufferSize,
    },
  });
  return { device, adapter: describeAdapter(adapter) };
}

const WG = 64;
const RWG = 256;
const CHUNK = 2048;
const NONE = 0xffffffff;

const fmt = (x: number): string => {
  if (!Number.isFinite(x)) throw new Error("non-finite kernel constant");
  let s = x.toPrecision(9);
  if (!/[.e]/.test(s)) s += ".0";
  return s + "f";
};

interface Csr {
  ptr: Uint32Array;
  col: Uint32Array;
  val: Float32Array;
}

function toCsr(t: Triplets, n: number): Csr {
  const count = new Uint32Array(n);
  for (let p = 0; p < t.i.length; p++) count[t.i[p] - 1]++;
  const ptr = new Uint32Array(n + 1);
  for (let r = 0; r < n; r++) ptr[r + 1] = ptr[r] + count[r];
  const fill = ptr.slice(0, n);
  const col = new Uint32Array(t.i.length);
  const val = new Float32Array(t.i.length);
  for (let p = 0; p < t.i.length; p++) {
    const q = fill[t.i[p] - 1]++;
    col[q] = t.j[p] - 1;
    val[q] = t.v[p];
  }
  return { ptr, col, val };
}

/** Sparse product in double. */
function spmv(t: Triplets, x: Float64Array, n: number): Float64Array {
  const y = new Float64Array(n);
  for (let p = 0; p < t.i.length; p++) y[t.i[p] - 1] += t.v[p] * x[t.j[p] - 1];
  return y;
}

/** Transposed sparse product in double. */
function spmvT(t: Triplets, x: Float64Array, n: number): Float64Array {
  const y = new Float64Array(n);
  for (let p = 0; p < t.i.length; p++) y[t.j[p] - 1] += t.v[p] * x[t.i[p] - 1];
  return y;
}

/** The weight matrices as one CSR pattern with three value arrays. */
function weightsCsr(
  ng: number,
  scale: Float64Array,
  WB: Triplets,
  WG: Triplets,
  WL: Triplets,
): { ptr: Uint32Array; col: Uint32Array; vb: Float32Array; vg: Float32Array; vl: Float32Array } {
  const rows: Map<number, { b: number; g: number; l: number }>[] = Array.from(
    { length: ng },
    () => new Map(),
  );
  const add = (t: Triplets, key: "b" | "g" | "l") => {
    for (let p = 0; p < t.i.length; p++) {
      const r = t.i[p] - 1;
      const c = t.j[p] - 1;
      let e = rows[r].get(c);
      if (!e) {
        e = { b: 0, g: 0, l: 0 };
        rows[r].set(c, e);
      }
      // the stencils are integers: their scales move into the weights
      e[key] += t.v[p] * scale[r] * scale[c];
    }
  };
  add(WB, "b");
  add(WG, "g");
  add(WL, "l");
  const ptr = new Uint32Array(ng + 1);
  for (let r = 0; r < ng; r++) ptr[r + 1] = ptr[r] + rows[r].size;
  const nnz = ptr[ng];
  const col = new Uint32Array(nnz);
  const vb = new Float32Array(nnz);
  const vg = new Float32Array(nnz);
  const vl = new Float32Array(nnz);
  let q = 0;
  for (let r = 0; r < ng; r++)
    for (const [c, e] of rows[r]) {
      col[q] = c;
      vb[q] = e.b;
      vg[q] = e.g;
      vl[q] = e.l;
      q++;
    }
  return { ptr, col, vb, vg, vl };
}

/** Listener entries grouped by delay bin: entry -> node, entry -> weight, in bin order. */
function binLists(
  bins: Float64Array,
  nodes: Float64Array,
  wgt: Float64Array,
  nbins: number,
): { ptr: Uint32Array; node: Uint32Array; wgt: Float32Array } {
  const n = bins.length;
  const count = new Uint32Array(nbins);
  for (let e = 0; e < n; e++) count[bins[e] - 1]++;
  const ptr = new Uint32Array(nbins + 1);
  for (let b = 0; b < nbins; b++) ptr[b + 1] = ptr[b] + count[b];
  const fill = ptr.slice(0, nbins);
  const node = new Uint32Array(n);
  const w = new Float32Array(n);
  for (let e = 0; e < n; e++) {
    const q = fill[bins[e] - 1]++;
    node[q] = nodes[e] - 1;
    w[q] = wgt[e];
  }
  return { ptr, node, wgt: w };
}

function pack<T extends Uint32Array | Float32Array>(
  parts: [string, ArrayLike<number>][],
  make: (n: number) => T,
): { data: T; off: Record<string, number> } {
  const off: Record<string, number> = {};
  let cursor = 0;
  for (const [name, arr] of parts) {
    off[name] = cursor;
    cursor += arr.length;
  }
  const data = make(Math.max(1, cursor));
  for (const [name, arr] of parts) data.set(arr, off[name]);
  return { data, off };
}

function buildShaders(su: FdLoopSetup) {
  const nx = su.nx | 0;
  const ny = su.ny | 0;
  const N = nx * ny;
  const nsteps = su.nsteps | 0;
  const nbins = su.nbins | 0;
  const na = su.na | 0;
  const nf = su.frameSteps.length;
  const nr = su.nr | 0;
  const nsup = su.nsup | 0;
  if (nr !== 0 && nr !== 3) throw new Error("the loop expects 0 or 3 rigid-body coordinates");
  if (su.sa.length !== N || su.keep.length !== N || su.Mvec.length !== N)
    throw new Error("setup arrays do not match the grid");
  if (su.F.length < nsteps) throw new Error("force history shorter than the step count");
  if (su.wsamp.length !== na) throw new Error("sample weights do not match the sample count");
  if (nr > 0 && (su.R.length !== N * nr || su.AR.length !== N * nr || su.LR.length !== N * nr))
    throw new Error("rigid-body arrays do not match the grid");
  if (su.s1f.length !== N) throw new Error("loss array does not match the grid");
  if (nsup > 8) throw new Error("at most 8 supports");

  // ---- per-step tables, decided in f64 exactly as the .m decides them
  const frameOf = new Uint32Array(nsteps).fill(NONE);
  for (let kf = 0; kf < nf; kf++) {
    const step = su.frameSteps[kf];
    if (step >= 1 && step <= nsteps) frameOf[step - 1] = kf;
  }
  const iaOf = new Uint32Array(nsteps);
  for (let k = 0; k < nsteps; k++) {
    const ia = Math.floor(k * su.dt * su.fs); // ia = floor((nstep - 1) dt fs) + 1, made 0-based
    if (ia >= na) throw new Error("output sample index out of range");
    iaOf[k] = ia;
  }
  const bl = binLists(su.lisBin, su.lisNode, su.lisWgt, nbins);
  const ng = su.ng | 0;
  if (su.Gscale.length !== ng) throw new Error("strain scales do not match the strain count");
  // the stencils as exact integers (their scale is folded into the weights), and the transpose
  const Gi: Triplets = {
    i: su.Gs.i,
    j: su.Gs.j,
    v: Float64Array.from(su.Gs.v, (v, p) => Math.round(v / su.Gscale[su.Gs.i[p] - 1])),
  };
  for (let p = 0; p < Gi.v.length; p++)
    if (Math.abs(Gi.v[p] - su.Gs.v[p] / su.Gscale[su.Gs.i[p] - 1]) > 1e-6)
      throw new Error("a stencil coefficient is not an integer");
  const G = toCsr(Gi, ng);
  const GT = toCsr({ i: Gi.j, j: Gi.i, v: Gi.v }, N);
  const W = weightsCsr(ng, su.Gscale, su.WB, su.WG, su.WL);
  const minv = Float32Array.from(su.Mvec, (m) => 1 / m);
  const f32 = (a: ArrayLike<number>) => Float32Array.from(a);
  const rmax = new Float32Array(3);
  for (let k = 0; k < nr; k++) {
    let m = 0;
    for (let i = 0; i < N; i++) m = Math.max(m, Math.abs(su.R[i + k * N]));
    rmax[k] = m;
  }
  const supPtr = Uint32Array.from(su.supPtr, (p) => p - 1);
  const supNodes = Uint32Array.from(su.supNodes, (p) => p - 1);
  const supKind = Uint32Array.from(su.supKind);

  const opsParts: [string, ArrayLike<number>][] = [
    ["SA", f32(su.sa)],
    ["F", f32(su.F.subarray(0, nsteps))],
    ["MINV", minv],
    ["KEEP", f32(su.keep)],
    ["INVA", f32(su.inva)],
    ["KSD", f32(su.ks)],
    ["GV", G.val],
    ["GTV", GT.val],
    ["WBV", W.vb],
    ["WGV", W.vg],
    ["WLV", W.vl],
    ["R", f32(nr ? su.R : [0])],
    ["AR", f32(nr ? su.AR : [0])],
    ["LR", f32(nr ? su.LR : [0])],
    ["S1F", f32(su.s1f)],
    ["KRR", f32(nr ? su.Krr : [0])],
    ["GINV", f32(nr ? su.Ginv : [0])],
    ["KB", f32(su.kb)],
    ["CB", f32(su.cb)],
    ["GB", f32(su.gb)],
    ["KT", f32(su.kt)],
    ["CT", f32(su.ct)],
    ["GT", f32(su.gt)],
    ["CS", f32(su.cs)],
    ["KS", f32(su.ks)],
    ["FSTATS", f32(su.fstatS)],
    ["FSTATB", f32(su.fstatB)],
    ["FSTATT", f32(su.fstatT)],
    ["WGT", bl.wgt],
    ["WSAMP", f32(su.wsamp)],
  ];
  const idxParts: [string, ArrayLike<number>][] = [
    ["GPTR", G.ptr],
    ["GCOL", G.col],
    ["GTPTR", GT.ptr],
    ["GTCOL", GT.col],
    ["WPTR", W.ptr],
    ["WCOL", W.col],
    ["BINPTR", bl.ptr],
    ["BINNODE", bl.node],
    ["SUPPTR", supPtr],
    ["SUPNODES", supNodes],
    ["SUPKIND", supKind],
    ["FRAME", frameOf],
    ["IA", iaOf],
  ];
  const ops = pack(opsParts, (n) => new Float32Array(n));
  const idx = pack(idxParts, (n) => new Uint32Array(n));

  // ---- state buffer: two slots each of phi, d = phi - phim and Lphi; B; f; fsup; rigid scalars
  const ST_PHI = 0;
  const ST_D = 2 * N;
  const ST_L = 4 * N;
  const ST_B = 6 * N;
  const ST_F = 7 * N;
  const ST_FS = 8 * N; // three arrays: spring, bottom and top support forces
  const ST_AUX = 11 * N; // c[2][3], dc[2][3], cacc[3], reserve
  const ST_E = ST_AUX + 32; // the strains, ng of them
  const stLen = ST_E + ng;
  const AUX_C = 0; // c slot s at AUX_C + 3 s
  const AUX_DC = 6;
  const AUX_CACC = 12;

  // ---- output buffer
  const OUT_WP = 0;
  const OUT_WXX = na;
  const OUT_WYY = 2 * na;
  const OUT_WXY = 3 * na;
  const OUT_C = 4 * na;
  const OUT_FS = OUT_C + nr * na;
  const OUT_PB = OUT_FS + nsup * na;
  const OUT_CF = OUT_PB + nbins * na;
  const outLen = OUT_CF + nr * nf;

  const s0dt = su.sigma0 * su.dt;
  const c: Record<string, number> = {
    INV_DX2: 1 / (su.dx * su.dx),
    INV_DY2: 1 / (su.dy * su.dy),
    INV_DXDY4: 1 / (4 * su.dx * su.dy),
    INV_DT: 1 / su.dt,
    DT2: su.dt2,
    CSIG: s0dt / (1 + s0dt),
    INV_DT2: 1 / su.dt2,
    BLOWUP: su.blowup,
    RMAX0: rmax[0],
    RMAX1: rmax[1],
    RMAX2: rmax[2],
  };
  const ip = su.ip - 1;
  const jp = su.jp - 1;
  if (ip < 1 || ip > ny - 2 || jp < 1 || jp > nx - 2)
    throw new Error("probe is not an interior node");

  const header = `
const NX = ${nx}u; const NY = ${ny}u; const N = ${N}u;
const NSTEPS = ${nsteps}u; const NA = ${na}u;
const NBINS = ${nbins}u; const NR = ${nr}u; const NSUP = ${nsup}u; const NF = ${nf}u; const NG = ${ng}u;
const IPROBE = ${ip + jp * ny}u;
${Object.entries(c)
  .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
  .join("\n")}
${Object.entries(ops.off)
  .map(([k, v]) => `const OP_${k} = ${v}u;`)
  .join("\n")}
${Object.entries(idx.off)
  .map(([k, v]) => `const IX_${k} = ${v}u;`)
  .join("\n")}
const ST_PHI = ${ST_PHI}u; const ST_D = ${ST_D}u; const ST_L = ${ST_L}u; const ST_B = ${ST_B}u;
const ST_F = ${ST_F}u; const ST_FS = ${ST_FS}u; const ST_AUX = ${ST_AUX}u; const ST_E = ${ST_E}u;
const AUX_C = ${AUX_C}u; const AUX_DC = ${AUX_DC}u; const AUX_CACC = ${AUX_CACC}u;
const OUT_WP = ${OUT_WP}u; const OUT_WXX = ${OUT_WXX}u; const OUT_WYY = ${OUT_WYY}u; const OUT_WXY = ${OUT_WXY}u;
const OUT_C = ${OUT_C}u; const OUT_FS = ${OUT_FS}u; const OUT_PB = ${OUT_PB}u; const OUT_CF = ${OUT_CF}u;

@group(0) @binding(0) var<storage, read> ops: array<f32>;
@group(0) @binding(1) var<storage, read> idx: array<u32>;
@group(0) @binding(2) var<storage, read_write> st: array<f32>;
@group(0) @binding(3) var<storage, read_write> frames: array<f32>;
@group(0) @binding(4) var<storage, read_write> outb: array<f32>;
@group(0) @binding(5) var<storage, read_write> ctl: array<u32>;
@group(0) @binding(6) var<uniform> uni: vec4<u32>;

// The blow-up flag holds the 1-based step at which the field passed the limit;
// the steps after it do nothing, the step itself still records, as the .m does.
fn stopped(k: u32) -> bool {
  let d = ctl[1];
  return d != 0u && d <= k;
}
// Slot of the current field (phi, d, Lphi, c, dc) at step k and of the one being written
fn cur(k: u32) -> u32 { return (k & 1u) * N; }
fn nxt(k: u32) -> u32 { return ((k + 1u) & 1u) * N; }
fn curc(k: u32) -> u32 { return (k & 1u) * 3u; }
fn nxtc(k: u32) -> u32 { return ((k + 1u) & 1u) * 3u; }
// Rigid coordinate j of the current step, zero without rigid coordinates
fn cget(k: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return st[ST_AUX + AUX_C + curc(k) + j]; }
fn dcget(k: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return st[ST_AUX + AUX_DC + curc(k) + j]; }
fn rget(i: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return ops[OP_R + i + j * N]; }
// Total field and its backward-difference velocity at node i, step k
fn wtot(k: u32, i: u32) -> f32 {
  var w = st[ST_PHI + cur(k) + i];
  if (NR > 0u) { w = w + rget(i, 0u) * cget(k, 0u) + rget(i, 1u) * cget(k, 1u) + rget(i, 2u) * cget(k, 2u); }
  return w;
}
fn vtot(k: u32, i: u32) -> f32 {
  var d = st[ST_D + cur(k) + i];
  if (NR > 0u) { d = d + rget(i, 0u) * dcget(k, 0u) + rget(i, 1u) * dcget(k, 1u) + rget(i, 2u) * dcget(k, 2u); }
  return d * INV_DT;
}
`;

  // Every difference stencil of the energy method applied to phi; the stencils
  // are integers, so a strain is a sum of nearby field values with small
  // integer coefficients, nearly exact in f32.
  const strain = `${header}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = ctl[0] + uni.x;
  if (k >= NSTEPS) { return; }
  let q = gid.x;
  if (q >= NG) { return; }
  if (stopped(k)) { return; }
  let pS = ST_PHI + cur(k);
  var e: f32 = 0.0;
  for (var p = idx[IX_GPTR + q]; p < idx[IX_GPTR + q + 1u]; p = p + 1u) {
    e = e + ops[OP_GV + p] * st[pS + idx[IX_GCOL + p]];
  }
  st[ST_E + q] = e;
}
`;

  // G' applied to the weighted strains: bending, prestress and Laplacian at once
  const force = `${header}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = ctl[0] + uni.x;
  if (k >= NSTEPS) { return; }
  let i = gid.x;
  if (i >= N) { return; }
  if (stopped(k)) { return; }
  var yb: f32 = 0.0; var yg: f32 = 0.0; var yl: f32 = 0.0;
  for (var p = idx[IX_GTPTR + i]; p < idx[IX_GTPTR + i + 1u]; p = p + 1u) {
    let q = idx[IX_GTCOL + p];
    let g = ops[OP_GTV + p];
    var sb: f32 = 0.0; var sg: f32 = 0.0; var sl: f32 = 0.0;
    for (var r = idx[IX_WPTR + q]; r < idx[IX_WPTR + q + 1u]; r = r + 1u) {
      let e = st[ST_E + idx[IX_WCOL + r]];
      sb = sb + ops[OP_WBV + r] * e;
      sg = sg + ops[OP_WGV + r] * e;
      sl = sl + ops[OP_WLV + r] * e;
    }
    yb = yb + g * sb; yg = yg + g * sg; yl = yl + g * sl;
  }
  let phi = st[ST_PHI + cur(k) + i];
  st[ST_B + i] = yb + (yg + ops[OP_KSD + i] * phi);
  st[ST_L + cur(k) + i] = yl * ops[OP_INVA + i];
}
`;

  // Nodal forces, the reductions of the rigid-body equations and the supports,
  // the rigid update, and this step's series, all in one workgroup.
  const dots = `${header}
var<workgroup> sh0: array<f32, ${RWG}>;
var<workgroup> sh1: array<f32, ${RWG}>;
var<workgroup> sh2: array<f32, ${RWG}>;
var<workgroup> sh3: array<f32, ${RWG}>;
var<workgroup> sh4: array<f32, ${RWG}>;
var<workgroup> sh5: array<f32, ${RWG}>;
var<workgroup> shs: array<f32, ${RWG}>;
fn reduce6(lid: u32) {
  for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
    if (lid < s) {
      sh0[lid] = sh0[lid] + sh0[lid + s]; sh1[lid] = sh1[lid] + sh1[lid + s]; sh2[lid] = sh2[lid] + sh2[lid + s];
      sh3[lid] = sh3[lid] + sh3[lid + s]; sh4[lid] = sh4[lid] + sh4[lid + s]; sh5[lid] = sh5[lid] + sh5[lid + s];
    }
    workgroupBarrier();
  }
}
@compute @workgroup_size(${RWG})
fn main(@builtin(local_invocation_index) lid: u32) {
  let k = ctl[0] + uni.x;
  let stop = stopped(k) || k >= NSTEPS;
  let pS = ST_PHI + cur(k);
  let lS = ST_L + cur(k); let lmS = ST_L + nxt(k);
  let Fk = ops[OP_F + min(k, NSTEPS - 1u)];
  var dc0: f32 = dcget(k, 0u); var dc1: f32 = dcget(k, 1u); var dc2: f32 = dcget(k, 2u);
  // nodal forces about rest, and the partial sums R' f and (AR)' phi
  var a0: f32 = 0.0; var a1: f32 = 0.0; var a2: f32 = 0.0;
  var g0: f32 = 0.0; var g1: f32 = 0.0; var g2: f32 = 0.0;
  for (var i = lid; i < N; i = i + ${RWG}u) {
    let w = wtot(k, i);
    let v = vtot(k, i);
    // the sigma1 loss on the total velocity: 2 sigma1 M Lm (w - wm) / dt
    var dlw = st[lS + i] - st[lmS + i];
    if (NR > 0u) { dlw = dlw + ops[OP_LR + i] * dc0 + ops[OP_LR + N + i] * dc1 + ops[OP_LR + 2u * N + i] * dc2; }
    var f = Fk * ops[OP_SA + i] + ops[OP_S1F + i] * dlw;
    var fsS = ops[OP_FSTATS + i] - ops[OP_KS + i] * w;
    var fsB = ops[OP_FSTATB + i];
    var fsT = ops[OP_FSTATT + i];
    let kb = ops[OP_KB + i];
    if (kb > 0.0) {
      let gb = ops[OP_GB + i];
      let comp = gb - w;
      var fb = kb * max(comp, 0.0) - kb * max(gb, 0.0);
      if (comp > 0.0) { fb = fb - ops[OP_CB + i] * v; }
      f = f + fb; fsB = fsB + fb;
    }
    let kt = ops[OP_KT + i];
    if (kt > 0.0) {
      let gt = ops[OP_GT + i];
      let comp = gt + w;
      var ft = -kt * max(comp, 0.0) + kt * max(gt, 0.0);
      if (comp > 0.0) { ft = ft - ops[OP_CT + i] * v; }
      f = f + ft; fsT = fsT + ft;
    }
    let cs = ops[OP_CS + i];
    if (cs > 0.0) { let fl = -cs * v; f = f + fl; fsS = fsS + fl; }
    st[ST_F + i] = f;
    st[ST_FS + i] = fsS;
    st[ST_FS + N + i] = fsB;
    st[ST_FS + 2u * N + i] = fsT;
    if (NR > 0u) {
      let phi = st[pS + i];
      a0 = a0 + rget(i, 0u) * f; a1 = a1 + rget(i, 1u) * f; a2 = a2 + rget(i, 2u) * f;
      g0 = g0 + ops[OP_AR + i] * phi; g1 = g1 + ops[OP_AR + N + i] * phi; g2 = g2 + ops[OP_AR + 2u * N + i] * phi;
    }
  }
  sh0[lid] = a0; sh1[lid] = a1; sh2[lid] = a2; sh3[lid] = g0; sh4[lid] = g1; sh5[lid] = g2;
  workgroupBarrier();
  reduce6(lid);
  // rigid-body update and the rigid series, on one thread
  if (lid == 0u && NR > 0u) {
    let cS = ST_AUX + AUX_C + curc(k); let cnS = ST_AUX + AUX_C + nxtc(k);
    let dS = ST_AUX + AUX_DC + curc(k); let dnS = ST_AUX + AUX_DC + nxtc(k);
    var rf = array<f32, 3>(sh0[0], sh1[0], sh2[0]);
    var g = array<f32, 3>(sh3[0], sh4[0], sh5[0]);
    var cc = array<f32, 3>(st[cS], st[cS + 1u], st[cS + 2u]);
    var dc = array<f32, 3>(st[dS], st[dS + 1u], st[dS + 2u]);
    var tmp = array<f32, 3>(0.0, 0.0, 0.0);
    for (var p = 0u; p < 3u; p = p + 1u) {
      var s = rf[p] - g[p];
      for (var q = 0u; q < 3u; q = q + 1u) { s = s - ops[OP_KRR + p + 3u * q] * cc[q]; }
      tmp[p] = s;
    }
    if (!stop) {
      let ia = idx[IX_IA + k];
      for (var p = 0u; p < 3u; p = p + 1u) {
        var acc: f32 = 0.0;
        for (var q = 0u; q < 3u; q = q + 1u) { acc = acc + ops[OP_GINV + p + 3u * q] * tmp[q]; }
        let rest = DT2 * acc;
        let dn = (dc[p] + rest) - CSIG * (2.0 * dc[p] + rest);
        st[ST_AUX + AUX_CACC + p] = acc;
        st[dnS + p] = dn;
        st[cnS + p] = cc[p] + dn;
        outb[OUT_C + p * NA + ia] = outb[OUT_C + p * NA + ia] + cc[p] * ops[OP_WSAMP + ia];
      }
    }
  }
  // the probe series of this step
  if (lid == 0u && !stop) {
    let ia = idx[IX_IA + k];
    let cw = st[pS + IPROBE];
    let e = st[pS + IPROBE + NY]; let w = st[pS + IPROBE - NY];
    let n = st[pS + IPROBE + 1u]; let so = st[pS + IPROBE - 1u];
    let ne = st[pS + IPROBE + NY + 1u]; let nw = st[pS + IPROBE - NY + 1u];
    let se = st[pS + IPROBE + NY - 1u]; let sw = st[pS + IPROBE - NY - 1u];
    outb[OUT_WP + ia] = outb[OUT_WP + ia] + wtot(k, IPROBE) * ops[OP_WSAMP + ia];
    outb[OUT_WXX + ia] = outb[OUT_WXX + ia] + (e - 2.0 * cw + w) * INV_DX2 * ops[OP_WSAMP + ia];
    outb[OUT_WYY + ia] = outb[OUT_WYY + ia] + (n - 2.0 * cw + so) * INV_DY2 * ops[OP_WSAMP + ia];
    outb[OUT_WXY + ia] = outb[OUT_WXY + ia] + (ne - se - nw + sw) * INV_DXDY4 * ops[OP_WSAMP + ia];
  }
  // support forces, one reduction per support
  for (var q = 0u; q < NSUP; q = q + 1u) {
    workgroupBarrier();
    let p0 = idx[IX_SUPPTR + q]; let p1 = idx[IX_SUPPTR + q + 1u];
    let base = ST_FS + idx[IX_SUPKIND + q] * N;
    var a: f32 = 0.0;
    for (var p = p0 + lid; p < p1; p = p + ${RWG}u) { a = a + st[base + idx[IX_SUPNODES + p]]; }
    shs[lid] = a;
    workgroupBarrier();
    for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
      if (lid < s) { shs[lid] = shs[lid] + shs[lid + s]; }
      workgroupBarrier();
    }
    if (lid == 0u && !stop) {
      let ia = idx[IX_IA + k];
      outb[OUT_FS + q * NA + ia] = outb[OUT_FS + q * NA + ia] + shs[0] * ops[OP_WSAMP + ia];
    }
  }
}
`;

  const update = `${header}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = ctl[0] + uni.x;
  if (k >= NSTEPS) { return; }
  let i = gid.x;
  if (i >= N) { return; }
  if (stopped(k)) { return; }
  let pS = ST_PHI + cur(k); let pnS = ST_PHI + nxt(k);
  let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
  var fi = st[ST_F + i] - st[ST_B + i];
  if (NR > 0u) {
    // the force of the rigid displacement on the plate, AR c, with c of the current step
    fi = fi - ops[OP_AR + i] * cget(k, 0u) - ops[OP_AR + N + i] * cget(k, 1u) - ops[OP_AR + 2u * N + i] * cget(k, 2u);
  }
  var acc = fi * ops[OP_MINV + i];
  if (NR > 0u) {
    acc = acc - rget(i, 0u) * st[ST_AUX + AUX_CACC] - rget(i, 1u) * st[ST_AUX + AUX_CACC + 1u] - rget(i, 2u) * st[ST_AUX + AUX_CACC + 2u];
  }
  let rest = DT2 * acc;
  let d = st[dS + i];
  var dn = (d + rest) - CSIG * (2.0 * d + rest);
  var pn = st[pS + i] + dn;
  let kp = ops[OP_KEEP + i];
  pn = pn * kp; dn = dn * kp;
  st[pnS + i] = pn;
  st[dnS + i] = dn;
  let f = idx[IX_FRAME + k];
  if (f != ${NONE}u) {
    frames[f * N + i] = pn;
    if (NR > 0u && i < 3u) { outb[OUT_CF + i * NF + f] = st[ST_AUX + AUX_C + nxtc(k) + i]; }
  }
}
`;

  // Workgroups 0..NBINS-1 sum the weighted acceleration of their bin's entries
  // into the audio sample of this step; workgroup NBINS takes the extent of the
  // new field for the blow-up test.
  const reduce = `${header}
var<workgroup> sh: array<f32, ${RWG}>;
@compute @workgroup_size(${RWG})
fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wgid: vec3<u32>) {
  let k = ctl[0] + uni.x;
  let stop = stopped(k) || k >= NSTEPS;
  let pnS = ST_PHI + nxt(k);
  let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
  let wg = wgid.x;
  if (wg < NBINS) {
    var dd = array<f32, 3>(0.0, 0.0, 0.0);
    if (NR > 0u) {
      for (var p = 0u; p < 3u; p = p + 1u) { dd[p] = st[ST_AUX + AUX_DC + nxtc(k) + p] - st[ST_AUX + AUX_DC + curc(k) + p]; }
    }
    let p0 = idx[IX_BINPTR + wg]; let p1 = idx[IX_BINPTR + wg + 1u];
    var a: f32 = 0.0;
    for (var p = p0 + lid; p < p1; p = p + ${RWG}u) {
      let i = idx[IX_BINNODE + p];
      var accw = st[dnS + i] - st[dS + i];
      if (NR > 0u) { accw = accw + rget(i, 0u) * dd[0] + rget(i, 1u) * dd[1] + rget(i, 2u) * dd[2]; }
      a = a + ops[OP_WGT + p] * (accw * INV_DT2);
    }
    sh[lid] = a;
    workgroupBarrier();
    for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
      if (lid < s) { sh[lid] = sh[lid] + sh[lid + s]; }
      workgroupBarrier();
    }
    if (lid == 0u && !stop) {
      let ia = idx[IX_IA + k];
      let o = OUT_PB + wg + ia * NBINS;
      outb[o] = outb[o] + sh[0] * ops[OP_WSAMP + ia];
    }
  } else {
    var m: f32 = 0.0;
    for (var i = lid; i < N; i = i + ${RWG}u) { m = max(m, abs(st[pnS + i])); }
    sh[lid] = m;
    workgroupBarrier();
    for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
      if (lid < s) { sh[lid] = max(sh[lid], sh[lid + s]); }
      workgroupBarrier();
    }
    if (lid == 0u && !stop) {
      var peak = sh[0];
      if (NR > 0u) {
        let cn = ST_AUX + AUX_C + nxtc(k);
        peak = peak + abs(st[cn]) * RMAX0 + abs(st[cn + 1u]) * RMAX1 + abs(st[cn + 2u]) * RMAX2;
      }
      if (!(peak < BLOWUP)) { ctl[1] = k + 1u; }
    }
  }
}
`;

  const sources = [strain, force, dots, update, reduce];
  const groups = [Math.ceil(ng / WG), Math.ceil(N / WG), 1, Math.ceil(N / WG), nbins + 1];

  // ---- initial state, in f64 as the .m has it
  const st0 = new Float32Array(stLen);
  const e0 = spmv(su.Gs, su.phi0, ng);
  const wl0 = spmv(su.WL, e0, ng);
  const lphi0 = spmvT(su.Gs, wl0, N);
  for (let i = 0; i < N; i++) lphi0[i] *= su.inva[i];
  for (let i = 0; i < N; i++) {
    st0[ST_PHI + i] = su.phi0[i];
    st0[ST_L + i] = lphi0[i];
    st0[ST_L + N + i] = lphi0[i];
  }
  for (let j = 0; j < nr; j++) st0[ST_AUX + AUX_C + j] = su.c0[j];

  return {
    N,
    nf,
    nr,
    nsup,
    nbins,
    na,
    outLen,
    stLen,
    st0,
    ops: ops.data,
    idx: idx.data,
    sources,
    groups,
    out: { OUT_WP, OUT_WXX, OUT_WYY, OUT_WXY, OUT_C, OUT_FS, OUT_PB, OUT_CF },
  };
}

export async function runFdLoopGpu(
  device: GPUDevice,
  su: FdLoopSetup,
  onProgress?: (stepsDone: number, nsteps: number) => void,
): Promise<FdLoopRecord> {
  const nsteps = su.nsteps | 0;
  const b = buildShaders(su);

  const bgl = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
      { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
      { binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
      {
        binding: 6,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 16 },
      },
    ],
  });
  const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
  const pipelines: GPUComputePipeline[] = [];
  for (const code of b.sources) {
    const module = device.createShaderModule({ code });
    const info = await module.getCompilationInfo();
    const errs = info.messages.filter((m) => m.type === "error");
    if (errs.length)
      throw new Error(
        "WGSL: " + errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join("; "),
      );
    pipelines.push(
      device.createComputePipeline({ layout, compute: { module, entryPoint: "main" } }),
    );
  }

  const mk = (size: number, usage: number) =>
    device.createBuffer({ size: Math.max(16, size), usage });
  const opsBuf = mk(b.ops.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(opsBuf, 0, b.ops);
  const idxBuf = mk(b.idx.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(idxBuf, 0, b.idx);
  const stBuf = mk(b.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(stBuf, 0, b.st0);
  const framesBuf = mk(b.N * b.nf * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const outBuf = mk(b.outLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const ctlBuf = mk(16, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
  const uniBuf = mk(CHUNK * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
  {
    const uni = new Uint32Array(CHUNK * 64);
    for (let s = 0; s < CHUNK; s++) uni[s * 64] = s;
    device.queue.writeBuffer(uniBuf, 0, uni);
  }
  const ctlStage = [0, 1].map(() => mk(16, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST));
  const readFrames = mk(b.N * b.nf * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
  const readOut = mk(b.outLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
  const all = [
    opsBuf,
    idxBuf,
    stBuf,
    framesBuf,
    outBuf,
    ctlBuf,
    uniBuf,
    ...ctlStage,
    readFrames,
    readOut,
  ];

  const bindGroup = device.createBindGroup({
    layout: bgl,
    entries: [
      { binding: 0, resource: { buffer: opsBuf } },
      { binding: 1, resource: { buffer: idxBuf } },
      { binding: 2, resource: { buffer: stBuf } },
      { binding: 3, resource: { buffer: framesBuf } },
      { binding: 4, resource: { buffer: outBuf } },
      { binding: 5, resource: { buffer: ctlBuf } },
      { binding: 6, resource: { buffer: uniBuf, offset: 0, size: 16 } },
    ],
  });

  try {
    const t0 = performance.now();
    let divergedStep = 0;
    const pending: { done: Promise<undefined>; stage: GPUBuffer; kend: number }[] = [];
    const settle = async () => {
      const p = pending.shift()!;
      await p.done;
      await p.stage.mapAsync(GPUMapMode.READ);
      const flag = new Uint32Array(p.stage.getMappedRange().slice(0))[1];
      p.stage.unmap();
      onProgress?.(p.kend, nsteps);
      if (flag !== 0) divergedStep = flag;
    };
    let kbase = 0;
    let c = 0;
    while (kbase < nsteps && divergedStep === 0) {
      const n = Math.min(CHUNK, nsteps - kbase);
      device.queue.writeBuffer(ctlBuf, 0, Uint32Array.of(kbase));
      const enc = device.createCommandEncoder();
      const pass = enc.beginComputePass();
      for (let s = 0; s < n; s++) {
        pass.setBindGroup(0, bindGroup, [s * 256]);
        for (let q = 0; q < pipelines.length; q++) {
          pass.setPipeline(pipelines[q]);
          pass.dispatchWorkgroups(b.groups[q]);
        }
      }
      pass.end();
      const stage = ctlStage[c % 2];
      enc.copyBufferToBuffer(ctlBuf, 0, stage, 0, 16);
      device.queue.submit([enc.finish()]);
      pending.push({ done: device.queue.onSubmittedWorkDone(), stage, kend: kbase + n });
      kbase += n;
      c++;
      if (pending.length >= 2) await settle();
    }
    while (pending.length) await settle();

    const enc = device.createCommandEncoder();
    enc.copyBufferToBuffer(framesBuf, 0, readFrames, 0, Math.max(16, b.N * b.nf * 4));
    enc.copyBufferToBuffer(outBuf, 0, readOut, 0, Math.max(16, b.outLen * 4));
    device.queue.submit([enc.finish()]);
    await readFrames.mapAsync(GPUMapMode.READ);
    await readOut.mapAsync(GPUMapMode.READ);
    const Wf = new Float32Array(readFrames.getMappedRange().slice(0)).subarray(0, b.N * b.nf);
    const out = new Float32Array(readOut.getMappedRange().slice(0));
    readFrames.unmap();
    readOut.unmap();
    const gpuMs = performance.now() - t0;
    const o = b.out;
    const na = b.na;
    return {
      w_p: out.slice(o.OUT_WP, o.OUT_WP + na),
      wxx_p: out.slice(o.OUT_WXX, o.OUT_WXX + na),
      wyy_p: out.slice(o.OUT_WYY, o.OUT_WYY + na),
      wxy_p: out.slice(o.OUT_WXY, o.OUT_WXY + na),
      C: out.slice(o.OUT_C, o.OUT_C + b.nr * na),
      Fsup: out.slice(o.OUT_FS, o.OUT_FS + b.nsup * na),
      Wf: Float32Array.from(Wf),
      Cf: out.slice(o.OUT_CF, o.OUT_CF + b.nr * b.nf),
      Pbins: out.slice(o.OUT_PB, o.OUT_PB + b.nbins * na),
      divergedStep,
      gpuMs,
    };
  } finally {
    for (const buf of all) buf.destroy();
  }
}