/**
 * 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 three compute dispatches for a rectangle with supported or
 * clamped edges, and three or five when the setup carries sparse operators
 * (free edges, or an ellipse of any edge condition); the dispatch boundaries
 * are the global synchronization points the stencil needs:
 *   op       Lw = Laplacian of w at every node (ghost layer by reflection), or,
 *            with sparse operators, Lw = Lm w / mvec and B = Km w / mvec
 *   update   the new field, wn = c_new (2 w - c_old wm - k2dt2 B + s1dt2 (Lw - Lw_m)
 *            + fdt2 F(n) s); edges zeroed for the dense rectangle, the nodes
 *            that are not unknowns (S.keep = 0) zeroed for the sparse form
 *   dots     (nr > 0) the rigid-body components Qb' wn
 *   project  (nr > 0) wn - Rb (Qb' wn)
 *   reduce   acceleration summed per listener delay bin into the audio bins,
 *            the probe displacement and curvatures, max |wn| 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 state is carried as
 * (w, d = w - wm) rather than (w, wm): the update is algebraically the .m's,
 *   dn = (d + rest) - cs (2 d + rest),   wn = w + dn,   cs = sigma0 dt / (1 + sigma0 dt),
 * with rest = -k2dt2 B + s1dt2 (Lw - Lw_m) + fdt2 F(n) s, so that no factor of
 * the form 1 - epsilon is stored in f32 and the rounding of w (about 6e-8 of a
 * micrometre-scale displacement) does not feed back into the velocity as it
 * would through 2 w - wm. The acceleration for the audio is (dn - d) / dt^2.
 */

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;
  bc: string;
  /** Load footprint, ny x nx column-major */
  s: Float64Array;
  /** Force at the nsteps + 1 time levels */
  F: Float64Array;
  /** Probe row and column */
  ip: number;
  jp: number;
  /** Steps at which the field is kept, sorted and unique */
  frameSteps: Float64Array;
  /** Listener delay bin and weight of every node */
  bins: Float64Array;
  wgt: Float64Array;
  nbins: number;
  na: number;
  /** The loop uses the sparse operators (free edges, or an ellipse) */
  sparse: boolean;
  /** Sparse form only: 1 at the unknowns and 0 at the nodes held at zero (N), node weights (N) */
  keep: Float64Array | null;
  mvec: Float64Array | null;
  /** Rigid-body motions to remove: basis and projector (N x nr), nr = 0, 2 or 3 */
  Rb: Float64Array | null;
  Qb: Float64Array | null;
  nr: number;
  Km: Triplets | null;
  Lm: Triplets | null;
  sigma0: number;
  k2dt2: number;
  s1dt2: number;
  fdt2: number;
  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;
  /** ny x nx x nf */
  Wf: 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 };
}

/** Nodes grouped by listener delay bin, as a CSR-like list. */
function binLists(bins: Float64Array, nbins: number): { ptr: Uint32Array; nodes: Uint32Array } {
  const n = bins.length;
  const count = new Uint32Array(nbins);
  for (let i = 0; i < n; i++) count[bins[i] - 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 nodes = new Uint32Array(n);
  for (let i = 0; i < n; i++) nodes[fill[bins[i] - 1]++] = i;
  return { ptr, nodes };
}

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 sp = su.sparse;
  const nr = sp ? su.nr : 0;
  // With a projection the frames are written after it; otherwise straight from the update
  const framesInUpdate = nr === 0;
  if (su.s.length !== N || su.bins.length !== N || su.wgt.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");

  // ---- 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 nstep = k + 1;
    const ia = Math.floor(nstep * su.dt * su.fs); // ia = floor(nstep * dt * fs) + 1, made 0-based
    if (ia >= na) throw new Error("audio bin index out of range");
    iaOf[k] = ia;
  }
  const bl = binLists(su.bins, nbins);

  const f32 = (a: ArrayLike<number>) => Float32Array.from(a);
  const opsParts: [string, ArrayLike<number>][] = [
    ["S", f32(su.s)],
    ["WGT", f32(su.wgt)],
    ["F", f32(su.F.subarray(0, nsteps))],
  ];
  const idxParts: [string, ArrayLike<number>][] = [
    ["BINPTR", bl.ptr],
    ["BINNODES", bl.nodes],
    ["FRAME", frameOf],
    ["IA", iaOf],
  ];
  if (sp) {
    if (!su.mvec || !su.keep || !su.Km || !su.Lm)
      throw new Error("the sparse form needs mvec, keep, Km and Lm");
    if (su.mvec.length !== N || su.keep.length !== N)
      throw new Error("node arrays do not match the grid");
    const K = toCsr(su.Km, N);
    const L = toCsr(su.Lm, N);
    opsParts.push(["MVEC", f32(su.mvec)], ["KEEP", f32(su.keep)], ["KV", K.val], ["LV", L.val]);
    idxParts.push(["KPTR", K.ptr], ["KCOL", K.col], ["LPTR", L.ptr], ["LCOL", L.col]);
    if (nr > 0) {
      if (!su.Rb || !su.Qb || su.Rb.length !== N * nr || su.Qb.length !== N * nr)
        throw new Error("rigid-body arrays do not match the grid");
      opsParts.push(["RB", f32(su.Rb)], ["QB", f32(su.Qb)]);
    }
  }
  const ops = pack(opsParts, (n) => new Float32Array(n));
  const idx = pack(idxParts, (n) => new Uint32Array(n));

  // ---- state buffer: two slots each of w, d = w - wm and Lw; B; a few scalars
  const ST_W = 0;
  const ST_D = 2 * N;
  const ST_L = 4 * N;
  const ST_B = 6 * N;
  const ST_AUX = 7 * N;
  const stLen = ST_AUX + 16;

  const nprobe = nsteps + 1;
  const OUT_PB = 4 * nprobe;
  const outLen = OUT_PB + nbins * na;

  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),
    K2DT2: su.k2dt2,
    S1DT2: su.s1dt2,
    FDT2: su.fdt2,
    CS: s0dt / (1 + s0dt),
    INV_DT2: 1 / (su.dt * su.dt),
    DTFS: su.dt * su.fs,
    BLOWUP: su.blowup,
    SGN: su.bc === "clamped" ? 1 : -1,
  };
  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 NXI: i32 = ${nx}; const NYI: i32 = ${ny};
const NSTEPS = ${nsteps}u; const NPROBE = ${nprobe}u;
const NBINS = ${nbins}u; const NA = ${na}u;
const IPROBE = ${ip + jp * ny}u;
const NR = ${nr}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_W = ${ST_W}u; const ST_D = ${ST_D}u; const ST_L = ${ST_L}u; const ST_B = ${ST_B}u; const ST_AUX = ${ST_AUX}u;
const OUT_PB = ${OUT_PB}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 |w| 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 (w, d, Lw) 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; }
`;

  // Padded field of plate_pad: the edge nodes are zero and one ghost layer
  // reflects the first interior value, odd for supported, even for clamped.
  const opSupported = `${header}
fn wat(wS: u32, i: i32, j: i32) -> f32 {
  var ii = i; var jj = j; var sg: f32 = 1.0;
  if (ii < 0) { ii = 1; sg = sg * SGN; } else if (ii >= NYI) { ii = NYI - 2; sg = sg * SGN; }
  if (jj < 0) { jj = 1; sg = sg * SGN; } else if (jj >= NXI) { jj = NXI - 2; sg = sg * SGN; }
  return sg * st[wS + u32(ii) + u32(jj) * NY];
}
@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 wS = ST_W + cur(k);
  let r = i32(i % NY); let cc = i32(i / NY);
  let w0 = st[wS + i];
  let lw = (wat(wS, r + 1, cc) - 2.0 * w0 + wat(wS, r - 1, cc)) * INV_DY2
         + (wat(wS, r, cc + 1) - 2.0 * w0 + wat(wS, r, cc - 1)) * INV_DX2;
  st[ST_L + cur(k) + i] = lw;
}
`;

  // Sparse form: Lw = Lm w / mvec and B = Km w / mvec from the sparse operators
  const opSparse = `${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 wS = ST_W + cur(k);
  var lsum: f32 = 0.0;
  for (var p = idx[IX_LPTR + i]; p < idx[IX_LPTR + i + 1u]; p = p + 1u) {
    lsum = lsum + ops[OP_LV + p] * st[wS + idx[IX_LCOL + p]];
  }
  var ksum: f32 = 0.0;
  for (var p = idx[IX_KPTR + i]; p < idx[IX_KPTR + i + 1u]; p = p + 1u) {
    ksum = ksum + ops[OP_KV + p] * st[wS + idx[IX_KCOL + p]];
  }
  let m = ops[OP_MVEC + i];
  st[ST_L + cur(k) + i] = lsum / m;
  st[ST_B + i] = ksum / m;
}
`;

  const biharmonic = sp
    ? `let b = st[ST_B + i];`
    : `let lS = ST_L + cur(k);
  let lw0 = st[lS + i];
  let b = (st[lS + i + 1u] - 2.0 * lw0 + st[lS + i - 1u]) * INV_DY2
        + (st[lS + i + NY] - 2.0 * lw0 + st[lS + i - NY]) * INV_DX2;`;
  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 r = i % NY; let cc = i / NY;
  let wS = ST_W + cur(k); let wnS = ST_W + nxt(k);
  let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
  var wn: f32 = 0.0; var dn: f32 = 0.0;
  ${sp ? "if (true) {" : "if (r > 0u && r < NY - 1u && cc > 0u && cc < NX - 1u) {"}
    ${biharmonic}
    let lw = st[ST_L + cur(k) + i];
    let lwm = st[ST_L + nxt(k) + i];
    let rest = -K2DT2 * b + S1DT2 * (lw - lwm) + FDT2 * ops[OP_F + k] * ops[OP_S + i];
    let d = st[dS + i];
    dn = (d + rest) - CS * (2.0 * d + rest);
    wn = st[wS + i] + dn;
  }
  ${sp ? "let kp = ops[OP_KEEP + i]; wn = wn * kp; dn = dn * kp;" : ""}
  st[wnS + i] = wn;
  st[dnS + i] = dn;
  ${
    framesInUpdate
      ? `let f = idx[IX_FRAME + k];
  if (f != ${NONE}u) { frames[f * N + i] = wn; }`
      : ""
  }
}
`;

  // Rigid-body components Qb' wn (the nodes held at zero already are)
  const dots = `${header}
var<workgroup> sh0: array<f32, ${RWG}>;
var<workgroup> sh1: array<f32, ${RWG}>;
var<workgroup> sh2: array<f32, ${RWG}>;
@compute @workgroup_size(${RWG})
fn main(@builtin(local_invocation_index) lid: u32) {
  let k = ctl[0] + uni.x;
  let wnS = ST_W + nxt(k);
  var a0: f32 = 0.0; var a1: f32 = 0.0; var a2: f32 = 0.0;
  for (var i = lid; i < N; i = i + ${RWG}u) {
    let wn = st[wnS + i];
    a0 = a0 + ops[OP_QB + i] * wn;
    a1 = a1 + ops[OP_QB + N + i] * wn;
    if (NR == 3u) { a2 = a2 + ops[OP_QB + 2u * N + i] * wn; }
  }
  sh0[lid] = a0; sh1[lid] = a1; sh2[lid] = a2;
  workgroupBarrier();
  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]; }
    workgroupBarrier();
  }
  if (lid == 0u) { st[ST_AUX] = sh0[0]; st[ST_AUX + 1u] = sh1[0]; st[ST_AUX + 2u] = sh2[0]; }
}
`;

  const project = `${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 wnS = ST_W + nxt(k); let dnS = ST_D + nxt(k);
  var corr = ops[OP_RB + i] * st[ST_AUX] + ops[OP_RB + N + i] * st[ST_AUX + 1u];
  if (NR == 3u) { corr = corr + ops[OP_RB + 2u * N + i] * st[ST_AUX + 2u]; }
  // Rb is zero at the nodes held at zero; the factor makes that exact in f32 as well
  let kp = ops[OP_KEEP + i];
  let wn = (st[wnS + i] - corr) * kp; let dn = (st[dnS + i] - corr) * kp;
  st[wnS + i] = wn; st[dnS + i] = dn;
  let f = idx[IX_FRAME + k];
  if (f != ${NONE}u) { frames[f * N + i] = wn; }
}
`;

  // Workgroups 0..NBINS-1 sum the weighted acceleration of their bin's nodes into
  // the audio sample of this step; workgroup NBINS takes max |wn| for the
  // blow-up test and records the probe.
  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 wnS = ST_W + nxt(k);
  let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
  let wg = wgid.x;
  if (wg < NBINS) {
    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_BINNODES + p];
      a = a + ops[OP_WGT + i] * ((st[dnS + i] - st[dS + i]) * 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] * DTFS;
    }
  } else {
    var m: f32 = 0.0;
    for (var i = lid; i < N; i = i + ${RWG}u) { m = max(m, abs(st[wnS + 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) {
      let c = st[wnS + IPROBE];
      let e = st[wnS + IPROBE + NY]; let w = st[wnS + IPROBE - NY];
      let n = st[wnS + IPROBE + 1u]; let so = st[wnS + IPROBE - 1u];
      let ne = st[wnS + IPROBE + NY + 1u]; let nw = st[wnS + IPROBE - NY + 1u];
      let se = st[wnS + IPROBE + NY - 1u]; let sw = st[wnS + IPROBE - NY - 1u];
      outb[k + 1u] = c;
      outb[NPROBE + k + 1u] = (e - 2.0 * c + w) * INV_DX2;
      outb[2u * NPROBE + k + 1u] = (n - 2.0 * c + so) * INV_DY2;
      outb[3u * NPROBE + k + 1u] = (ne - se - nw + sw) * INV_DXDY4;
      if (sh[0] >= BLOWUP) { ctl[1] = k + 1u; }
    }
  }
}
`;

  const op = sp ? opSparse : opSupported;
  const sources = nr > 0 ? [op, update, dots, project, reduce] : [op, update, reduce];
  const groups =
    nr > 0
      ? [Math.ceil(N / WG), Math.ceil(N / WG), 1, Math.ceil(N / WG), nbins + 1]
      : [Math.ceil(N / WG), Math.ceil(N / WG), nbins + 1];
  return { N, nf, nprobe, nbins, na, outLen, stLen, ops: ops.data, idx: idx.data, sources, groups };
}

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);
  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 np = b.nprobe;
    return {
      w_p: out.slice(0, np),
      wxx_p: out.slice(np, 2 * np),
      wyy_p: out.slice(2 * np, 3 * np),
      wxy_p: out.slice(3 * np, 4 * np),
      Wf: Float32Array.from(Wf),
      Pbins: out.slice(4 * np, 4 * np + b.nbins * b.na),
      divergedStep,
      gpuMs,
    };
  } finally {
    for (const buf of all) buf.destroy();
  }
}