// Multi-dispatch step: each stage its own dispatch at full GPU width.
// Representative step = 2 stencil stages + 4 transforms + 2 page solves,
// recorded 512 steps per command buffer.
import { installNodeWebGpu, requestDevice } from './gpu.mjs';

const NRT = 26, NTH = 60, N = NRT * NTH;
const WG = Number(process.env.WG ?? 128);
const NWG = Math.ceil(N / WG);
const STEPS = 512;

const DYNOFF = !!process.env.DYNOFF;
const GATE = !!process.env.GATE;
const HDR = `
const NRT = ${NRT}u; const NTH = ${NTH}u; const N = ${N}u;
@group(0) @binding(0) var<storage, read> ops: array<f32>;
@group(0) @binding(1) var<storage, read_write> st: array<f32>;
${DYNOFF ? '@group(0) @binding(2) var<uniform> ctl: vec4<u32>;' : ''}
`;

const KERNELS = {
  stencil: `${HDR}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= N) { return; }
  ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
  let l = i % NRT; let j = i / NRT;
  let jm = (j + NTH - 1u) % NTH; let jp = (j + 1u) % NTH;
  var v: f32 = 0.0;
  if (l >= 1u && l <= NRT - 2u) {
    v = (st[N + (l - 1u) + j * NRT] - 2.0 * st[N + l + j * NRT] + st[N + (l + 1u) + j * NRT]) * 0.5
      + (st[N + l + jm * NRT] - st[N + l + jp * NRT]) * 0.25;
  }
  st[2u * N + i] = v;
}`,
  transform: `${HDR}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= N) { return; }
  ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
  let l = i % NRT; let j = i / NRT;
  var sum: f32 = 0.0;
  for (var q: u32 = 0u; q < NTH; q = q + 1u) {
    sum = sum + st[N + l + q * NRT] * ops[q + j * NTH];
  }
  st[2u * N + i] = sum + f32(${DYNOFF ? 'ctl.x' : '0u'});
}`,
  page: `${HDR}
@compute @workgroup_size(${WG})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= N) { return; }
  ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
  let l = i % NRT; let j = i / NRT;
  let base = j * NRT * NRT;
  var sum: f32 = 0.0;
  for (var p: u32 = 0u; p < NRT; p = p + 1u) {
    sum = sum + ops[base + l + p * NRT] * st[N + p + j * NRT];
  }
  st[2u * N + i] = sum;
}`,
};

async function main() {
  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}, wg size ${WG} x ${NWG} workgroups`);

  const BIG = Number(process.env.BIG ?? 1);
  const ops = device.createBuffer({ size: 4 * NRT * NRT * NTH * BIG, usage: GPUBufferUsage.STORAGE });
  const stb = device.createBuffer({ size: 4 * N * 4 * BIG, usage: GPUBufferUsage.STORAGE });
  const uni = device.createBuffer({ size: STEPS * 256, usage: GPUBufferUsage.UNIFORM });

  const pipes = {};
  for (const [name, code] of Object.entries(KERNELS)) {
    const module = device.createShaderModule({ code });
    const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
    if (errs.length) throw new Error(name + ': ' + errs.map((m) => m.message).join('\n'));
    const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
    const entries = [
      { binding: 0, resource: { buffer: ops } },
      { binding: 1, resource: { buffer: stb } },
    ];
    if (DYNOFF && name === 'transform') entries.push({ binding: 2, resource: { buffer: uni, offset: 0, size: 16 } });
    const bg = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries,
    });
    pipes[name] = { pipeline, bg };
  }

  let stepSeq = ['stencil', 'transform', 'page', 'transform', 'stencil', 'transform', 'page', 'transform'];
  if (process.env.SEQ1) stepSeq = ['transform'];
  if (process.env.SEQN) stepSeq = Array(Number(process.env.SEQN)).fill('transform');
  const record = () => {
    const enc = device.createCommandEncoder();
    const pass = enc.beginComputePass();
    for (let s = 0; s < STEPS; s++) {
      for (const name of stepSeq) {
        const p = pipes[name];
        pass.setPipeline(p.pipeline);
        pass.setBindGroup(0, p.bg);
        pass.dispatchWorkgroups(NWG);
      }
    }
    pass.end();
    return enc.finish();
  };

  device.queue.writeBuffer(stb, 4, new Float32Array([1]));
  // warmup
  device.queue.submit([record()]);
  await device.queue.onSubmittedWorkDone();

  let tRec = 0;
  const t0 = performance.now();
  const reps = Number(process.env.REPS ?? 4);
  for (let r = 0; r < reps; r++) {
    const tr = performance.now();
    const cb = record();
    tRec += performance.now() - tr;
    device.queue.submit([cb]);
  }
  await device.queue.onSubmittedWorkDone();
  const ms = (performance.now() - t0) / reps;
  console.log(`8-dispatch step: ${(1000 * ms / STEPS).toFixed(2)} us/step total (cpu recording ${(1000 * tRec / reps / STEPS).toFixed(2)} us/step)`);
}

main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });