import type { BeadSystem, Excitation, SimOptions, SimResult } from './types';
import { timeStep, type Progress } from './solver';

// The same scheme as simulate() in solver.ts (velocity Verlet, same time step,
// rigid-motion removal, contact, delayed pickup), run on the GPU in single
// precision. Each phase of a time step is its own small kernel, and a command
// buffer holds the dispatches for many samples; time and sample counters live
// on the GPU, so nothing returns to JavaScript until a chunk is done. (One
// persistent workgroup running whole time steps was tried first: its large
// kernel made every barrier cost several microseconds on an Intel iGPU, and it
// could use only one of the GPU's subslices.) Forces are gathered per bead from
// fixed-width adjacency lists, so no atomics are needed. Sums over beads go
// through per-workgroup partials, which each workgroup of the next kernel adds
// up itself (cheaper than a dispatch of its own). Three buffers only: the
// per-dispatch cost grows with the number bound. To keep small vibrations from vanishing in
// f32, the state is the displacement from rest, and spring extensions and
// hinge angles are built from precomputed rest differences.

const WGB = 64; // invocations per workgroup

interface Packed {
  wgsl: string;
  ro: Uint32Array; // parameters (f32 bits) and indices
  rw: Float32Array; // state, shared scalars, per-workgroup partial sums
  outLen: number; // audio, contact force, frame scalars, mallet positions, then frames
  framesAt: number;
  nSamples: number;
  nFrames: number;
  samplesPerFrame: number;
  dt: number;
  stepsPerSample: number;
  omegaMax: number;
  nBead: number; // workgroups over beads
  nHinge: number; // workgroups over hinges
  hasMallet: boolean;
  freeRigid: boolean;
  velForces: boolean;
}

function f(x: number): string {
  if (!Number.isFinite(x)) throw new Error(`non-finite constant ${x}`);
  return x.toExponential(9);
}

class Pack<T extends Float32Array | Uint32Array> {
  parts: ArrayLike<number>[] = [];
  len = 0;
  constructor(private Ctor: new (n: number) => T) {}
  add(a: ArrayLike<number>): number {
    const off = this.len;
    this.parts.push(a);
    this.len += a.length;
    return off;
  }
  build(): T {
    const out = new this.Ctor(Math.max(4, this.len));
    let o = 0;
    for (const p of this.parts) {
      out.set(p, o);
      o += p.length;
    }
    return out;
  }
}

// rows of variable length -> slot-major table [slot][row] padded with `pad`
function padded<T>(rows: T[][], pad: T): { width: number; table: T[] } {
  const width = Math.max(1, ...rows.map((r) => r.length));
  const table: T[] = [];
  for (let q = 0; q < width; q++) for (const r of rows) table.push(q < r.length ? r[q] : pad);
  return { width, table };
}

function pack(sys: BeadSystem, exc: Excitation, opt: SimOptions, sg: boolean): Packed {
  const { n, si, sj, sk, sL0, sc, mass, pinned, anchors, rest } = sys;
  const mallet = exc.mallet;
  const { omegaMax, dt, stepsPerSample } = timeStep(sys, mallet, opt.fs, opt.cfl);
  const nSamples = Math.round(opt.duration * opt.fs);
  const samplesPerFrame = Math.max(1, Math.round(opt.fs / opt.frameRate));
  const nFrames = Math.floor((nSamples - 1) / samplesPerFrame) + 1;

  const prm = new Pack(Float32Array);
  const idx = new Pack(Uint32Array);

  // beads
  const invMass = Float64Array.from(mass, (m, i) => (pinned[i] ? 0 : 1 / m));
  let M = 0;
  const cm = [0, 0, 0];
  for (let i = 0; i < n; i++) {
    M += mass[i];
    for (let c = 0; c < 3; c++) cm[c] += mass[i] * rest[3 * i + c];
  }
  const r0 = new Float64Array(3 * n);
  for (let i = 0; i < n; i++) for (let c = 0; c < 3; c++) r0[3 * i + c] = rest[3 * i + c] - cm[c] / M;
  const i0 = new Uint32Array(n);
  const fr = new Float64Array(n);
  if (opt.pickupDelay)
    for (let i = 0; i < n; i++) {
      i0[i] = Math.floor(opt.pickupDelay[i]);
      fr[i] = opt.pickupDelay[i] - i0[i];
    }
  const H = Math.max(...i0) + 2;
  const pw = mallet ? mallet.weights : new Float64Array(n);
  const patch: number[] = [];
  for (let i = 0; i < n; i++) if (pw[i] > 0) patch.push(i);
  let mEff = 1;
  if (mallet) {
    let inv = 0;
    for (const i of patch) inv += (pw[i] * pw[i]) / mass[i];
    mEff = 1 / (1 / mallet.mass + inv);
  }

  const P_MASS = prm.add(mass);
  const P_IM = prm.add(invMass);
  const P_R0 = prm.add(r0);
  const P_W = prm.add(opt.pickup);
  const P_FR = prm.add(fr);
  const P_PW = prm.add(pw);
  const P_FS = prm.add(exc.force ? exc.force.shape : new Float64Array(3 * n));

  // springs, gathered per bead: [k, c, rest difference (3), rest length, e0, weight]
  // e0 = rest length - natural length (pre-tension); the weight counts a spring
  // between two moving beads half at each end. Padding: a zero spring to a fixed point.
  const FIXED = 0xffffffff;
  const SLIDE = 0x80000000;
  const code = (node: number) => (node < n ? node : anchors[node - n].slideBead < 0 ? FIXED : SLIDE + anchors[node - n].slideBead);
  const nodePos = (node: number, c: number) => {
    if (node < n) return rest[3 * node + c];
    const a = anchors[node - n];
    return a.slideBead < 0 || c === 0 ? a.pos[c] : rest[3 * a.slideBead + c];
  };
  type SpringEntry = { other: number; vals: number[] };
  const byBead: SpringEntry[][] = Array.from({ length: n }, () => []);
  for (let s = 0; s < si.length; s++) {
    for (const [me, ot] of [[si[s], sj[s]], [sj[s], si[s]]]) {
      if (me >= n || pinned[me]) continue;
      const rd = [0, 1, 2].map((c) => nodePos(ot, c) - nodePos(me, c));
      const lr = Math.hypot(rd[0], rd[1], rd[2]);
      const moving = ot < n && !pinned[ot];
      byBead[me].push({ other: code(ot), vals: [sk[s], sc[s], rd[0], rd[1], rd[2], lr, lr - sL0[s], moving ? 0.5 : 1] });
    }
  }
  const sp = padded(byBead, { other: FIXED, vals: [0, 0, 1, 0, 0, 1, 0, 0] });
  const P_SF = prm.add(sp.table.flatMap((e) => e.vals));
  const I_SO = idx.add(sp.table.map((e) => e.other));

  // bending: hinge nh is a dummy with zero angle, rate, moment and gradient
  const b = sys.bending;
  const nh = b ? b.w0.length : 0;
  let P_HF = 0, P_HC = 0, I_HN = 0, I_HCC = 0, I_BHE = 0;
  let MAXHC = 1, MAXBH = 1;
  if (b) {
    const hf: number[] = [];
    const hn: number[] = [];
    const d = (p: number, q: number) => [0, 1, 2].map((c) => rest[3 * p + c] - rest[3 * q + c]);
    for (let e = 0; e < nh; e++) {
      const [w0, w1, e0, e1] = [b.w0[e], b.w1[e], b.e0[e], b.e1[e]];
      if (Math.max(w0, w1, e0, e1) >= n) throw new Error('hinge on an anchor');
      hn.push(w0, w1, e0, e1);
      hf.push(b.sign[e], b.theta0[e], ...d(e1, e0), ...d(w0, e0), ...d(w1, e1));
    }
    // moment of hinge a = sum over its triangles of M[a][c] phi_c
    const rows: Map<number, number>[] = Array.from({ length: nh }, () => new Map());
    for (let t = 0; t < b.tri.length / 3; t++)
      for (let a = 0; a < 3; a++) {
        const ea = b.tri[3 * t + a];
        if (ea < 0) continue;
        for (let c = 0; c < 3; c++) {
          const ec = b.tri[3 * t + c];
          if (ec >= 0) rows[ea].set(ec, (rows[ea].get(ec) ?? 0) + b.M[9 * t + 3 * a + c]);
        }
      }
    const hc = padded(rows.map((r) => [...r]), [nh, 0]);
    MAXHC = hc.width;
    // per bead: the hinges it belongs to (hinge * 4 + slot)
    const bh: number[][] = Array.from({ length: n }, () => []);
    for (let e = 0; e < nh; e++) for (let k = 0; k < 4; k++) if (!pinned[hn[4 * e + k]]) bh[hn[4 * e + k]].push(4 * e + k);
    const bhp = padded(bh, 4 * nh);
    MAXBH = bhp.width;
    P_HF = prm.add(hf);
    P_HC = prm.add(hc.table.map(([, v]) => v));
    I_HN = idx.add(hn);
    I_HCC = idx.add(hc.table.map(([c]) => c));
    I_BHE = idx.add(bhp.table);
  }
  const I_PATCH = idx.add(patch);
  const I_I0 = idx.add(i0);

  // state
  const S_U = 0;
  const S_V = S_U + 3 * n;
  const S_A = S_V + 3 * n;
  const S_G = S_A + 3 * n;
  const S_PHI = S_G + 12 * (nh + 1);
  const S_RATE = S_PHI + nh + 1;
  const S_MOM = S_RATE + nh + 1;
  const S_MEL = S_MOM + nh + 1;
  const S_HIST = S_MEL + nh + 1;
  const state = new Float32Array(S_HIST + n * H);
  for (let i = 0; i < n; i++)
    if (!pinned[i])
      for (let c = 0; c < 3; c++) {
        state[S_U + 3 * i + c] = exc.u0[3 * i + c];
        state[S_V + 3 * i + c] = exc.v0[3 * i + c];
      }
  const nBead = Math.ceil(n / WGB);
  const nHinge = Math.ceil(nh / WGB);
  // read-write buffer: the state, then shared scalars, then per-workgroup partial sums
  const SC = state.length;
  const Q_BF = SC + 64; // 17 per bead workgroup
  const Q_HD = Q_BF + 17 * nBead; // 1 per hinge workgroup (hinge damping power)
  const Q_KC = Q_HD + Math.max(1, nHinge); // 6 per bead workgroup
  const Q_REC = Q_KC + 6 * nBead; // 3 per bead or hinge workgroup: audio, energy, largest strain
  const Q_IRV = Q_REC + 3 * (nBead + nHinge); // 15 per bead workgroup
  const rw = new Float32Array(Q_IRV + 15 * nBead);
  rw.set(state);
  if (mallet) rw.set(mallet.vel0, SC + 3);
  // read-only buffer: parameters (as f32 bits), then indices
  const prmArr = prm.build();
  const idxArr = idx.build();
  const ro = new Uint32Array(prmArr.length + idxArr.length);
  ro.set(new Uint32Array(prmArr.buffer));
  ro.set(idxArr, prmArr.length);
  const IO = prmArr.length;

  const O_AUD = 0;
  const O_CF = nSamples;
  const O_FSC = 2 * nSamples;
  const O_MP = O_FSC + 6 * nFrames;
  const OFR = O_MP + 3 * nFrames;

  const force = exc.force;
  const sigma1 = b?.sigma1 ?? 0;
  const anyDash = sc.some((c) => c !== 0);
  const velForces = sys.sigma0 !== 0 || anyDash || sigma1 !== 0 || (mallet?.mu ?? 0) > 0;
  const u = (x: number) => `${x}u`;
  const consts = `
const N = ${u(n)};
const NH = ${u(nh)};
const NP = ${u(patch.length)};
const NBW = ${u(nBead)};
const NHW = ${u(nHinge)};
const MAXS = ${u(sp.width)};
const MAXBH = ${u(MAXBH)};
const MAXHC = ${u(MAXHC)};
const H = ${u(H)};
const SPF = ${u(samplesPerFrame)};
const NF = ${u(nFrames)};
const DT = ${f(dt)};
const S0 = ${f(sys.sigma0)};
const S1 = ${f(sigma1)};
const MTOT = ${f(M)};
const HAS_MALLET = ${!!mallet};
const FREE_RIGID = ${!!sys.freeRigid};
const HAS_PULSE = ${!!force};
const ANY_DASH = ${anyDash};
const MM = ${f(mallet?.mass ?? 1)};
const MK = ${f(mallet?.K ?? 0)};
const MP = ${f(mallet?.p ?? 1)};
const MU = ${f(mallet?.mu ?? 0)};
const REACH2 = ${f(mallet?.reach !== undefined ? mallet.reach * mallet.reach : -1)};
const MEFF = ${f(mEff)};
const NRM = vec3f(${(mallet?.normal ?? [0, 0, 0]).map(f).join(', ')});
const F0 = ${f(force?.F0 ?? 0)};
const TAU = ${f(force?.tau ?? 1)};
const P_MASS = ${u(P_MASS)}; const P_IM = ${u(P_IM)}; const P_R0 = ${u(P_R0)}; const P_W = ${u(P_W)};
const P_FR = ${u(P_FR)}; const P_PW = ${u(P_PW)}; const P_FS = ${u(P_FS)}; const P_SF = ${u(P_SF)};
const P_HF = ${u(P_HF)}; const P_HC = ${u(P_HC)};
const I_SO = ${u(I_SO + IO)}; const I_HN = ${u(I_HN + IO)}; const I_HCC = ${u(I_HCC + IO)}; const I_BHE = ${u(I_BHE + IO)};
const I_PATCH = ${u(I_PATCH + IO)}; const I_I0 = ${u(I_I0 + IO)};
const S_U = ${u(S_U)}; const S_V = ${u(S_V)}; const S_A = ${u(S_A)}; const S_G = ${u(S_G)};
const S_PHI = ${u(S_PHI)}; const S_RATE = ${u(S_RATE)}; const S_MOM = ${u(S_MOM)}; const S_MEL = ${u(S_MEL)};
const S_HIST = ${u(S_HIST)}; const SC = ${u(SC)};
const Q_BF = ${u(Q_BF)}; const Q_HD = ${u(Q_HD)}; const Q_KC = ${u(Q_KC)}; const Q_REC = ${u(Q_REC)}; const Q_IRV = ${u(Q_IRV)};
const O_AUD = ${u(O_AUD)}; const O_CF = ${u(O_CF)}; const O_FSC = ${u(O_FSC)}; const O_MP = ${u(O_MP)}; const OFR = ${u(OFR)};
`;
  return {
    wgsl: (sg ? 'enable subgroups;\n' : '') + consts + reducers(sg) + KERNELS.replaceAll('/*SG*/', sg ? ', @builtin(subgroup_invocation_id) sgiIn: u32, @builtin(subgroup_size) sgsIn: u32' : '').replaceAll('/*SGI*/', sg ? 'sgi = sgiIn; sgs = sgsIn;' : ''),
    ro,
    rw,
    outLen: OFR + 3 * n * nFrames,
    framesAt: OFR,
    nSamples,
    nFrames,
    samplesPerFrame,
    dt,
    stepsPerSample,
    omegaMax,
    nBead,
    nHinge,
    hasMallet: !!mallet,
    freeRigid: !!sys.freeRigid,
    velForces,
  };
}

// Workgroup sums: each invocation puts its values, then wsum(K) leaves the
// sums in wres[0..K] (all invocations must call it). wmax: the largest value
// put at 0. With subgroups: a subgroup sum per value, then one invocation per
// value adds the subgroups' sums (subgroups taken to be consecutive
// invocations); else a tree in workgroup memory.
function reducers(sg: boolean): string {
  const common = `
var<workgroup> wred: array<f32, 17u * 64u>;
var<workgroup> wres: array<f32, 17>;
var<private> sgi: u32;
var<private> sgs: u32;
fn put(l: u32, k: u32, v: f32) { wred[k * 64u + l] = v; }
fn put3(l: u32, k: u32, v: vec3f) { put(l, k, v.x); put(l, k + 1u, v.y); put(l, k + 2u, v.z); }
`;
  if (sg)
    return common + `
var<workgroup> wsg: array<f32, 17u * 16u>;
fn wsum(l: u32, K: u32) {
  workgroupBarrier();
  let g = l / sgs;
  for (var k = 0u; k < K; k++) {
    let s = subgroupAdd(wred[k * 64u + l]);
    if (sgi == 0u) { wsg[k * 16u + g] = s; }
  }
  workgroupBarrier();
  if (l < K) {
    var t = 0.0;
    for (var j = 0u; j < 64u / sgs; j++) { t += wsg[l * 16u + j]; }
    wres[l] = t;
  }
  workgroupBarrier();
}
fn wmax(l: u32) -> f32 {
  workgroupBarrier();
  let s = subgroupMax(wred[l]);
  if (sgi == 0u) { wsg[l / sgs] = s; }
  workgroupBarrier();
  var m = 0.0;
  for (var j = 0u; j < 64u / sgs; j++) { m = max(m, wsg[j]); }
  workgroupBarrier();
  return m;
}
`;
  return common + `
fn wsum(l: u32, K: u32) {
  workgroupBarrier();
  // each level adds pairs (k, j) and (k, j + s), spread over all invocations
  for (var s = 32u; s > 0u; s = s / 2u) {
    for (var x = l; x < K * s; x += 64u) {
      let k = x / s;
      let j = x - k * s;
      wred[k * 64u + j] += wred[k * 64u + j + s];
    }
    workgroupBarrier();
  }
  if (l < K) { wres[l] = wred[l * 64u]; }
  workgroupBarrier();
}
fn wmax(l: u32) -> f32 {
  workgroupBarrier();
  for (var s = 32u; s > 0u; s = s / 2u) {
    if (l < s) { wred[l] = max(wred[l], wred[l + s]); }
    workgroupBarrier();
  }
  let m = wred[0];
  workgroupBarrier();
  return m;
}
`;
}

// Shared scalars (at SC), written by invocation 0 of workgroup 0:
//  0-2 mallet position q, 3-5 velocity, 6-8 acceleration, 9-11 force on the
//  mallet, 12 contact force, 13 contact energy, 14 friction power,
//  17-18 dissipated since the last frame (Kahan sum), 19-20 work in since the
//  last frame, 44 time steps taken and 45 next sample to record (u32 bits).
// Each workgroup keeps its own copy of the rigid-motion correction (wcor).
const KERNELS = /* wgsl */ `
@group(0) @binding(0) var<storage, read> ro: array<u32>;
@group(0) @binding(1) var<storage, read_write> rw: array<f32>;
@group(0) @binding(2) var<storage, read_write> ob: array<f32>;

fn P(k: u32) -> f32 { return bitcast<f32>(ro[k]); }
fn pv(o: u32) -> vec3f { return vec3f(P(o), P(o + 1u), P(o + 2u)); }
fn ld(o: u32) -> vec3f { return vec3f(rw[o], rw[o + 1u], rw[o + 2u]); }
fn sv(o: u32, x: vec3f) { rw[o] = x.x; rw[o + 1u] = x.y; rw[o + 2u] = x.z; }
fn s3(k: u32) -> vec3f { return ld(SC + k); }
fn ss3(k: u32, x: vec3f) { sv(SC + k, x); }
fn count(k: u32) -> u32 { return bitcast<u32>(rw[SC + k]); }
fn gU(i: u32) -> vec3f { return ld(S_U + 3u * i); }
fn gV(i: u32) -> vec3f { return ld(S_V + 3u * i); }
fn gA(i: u32) -> vec3f { return ld(S_A + 3u * i); }
fn sU(i: u32, x: vec3f) { sv(S_U + 3u * i, x); }
fn sV(i: u32, x: vec3f) { sv(S_V + 3u * i, x); }
fn sA(i: u32, x: vec3f) { sv(S_A + 3u * i, x); }
// position of bead i relative to the rest centre of mass
fn pos(i: u32) -> vec3f { return pv(P_R0 + 3u * i) + gU(i); }

// displacement or velocity of a spring's far node: a bead, a fixed anchor (0)
// or a sliding anchor that follows bead b across its rod
fn otherU(o: u32) -> vec3f {
  if (o == 0xffffffffu) { return vec3f(0.0); }
  if (o >= 0x80000000u) { let x = gU(o - 0x80000000u); return vec3f(0.0, x.y, x.z); }
  return gU(o);
}
fn otherV(o: u32) -> vec3f {
  if (o == 0xffffffffu) { return vec3f(0.0); }
  if (o >= 0x80000000u) { let x = gV(o - 0x80000000u); return vec3f(0.0, x.y, x.z); }
  return gV(o);
}

fn kahan(k: u32, x: f32) {
  let o = SC + k;
  let y = x - rw[o + 1u];
  let t = rw[o] + y;
  rw[o + 1u] = (t - rw[o]) - y;
  rw[o] = t;
}

// sums of per-workgroup partials rw[q + w * stride + k] over w < nw into wres[0..K]
// (the loads are spread over the invocations: global memory latency, not
// arithmetic, is what costs here)
fn psum(l: u32, q: u32, K: u32, stride: u32, nw: u32) {
  for (var k = 0u; k < K; k++) { put(l, k, 0.0); }
  for (var x = l; x < nw * K; x += 64u) {
    let w = x / K;
    let k = x - w * K;
    wred[k * 64u + l] += rw[q + w * stride + k];
  }
  wsum(l, K);
}

// sums for the rigid part of a per-bead field x: mass-weighted geometry
// (m u, m s s^T as diagonal and off-diagonal) and momentum (m x, m s x x),
// s = position relative to the rest centre of mass
struct Rig { mu: vec3f, sd: vec3f, so: vec3f, p: vec3f, l: vec3f };
fn rigAdd(r: Rig, m: f32, u: vec3f, s: vec3f, x: vec3f) -> Rig {
  return Rig(r.mu + m * u, r.sd + m * s * s, r.so + m * vec3f(s.x * s.y, s.x * s.z, s.y * s.z), r.p + m * x, r.l + m * cross(s, x));
}
fn rigZero() -> Rig { return Rig(vec3f(0.0), vec3f(0.0), vec3f(0.0), vec3f(0.0), vec3f(0.0)); }
fn rigPut(l: u32, r: Rig) { put3(l, 0u, r.mu); put3(l, 3u, r.sd); put3(l, 6u, r.so); put3(l, 9u, r.p); put3(l, 12u, r.l); }

// This workgroup's rigid-motion correction: 0-2 centre of mass, 3-8 inverse
// inertia, 9-11 rigid translation rate, 12-14 rotation rate, 15 kinetic
// energy of the rigid part, 16 damping power, 17 external power (from the
// bead-force sums).
var<workgroup> wcor: array<f32, 18>;
fn w3(k: u32) -> vec3f { return vec3f(wcor[k], wcor[k + 1u], wcor[k + 2u]); }
fn sw3(k: u32, x: vec3f) { wcor[k] = x.x; wcor[k + 1u] = x.y; wcor[k + 2u] = x.z; }
// invocation 0: centre of mass and inverse inertia from geometry sums in wres[0..9]
fn rigidGeometry() {
  let c = vec3f(wres[0], wres[1], wres[2]) / MTOT;
  let sxx = wres[3] - MTOT * c.x * c.x;
  let syy = wres[4] - MTOT * c.y * c.y;
  let szz = wres[5] - MTOT * c.z * c.z;
  let sxy = wres[6] - MTOT * c.x * c.y;
  let sxz = wres[7] - MTOT * c.x * c.z;
  let syz = wres[8] - MTOT * c.y * c.z;
  let a = syy + szz; let b = -sxy; let cc = -sxz;
  let d = sxx + szz; let e = -syz; let ff = sxx + syy;
  let A00 = d * ff - e * e; let A01 = cc * e - b * ff; let A02 = b * e - cc * d;
  let A11 = a * ff - cc * cc; let A12 = b * cc - a * e; let A22 = a * d - b * b;
  let det = a * A00 + b * A01 + cc * A02;
  let k = select(0.0, 1.0 / det, abs(det) > 0.0);
  wcor[3] = A00 * k; wcor[4] = A01 * k; wcor[5] = A02 * k;
  wcor[6] = A11 * k; wcor[7] = A12 * k; wcor[8] = A22 * k;
  sw3(0u, c);
}
// invocation 0: rigid translation and rotation rate of a field from its
// momentum sums in wres[o..o+6], and its kinetic energy
fn rigidPart(o: u32) {
  let P = vec3f(wres[o], wres[o + 1u], wres[o + 2u]);
  let L = vec3f(wres[o + 3u], wres[o + 4u], wres[o + 5u]) - cross(w3(0u), P);
  let w = vec3f(wcor[3] * L.x + wcor[4] * L.y + wcor[5] * L.z, wcor[4] * L.x + wcor[6] * L.y + wcor[7] * L.z, wcor[5] * L.x + wcor[7] * L.y + wcor[8] * L.z);
  let vc = P / MTOT;
  sw3(9u, vc);
  sw3(12u, w);
  wcor[15] = 0.5 * dot(P, vc) + 0.5 * dot(w, L);
}
fn rigidAt(s: vec3f) -> vec3f { return w3(9u) + cross(w3(12u), s - w3(0u)); }

// all invocations: the geometry and the rigid part of the accelerations from
// the bead-force partials, and the damping and external power
fn accelRigid(l: u32) {
  psum(l, Q_BF, 17u, 17u, NBW);
  if (l == 0u) {
    rigidGeometry();
    rigidPart(9u);
    wcor[16] = wres[15];
    wcor[17] = wres[16];
  }
  workgroupBarrier();
}
// all invocations: the geometry from the bead-force partials and the rigid
// part of the velocities from the kick partials
fn velRigid(l: u32) {
  psum(l, Q_BF, 9u, 17u, NBW);
  if (l == 0u) { rigidGeometry(); }
  workgroupBarrier();
  psum(l, Q_KC, 6u, 6u, NBW);
  if (l == 0u) { rigidPart(0u); }
  workgroupBarrier();
}

// dihedral angle change and its gradient (Bridson et al. 2003; see hinge.ts)
fn hinge(e: u32) {
  let w0 = ro[I_HN + 4u * e];
  let w1 = ro[I_HN + 4u * e + 1u];
  let e0 = ro[I_HN + 4u * e + 2u];
  let e1 = ro[I_HN + 4u * e + 3u];
  let hb = P_HF + 11u * e;
  let sg = P(hb);
  let ue0 = gU(e0);
  let ue1 = gU(e1);
  let E = pv(hb + 2u) + (ue1 - ue0);
  let P1 = pv(hb + 5u) + (gU(w0) - ue0);
  let P2 = pv(hb + 8u) + (gU(w1) - ue1);
  let Q1 = P1 - E;
  let Q2 = P2 + E;
  let N1 = cross(P1, Q1);
  let N2 = cross(P2, Q2);
  let n1sq = dot(N1, N1);
  let n2sq = dot(N2, N2);
  let eLen = length(E);
  let nn = sqrt(n1sq) * sqrt(n2sq);
  let theta = atan2(dot(cross(N1, N2), E) / (nn * eLen), dot(N1, N2) / nn);
  let gw0 = (-eLen / n1sq) * N1;
  let gw1 = (-eLen / n2sq) * N2;
  let t1 = dot(Q1, E) / eLen / n1sq;
  let t2 = dot(P2, E) / eLen / n2sq;
  let r1 = dot(P1, E) / eLen / n1sq;
  let r2 = dot(Q2, E) / eLen / n2sq;
  let ge0 = -(t1 * N1 + t2 * N2);
  let ge1 = r1 * N1 + r2 * N2;
  let g = S_G + 12u * e;
  sv(g, sg * gw0);
  sv(g + 3u, sg * gw1);
  sv(g + 6u, sg * ge0);
  sv(g + 9u, sg * ge1);
  rw[S_PHI + e] = sg * (theta - P(hb + 1u));
  if (S1 != 0.0) {
    rw[S_RATE + e] = sg * (dot(gw0, gV(w0)) + dot(gw1, gV(w1)) + dot(ge0, gV(e0)) + dot(ge1, gV(e1)));
  }
}

// mallet contact force from the patch-weighted displacement U and velocity Vp
fn contact(U: vec3f, Vp: vec3f) {
  let q = s3(0u);
  let qv = s3(3u);
  var fm = vec3f(0.0);
  var fc = 0.0;
  var ec = 0.0;
  var pf = 0.0;
  let d = dot(q - U, NRM);
  if (d > 0.0) {
    let o = q - U - d * NRM;
    if (REACH2 < 0.0 || dot(o, o) <= REACH2) {
      let fnorm = MK * pow(d, MP);
      fc = fnorm;
      ec = fnorm * d / (MP + 1.0);
      fm = -fnorm * NRM;
      if (MU > 0.0) {
        var r = qv - Vp;
        r -= dot(r, NRM) * NRM;
        let slip = length(r);
        let eps = max(1e-3, 4.0 * MU * fnorm * DT / MEFF);
        let g = MU * fnorm / sqrt(slip * slip + eps * eps);
        fm -= g * r;
        pf = g * slip * slip;
      }
    }
  }
  ss3(6u, fm / MM);
  ss3(9u, fm);
  rw[SC + 12u] = fc;
  rw[SC + 13u] = ec;
  rw[SC + 14u] = pf;
}

// invocation 0 of workgroup 0 (after the mallet's position and velocity
// are final for the step): the mallet's half kick and drift, count the step
fn malletDrift() {
  if (HAS_MALLET) {
    let qv = s3(3u) + 0.5 * DT * s3(6u);
    ss3(3u, qv);
    ss3(0u, s3(0u) + DT * qv);
  }
  rw[SC + 44u] = bitcast<f32>(count(44u) + 1u);
}

// workgroup 0: sum the record partials and write the sample (and frame),
// then move to the next sample
fn recordSums(l: u32) {
  var mx = 0.0;
  for (var w = l; w < NBW + NHW; w += 64u) { mx = max(mx, rw[Q_REC + 3u * w + 2u]); }
  put(l, 0u, mx);
  let largest = wmax(l);
  psum(l, Q_REC, 3u, 3u, NBW + NHW);
  if (l == 0u) {
    let smp = count(45u);
    ob[O_AUD + smp] = wres[0];
    ob[O_CF + smp] = rw[SC + 12u];
    if (smp % SPF == 0u && smp / SPF < NF) {
      let fi = smp / SPF;
      let b = O_FSC + 6u * fi;
      let qv = s3(3u);
      ob[b] = wres[1];
      ob[b + 1u] = select(0.0, 0.5 * MM * dot(qv, qv) + rw[SC + 13u], HAS_MALLET);
      ob[b + 2u] = rw[SC + 17u];
      ob[b + 3u] = rw[SC + 19u];
      ob[b + 4u] = largest;
      rw[SC + 17u] = 0.0; rw[SC + 18u] = 0.0; rw[SC + 19u] = 0.0; rw[SC + 20u] = 0.0;
      let q = s3(0u);
      ob[O_MP + 3u * fi] = q.x;
      ob[O_MP + 3u * fi + 1u] = q.y;
      ob[O_MP + 3u * fi + 2u] = q.z;
    }
    rw[SC + 45u] = bitcast<f32>(smp + 1u);
  }
}

// first kick and drift of a sample's steps; workgroup 0 also writes the
// sample recorded just before
@compute @workgroup_size(64)
fn kickDriftRec(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  if (wg.x == 0u) {
    recordSums(l);
    if (l == 0u) { malletDrift(); }
  }
  let i = wg.x * 64u + l;
  if (i >= N || P(P_IM + i) == 0.0) { return; }
  let v = gV(i) + 0.5 * DT * gA(i);
  sV(i, v);
  sU(i, gU(i) + DT * v);
}

// later kicks and drifts: first remove the rigid part of the velocities left
// by the previous step
@compute @workgroup_size(64)
fn kickDrift(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  if (FREE_RIGID) { velRigid(l); }
  if (wg.x == 0u && l == 0u) {
    if (FREE_RIGID) { kahan(17u, wcor[15]); }
    malletDrift();
  }
  let i = wg.x * 64u + l;
  if (i >= N || P(P_IM + i) == 0.0) { return; }
  var v = gV(i);
  if (FREE_RIGID) { v -= rigidAt(pos(i)); }
  v += 0.5 * DT * gA(i);
  sV(i, v);
  sU(i, gU(i) + DT * v);
}

// after the last step of a sample: remove the rigid part of the velocities
@compute @workgroup_size(64)
fn velCorrect(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  velRigid(l);
  if (wg.x == 0u && l == 0u) { kahan(17u, wcor[15]); }
  let i = wg.x * 64u + l;
  if (i >= N) { return; }
  sV(i, gV(i) - rigidAt(pos(i)));
}

// hinge angles, gradients and rates; the last workgroup does the mallet contact
@compute @workgroup_size(64)
fn hinges(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  if (wg.x < NHW) {
    let e = wg.x * 64u + l;
    if (e < NH) { hinge(e); }
    return;
  }
  var su = vec3f(0.0);
  var sv = vec3f(0.0);
  for (var j = l; j < NP; j += 64u) {
    let i = ro[I_PATCH + j];
    let w = P(P_PW + i);
    su += w * gU(i);
    sv += w * gV(i);
  }
  put3(l, 0u, su);
  put3(l, 3u, sv);
  wsum(l, 6u);
  if (l == 0u) { contact(vec3f(wres[0], wres[1], wres[2]), vec3f(wres[3], wres[4], wres[5])); }
}

// hinge moments, elastic and with damping; per-workgroup hinge damping power
@compute @workgroup_size(64)
fn moments(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  let e = wg.x * 64u + l;
  var pd = 0.0;
  if (e < NH) {
    var m = 0.0;
    var r = 0.0;
    for (var q = 0u; q < MAXHC; q++) {
      let j = q * NH + e;
      let c = ro[I_HCC + j];
      let k = P(P_HC + j);
      m += k * rw[S_PHI + c];
      if (S1 != 0.0) { r += k * rw[S_RATE + c]; }
    }
    rw[S_MEL + e] = m;
    rw[S_MOM + e] = m + 2.0 * S1 * r;
    if (S1 != 0.0) { pd = 2.0 * S1 * rw[S_RATE + e] * r; }
  }
  put(l, 0u, pd);
  wsum(l, 1u);
  if (l == 0u) { rw[Q_HD + wg.x] = wres[0]; }
}

// accelerations (before rigid removal); per-workgroup damping and external
// power and sums for the rigid part of the accelerations
@compute @workgroup_size(64)
fn beadForces(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  let i = wg.x * 64u + l;
  let t = f32(count(44u)) * DT;
  var pd = 0.0;
  var pe = 0.0;
  var rig = rigZero();
  if (i < N) {
    let im = P(P_IM + i);
    let ui = gU(i);
    var ai = vec3f(0.0);
    if (im != 0.0) {
      let vi = gV(i);
      var fi = vec3f(0.0);
      for (var q = 0u; q < MAXS; q++) {
        let j = q * N + i;
        let o = ro[I_SO + j];
        let b = P_SF + 8u * j;
        let rd = pv(b + 2u);
        let du = otherU(o) - ui;
        let d = rd + du;
        let len = length(d);
        let dl = (2.0 * dot(rd, du) + dot(du, du)) / (len + P(b + 5u));
        fi += (P(b) * (dl + P(b + 6u)) / len) * d;
        if (ANY_DASH) {
          let c = P(b + 1u);
          let vr = otherV(o) - vi;
          fi += c * vr;
          pd += P(b + 7u) * c * dot(vr, vr);
        }
      }
      if (NH > 0u) {
        for (var q = 0u; q < MAXBH; q++) {
          let ent = ro[I_BHE + q * N + i];
          let e = ent >> 2u;
          fi -= rw[S_MOM + e] * ld(S_G + 12u * e + 3u * (ent & 3u));
        }
      }
      if (HAS_MALLET) { fi -= P(P_PW + i) * s3(9u); }
      if (S0 != 0.0) {
        let m2 = 2.0 * S0 * P(P_MASS + i);
        fi -= m2 * vi;
        pd += m2 * dot(vi, vi);
      }
      if (HAS_PULSE && t >= 0.0 && t <= TAU) {
        let fe = F0 * 0.5 * (1.0 - cos(6.283185307 * t / TAU)) * pv(P_FS + 3u * i);
        fi += fe;
        pe += dot(fe, vi);
      }
      ai = fi * im;
    }
    sA(i, ai);
    if (FREE_RIGID) { rig = rigAdd(rig, P(P_MASS + i), ui, pv(P_R0 + 3u * i) + ui, ai); }
  }
  rigPut(l, rig);
  put(l, 15u, pd);
  put(l, 16u, pe);
  wsum(l, 17u);
  if (l < 17u) { rw[Q_BF + wg.x * 17u + l] = wres[l]; }
}

// remove the rigid part of the accelerations, second half kick; per-workgroup
// sums for the rigid part of the velocities. Workgroup 0: the mallet's second
// half kick, dissipated energy and work in.
@compute @workgroup_size(64)
fn kick(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  if (FREE_RIGID || wg.x == 0u) { accelRigid(l); }
  if (wg.x == 0u && l == 0u) {
    var hd = 0.0;
    for (var w = 0u; w < NHW; w++) { hd += rw[Q_HD + w]; }
    if (HAS_MALLET) { ss3(3u, s3(3u) + 0.5 * DT * s3(6u)); }
    kahan(17u, DT * (wcor[16] + hd + rw[SC + 14u]));
    kahan(19u, DT * wcor[17]);
  }
  let i = wg.x * 64u + l;
  var p = vec3f(0.0);
  var lm = vec3f(0.0);
  if (i < N) {
    var a = gA(i);
    var v = gV(i);
    if (FREE_RIGID) {
      let s = pos(i);
      a -= rigidAt(s);
      sA(i, a);
      if (P(P_IM + i) != 0.0) { v += 0.5 * DT * a; }
      let m = P(P_MASS + i);
      p = m * v;
      lm = m * cross(s, v);
    } else if (P(P_IM + i) != 0.0) {
      v += 0.5 * DT * a;
    }
    sV(i, v);
  }
  if (FREE_RIGID) {
    put3(l, 0u, p);
    put3(l, 3u, lm);
    wsum(l, 6u);
    if (l < 6u) { rw[Q_KC + wg.x * 6u + l] = wres[l]; }
  }
}

// per-workgroup audio and (at frames) energy and largest strain for the next
// sample: bead workgroups first, then hinge workgroups (hinge energy). With
// correct = true the accelerations still hold their rigid part (forces were
// just recomputed) and it is removed first.
fn record(wg: u32, l: u32, correct: bool) {
  let smp = count(45u);
  let isFrame = smp % SPF == 0u && smp / SPF < NF;
  var ap = 0.0;
  var ep = 0.0;
  var mx = 0.0;
  if (wg < NBW) {
    if (correct) { accelRigid(l); }
    let i = wg * 64u + l;
    if (i < N) {
      var a = gA(i);
      if (correct) {
        a -= rigidAt(pos(i));
        sA(i, a);
      }
      let hb = S_HIST + i * H;
      rw[hb + smp % H] = dot(pv(P_W + 3u * i), a);
      let i0 = ro[I_I0 + i];
      let fr = P(P_FR + i);
      if (smp >= i0) { ap += (1.0 - fr) * rw[hb + (smp - i0) % H]; }
      if (smp >= i0 + 1u) { ap += fr * rw[hb + (smp - i0 - 1u) % H]; }
      if (isFrame) {
        let ui = gU(i);
        let fo = OFR + (smp / SPF) * 3u * N + 3u * i;
        ob[fo] = ui.x;
        ob[fo + 1u] = ui.y;
        ob[fo + 2u] = ui.z;
        let vi = gV(i);
        ep += 0.5 * P(P_MASS + i) * dot(vi, vi);
        if (P(P_IM + i) != 0.0) {
          for (var q = 0u; q < MAXS; q++) {
            let j = q * N + i;
            let o = ro[I_SO + j];
            let b = P_SF + 8u * j;
            let rd = pv(b + 2u);
            let lr = P(b + 5u);
            let e0 = P(b + 6u);
            let du = otherU(o) - ui;
            let len = length(rd + du);
            let dl = (2.0 * dot(rd, du) + dot(du, du)) / (len + lr);
            // energy above the pre-tensioned equilibrium: k/2 ((dl + e0)^2 - e0^2)
            ep += P(b + 7u) * 0.5 * P(b) * dl * (dl + 2.0 * e0);
            if (o < N) { mx = max(mx, abs(dl + e0) / (lr - e0)); }
          }
        }
      }
    }
  } else if (isFrame) {
    let e = (wg - NBW) * 64u + l;
    if (e < NH) { ep = 0.5 * rw[S_PHI + e] * rw[S_MEL + e]; }
  }
  put(l, 0u, mx);
  let largest = wmax(l);
  put(l, 0u, ap);
  put(l, 1u, ep);
  wsum(l, 2u);
  if (l == 0u) {
    rw[Q_REC + wg * 3u] = wres[0];
    rw[Q_REC + wg * 3u + 1u] = wres[1];
    rw[Q_REC + wg * 3u + 2u] = largest;
  }
}
@compute @workgroup_size(64)
fn recordSample(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/ record(wg.x, l, false); }
@compute @workgroup_size(64)
fn recordSampleCorrect(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/ record(wg.x, l, FREE_RIGID); }
// writes the last sample
@compute @workgroup_size(64)
fn recordLast(@builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/ recordSums(l); }

// at the start: per-workgroup sums for the rigid part of the initial velocities
@compute @workgroup_size(64)
fn initRigid(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  let i = wg.x * 64u + l;
  var rig = rigZero();
  if (i < N) {
    let u = gU(i);
    rig = rigAdd(rig, P(P_MASS + i), u, pv(P_R0 + 3u * i) + u, gV(i));
  }
  rigPut(l, rig);
  wsum(l, 15u);
  if (l < 15u) { rw[Q_IRV + wg.x * 15u + l] = wres[l]; }
}
// and its removal
@compute @workgroup_size(64)
fn initVel(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
  /*SGI*/
  psum(l, Q_IRV, 15u, 15u, NBW);
  if (l == 0u) {
    rigidGeometry();
    rigidPart(9u);
  }
  workgroupBarrier();
  let i = wg.x * 64u + l;
  if (i >= N) { return; }
  sV(i, gV(i) - rigidAt(pos(i)));
}
`;

const ENTRY = ['kickDriftRec', 'kickDrift', 'velCorrect', 'hinges', 'moments', 'beadForces', 'kick', 'recordSample', 'recordSampleCorrect', 'recordLast', 'initRigid', 'initVel'] as const;
type Entry = (typeof ENTRY)[number];

export async function simulateGPU(
  sys: BeadSystem,
  exc: Excitation,
  opt: SimOptions,
  progress?: Progress,
): Promise<SimResult & { device: string; loopSeconds: number }> {
  const t0 = performance.now();
  const gpu = navigator.gpu;
  if (!gpu) throw new Error('WebGPU is not available in this browser');
  const adapter = await gpu.requestAdapter({ powerPreference: 'high-performance' });
  if (!adapter) throw new Error('no WebGPU adapter');
  const lim = adapter.limits;
  const sg = adapter.features.has('subgroups');
  const pk = pack(sys, exc, opt, sg);
  const biggest = Math.max(pk.ro.byteLength, pk.rw.byteLength, pk.outLen * 4);
  if (biggest > lim.maxStorageBufferBindingSize || biggest > lim.maxBufferSize)
    throw new Error(`a GPU buffer would need ${(biggest / 1e6).toFixed(0)} MB (limit ${(Math.min(lim.maxStorageBufferBindingSize, lim.maxBufferSize) / 1e6).toFixed(0)} MB); lower the replay frame rate or the duration`);
  const device = await adapter.requestDevice({
    requiredFeatures: sg ? ['subgroups' as GPUFeatureName] : [],
    requiredLimits: {
      maxStorageBufferBindingSize: Math.max(biggest, 1 << 20),
      maxBufferSize: Math.max(biggest, 1 << 20),
    },
  });
  let lost = '';
  device.lost.then((info) => (lost = info.message || 'GPU device lost'));
  try {
    device.pushErrorScope('validation');
    const module = device.createShaderModule({ code: pk.wgsl });
    const info = await module.getCompilationInfo();
    const errs = info.messages.filter((m) => m.type === 'error');
    if (errs.length) throw new Error(`shader: ${errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join('; ')}`);
    const bgl = device.createBindGroupLayout({
      entries: (['read-only-storage', 'storage', 'storage'] as const).map((type, binding) => ({ binding, visibility: GPUShaderStage.COMPUTE, buffer: { type } })),
    });
    const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
    const pipes = {} as Record<Entry, GPUComputePipeline>;
    await Promise.all(ENTRY.map(async (e) => (pipes[e] = await device.createComputePipelineAsync({ layout, compute: { module, entryPoint: e } }))));
    const mk = (data: Float32Array | Uint32Array | null, size: number) => {
      const buf = device.createBuffer({ size: Math.max(16, size), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC });
      if (data) device.queue.writeBuffer(buf, 0, data as Float32Array<ArrayBuffer>);
      return buf;
    };
    const outB = mk(null, pk.outLen * 4);
    const bind = device.createBindGroup({
      layout: bgl,
      entries: [mk(pk.ro, pk.ro.byteLength), mk(pk.rw, pk.rw.byteLength), outB].map((buffer, binding) => ({ binding, resource: { buffer } })),
    });

    const { nBead, nHinge, hasMallet, freeRigid, velForces, stepsPerSample } = pk;
    const nh = sys.bending ? sys.bending.w0.length : 0;
    type D = [Entry, number];
    const forces: D[] = [];
    if (nHinge + (hasMallet ? 1 : 0) > 0) forces.push(['hinges', nHinge + (hasMallet ? 1 : 0)]);
    if (nh > 0) forces.push(['moments', nHinge]);
    forces.push(['beadForces', nBead]);
    // one sample: record it, then the steps to the next one; the accelerations
    // hold their rigid part after a force evaluation until a kick removes it
    const sample = (first: boolean, last: boolean): D[] => {
      const raw = first || velForces;
      const list: D[] = [[raw ? 'recordSampleCorrect' : 'recordSample', nBead + nHinge]];
      if (last) return [...list, ['recordLast', 1]];
      for (let j = 0; j < stepsPerSample; j++) list.push([j === 0 ? 'kickDriftRec' : 'kickDrift', nBead], ...forces, ['kick', nBead]);
      if (freeRigid) list.push(['velCorrect', nBead]);
      if (velForces) list.push(...forces);
      return list;
    };
    const start: D[] = [...(freeRigid ? ([['initRigid', nBead], ['initVel', nBead]] as D[]) : []), ...forces];
    const encode = (pass: GPUComputePassEncoder, list: D[]) => {
      for (const [e, w] of list) {
        pass.setPipeline(pipes[e]);
        pass.dispatchWorkgroups(w);
      }
    };
    const submit = (fill: (pass: GPUComputePassEncoder) => void) => {
      const enc = device.createCommandEncoder();
      const pass = enc.beginComputePass();
      pass.setBindGroup(0, bind);
      fill(pass);
      pass.end();
      device.queue.submit([enc.finish()]);
      return device.queue.onSubmittedWorkDone();
    };
    await submit((pass) => encode(pass, start));
    const err = await device.popErrorScope();
    if (err) throw new Error(err.message);

    // chunks of samples, about 50 ms of GPU work each, two in flight so the
    // GPU does not wait while the next one is encoded
    const tLoop = performance.now();
    const mid = sample(false, false);
    let chunk = 16;
    let pending: Promise<void> | null = null;
    let pendingEnd = 0;
    let tPrev = performance.now();
    for (let s = 0; s < pk.nSamples; ) {
      const e = Math.min(pk.nSamples, s + chunk);
      const done = submit((pass) => {
        for (let k = s; k < e; k++) encode(pass, k === 0 || k === pk.nSamples - 1 ? sample(k === 0, k === pk.nSamples - 1) : mid);
      });
      if (pending) {
        await pending;
        if (lost) throw new Error(lost);
        const now = performance.now();
        chunk = Math.max(1, Math.min(Math.round(chunk * Math.min(2, 50 / Math.max(now - tPrev, 1))), 4096));
        tPrev = now;
        progress?.(pendingEnd / pk.nSamples);
      }
      pending = done;
      pendingEnd = e;
      s = e;
    }
    await pending;
    if (lost) throw new Error(lost);
    progress?.(1);
    const loopSeconds = (performance.now() - tLoop) / 1000;

    const stage = device.createBuffer({ size: Math.max(16, pk.outLen * 4), usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    const enc = device.createCommandEncoder();
    enc.copyBufferToBuffer(outB, 0, stage, 0, Math.max(16, pk.outLen * 4));
    device.queue.submit([enc.finish()]);
    await stage.mapAsync(GPUMapMode.READ);
    const out = new Float32Array(stage.getMappedRange().slice(0, pk.outLen * 4));
    stage.destroy();
    if (lost) throw new Error(lost);

    const { nSamples, nFrames, samplesPerFrame, dt } = pk;
    const audio = out.slice(0, nSamples);
    for (let k = 0; k < audio.length; k++) if (!Number.isFinite(audio[k])) throw new Error('the GPU run blew up (non-finite output)');
    const contactForce = exc.mallet ? out.slice(nSamples, 2 * nSamples) : undefined;
    const O_FSC = 2 * nSamples;
    const O_MP = O_FSC + 6 * nFrames;
    const frameTimes = new Float64Array(nFrames);
    const energy = new Float64Array(nFrames);
    const balance = new Float64Array(nFrames);
    let maxStretch = 0;
    let diss = 0;
    let work = 0;
    let E0 = 0;
    for (let k = 0; k < nFrames; k++) {
      const b = O_FSC + 6 * k;
      frameTimes[k] = k * samplesPerFrame * stepsPerSample * dt;
      energy[k] = out[b];
      const total = out[b] + out[b + 1];
      diss += out[b + 2];
      work += out[b + 3];
      if (k === 0) E0 = total;
      balance[k] = total + diss - work - E0;
      maxStretch = Math.max(maxStretch, out[b + 4]);
    }
    return {
      audio,
      fs: opt.fs,
      frames: out.slice(pk.framesAt),
      frameTimes,
      energy,
      balance,
      dt,
      stepsPerSample,
      omegaMax: pk.omegaMax,
      wallSeconds: (performance.now() - t0) / 1000,
      maxStretch,
      contactForce,
      malletPos: exc.mallet ? out.slice(O_MP, O_MP + 3 * nFrames) : undefined,
      loopSeconds,
      device: [adapter.info?.vendor, adapter.info?.architecture].filter(Boolean).join(' ') || 'GPU',
    };
  } finally {
    device.destroy();
  }
}