// vec4 variants with the radial dimension padded to 28 (7 vec4 lanes).
// Transform: out[l,j] = sum_q X[l,q] M[q,j]; X stored 28 x 60, M 60 x 60.
// Page: out[l,j] = sum_p A[l,p,j] X[p,j]; A pages 28 x 28.
import { installNodeWebGpu, requestDevice } from './gpu.mjs';

const PADL = 28, NTH = 60, NP = PADL * NTH; // 1680
const NPV = NP / 4;
const WS = 256;
const STEPS = 2048;

function wrap(body) {
  return `
const PADL = ${PADL}u; const NTH = ${NTH}u; const NP = ${NP}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
@group(0) @binding(0) var<storage, read> ops: array<f32>;
@group(0) @binding(1) var<storage, read> opsv: array<vec4<f32>>;
@group(0) @binding(2) var<storage, read_write> stv: array<vec4<f32>>;
fn ld(i: u32) -> f32 { return stv[i / 4u][i % 4u]; }
@compute @workgroup_size(${WS})
fn main(@builtin(local_invocation_index) lid: u32) {
  for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
${body}
    workgroupBarrier();
    storageBarrier();
  }
  if (lid == 0u) { stv[0].x = stv[0].y + 1.0; }
}
`;
}

const CASES = {
  // 60 threads, 7 vec4 accumulators, X read as vec4, M scalar broadcast
  colRegsVec4: wrap(`
    if (lid < NTH) {
      let j = lid;
      var acc: array<vec4<f32>, 7>;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        let xb = (NP + q * PADL) / 4u;
        for (var v: u32 = 0u; v < 7u; v = v + 1u) {
          acc[v] = acc[v] + stv[xb + v] * mq;
        }
      }
      let ob = (2u * NP + j * PADL) / 4u;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
    }`),

  // 120 threads: j column plus upper/lower half (4 vec4s h=0, 3 vec4s h=1)
  colRegsVec4x120: wrap(`
    if (lid < 120u) {
      let j = lid % NTH; let h = lid / NTH;
      let v0 = h * 4u;
      let nv = 4u - h;   // 4 or 3
      var acc: array<vec4<f32>, 4>;
      for (var v: u32 = 0u; v < 4u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        let xb = (NP + q * PADL) / 4u + v0;
        for (var v: u32 = 0u; v < 4u; v = v + 1u) {
          if (v < nv) { acc[v] = acc[v] + stv[xb + v] * mq; }
        }
      }
      let ob = (2u * NP + j * PADL) / 4u + v0;
      for (var v: u32 = 0u; v < 4u; v = v + 1u) {
        if (v < nv) { stv[ob + v] = acc[v]; }
      }
    }`),

  // page solve: 60 threads, A columns read as vec4 (28-padded), X broadcast
  pageVec4: wrap(`
    if (lid < NTH) {
      let j = lid;
      let base = j * PADL * PADL;
      var acc: array<vec4<f32>, 7>;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
      for (var p: u32 = 0u; p < PADL; p = p + 1u) {
        let xq = ld(NP + p + j * PADL);
        let ab = (base + p * PADL) / 4u;
        for (var v: u32 = 0u; v < 7u; v = v + 1u) {
          acc[v] = acc[v] + opsv[ab + v] * xq;
        }
      }
      let ob = (2u * NP + j * PADL) / 4u;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
    }`),

  // triple page solve (the shell update): same but three A blocks, three X
  page3Vec4: wrap(`
    if (lid < NTH) {
      let j = lid;
      let base = j * PADL * PADL;
      var acc: array<vec4<f32>, 7>;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
      for (var p: u32 = 0u; p < PADL; p = p + 1u) {
        let x1 = ld(NP + p + j * PADL);
        let x2 = ld(2u * NP + p + j * PADL);
        let x3 = ld(3u * NP + p + j * PADL);
        let ab = (base + p * PADL) / 4u;
        let s2 = ${PADL * PADL * NTH / 4}u;
        for (var v: u32 = 0u; v < 7u; v = v + 1u) {
          acc[v] = acc[v] + opsv[ab + v] * x1 + opsv[ab + v + s2] * x2 + opsv[ab + v + 2u * s2] * x3;
        }
      }
      let ob = (j * PADL) / 4u;
      for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
    }`),
};

async function time(device, code, label) {
  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(label + ':\n' + errs.map((m) => `${m.lineNum}: ${m.message}`).join('\n'));
  const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
  const ops = device.createBuffer({ size: 4 * PADL * PADL * NTH * 3, usage: GPUBufferUsage.STORAGE });
  const stb = device.createBuffer({ size: 4 * NP * 4, usage: GPUBufferUsage.STORAGE });
  const entries = [];
  if (/\bops\[/.test(code)) entries.push({ binding: 0, resource: { buffer: ops } });
  if (/\bopsv\[/.test(code)) entries.push({ binding: 1, resource: { buffer: ops } });
  entries.push({ binding: 2, resource: { buffer: stb } });
  const bg = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries,
  });
  const run = async (reps) => {
    const enc = device.createCommandEncoder();
    for (let r = 0; r < reps; r++) {
      const pass = enc.beginComputePass();
      pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
    }
    device.queue.submit([enc.finish()]);
    await device.queue.onSubmittedWorkDone();
  };
  await run(1);
  const t0 = performance.now();
  await run(4);
  const ms = (performance.now() - t0) / 4;
  console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us per stage-step`);
  ops.destroy(); stb.destroy();
}

async function main() {
  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}`);
  for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
}

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