// Multi-dispatch backend: each stage of a time step is its own compute
// dispatch at full GPU width, and one step is 8 dispatches:
//   AB   hand state machine + contact + beam (workgroup 0), and the shell
//        derivative stencils / Airy right-hand side (all workgroups)
//   C    RHSAF = RHSA * F           (azimuthal analysis)
//   D    PH3   = Fi_j RHSAF_j       (Airy solve per wavenumber)
//   E    PHI   = PH3 * Gi           (azimuthal synthesis)
//   FG   NL from PHI (stress-function stencils recomputed inline), forcing
//   H    NLF   = NL * F
//   I    UH3N  = AiB UH3 + AiC UH13 + Ai NLF   (implicit update per column)
//   J    UNEW  = UH3N * Gi, pickup acceleration
//
// The step index k is kbase (state cell 0, written by the host once per
// chunk) plus s from a dynamic-offset uniform, so every recorded dispatch is
// identical up to that offset and slot rotation is a pure function of k.
// U/Uh slots rotate every step whether or not the shell is active: before
// first contact all slots are zero, so the rotation is harmless, and after
// first contact every step writes. The shell-active flag latches inside AB;
// stages of the same step read it a dispatch later, so they always see the
// latch, and the benign race inside AB itself can only cause zero-work on a
// zero field.

import { fmt } from './kernel.mjs';

export function buildStages(L, opts = {}) {
  const wg = opts.wg ?? Number(process.env.WG ?? 64);
  const { nrt, nth, N, nf, K, nsteps, off, A, c } = L;
  const { SC, ST_CVEL, ST_D, ST_U, ST_UH } = L.layout;

  const constLines = Object.entries(c)
    .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
    .join('\n');

  const header = `
const NRT = ${nrt}u;
const NTH = ${nth}u;
const N = ${N}u;
const NF = ${nf}u;
const NSTEPS = ${nsteps}u;
const KS = ${K}u;
const NRU = ${L.Nr}u;
const LS0 = ${L.LS0}u; const JS0 = ${L.JS0}u;
const LP0 = ${L.LP0}u; const JP0 = ${L.JP0}u;
const TIPI = ${L.TIPI0}u;
const NLR0 = ${L.NLR0}u;
const HOLDSTEPS = ${L.HOLDSTEPS};
${constLines}
${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')}
const ST_CVEL = ${ST_CVEL}u;
const ST_D = ${ST_D}u;
const ST_U = ${ST_U}u;
const ST_UH = ${ST_UH}u;
${L.scratchNames.map((n) => `const A_${n} = ${A[n]}u;`).join('\n')}

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

fn thm(j: u32) -> u32 { return (j + NTH - 1u) % NTH; }
fn thp(j: u32) -> u32 { return (j + 1u) % NTH; }
`;


  const stages = {};

  stages.AB = `${header}
var<workgroup> wgF: f32;
var<workgroup> wgHA: f32;
var<workgroup> wgHB: f32;
var<workgroup> wgHC: f32;
var<workgroup> shD: array<f32, ${nf}>;
var<workgroup> shD1: array<f32, ${nf}>;

fn u_dr(b: u32, l: u32, j: u32) -> f32 {
  if (l < 1u || l > NRT - 2u) { return 0.0; }
  return (st[b + (l + 1u) + j * NRT] - st[b + (l - 1u) + j * NRT]) * C1R;
}

@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
        @builtin(local_invocation_index) lid: u32,
        @builtin(workgroup_id) wgid: vec3<u32>) {
  // No early-out on k here: the host records exactly nsteps steps, and the
  // workgroupBarrier below requires uniform control flow up to this point.
  let k = u32(st[0u]) + ctl.x;
  let uS = ST_U + (k % 3u) * N;
  if (wgid.x == 0u) {
    if (lid == 0u) {
      var shellA = u32(st[1u]); var inC = u32(st[2u]); var phase = u32(st[3u]);
      var strokeIdx = u32(st[4u]); var holdLeft = i32(st[5u]); var vtar = st[6u];
      var vH = st[7u]; var uHm = st[8u]; var uH0 = st[9u]; var uHp = st[10u];
      var peakXi = st[12u]; var peakDisp = st[13u];
      let dCur = ST_D + ((k + 2u) % 3u) * NF;
      let dOld = ST_D + ((k + 1u) % 3u) * NF;
      if (shellA != 0u && k > 0u && (k % 512u) == 0u) {
        var pd: f32 = 0.0;
        for (var l: u32 = 0u; l < NRT; l = l + 1u) {
          pd = max(pd, abs(st[uS + l + JS0 * NRT]));
        }
        if (pd > peakDisp) { peakDisp = pd; }
      }
      let xi = st[uS + LS0 + JS0 * NRT] - st[dCur + TIPI] - GAP;
      var fM: f32 = 0.0;
      if (xi > 0.0) {
        fM = KM * pow(xi, ALM);
        shellA = 1u;
        if (inC == 0u && strokeIdx < KS && st[ST_CVEL + strokeIdx] == 0.0) {
          st[ST_CVEL + strokeIdx] = (st[dCur + TIPI] - st[dOld + TIPI]) / HT;
        }
        inC = 1u;
        if (xi > peakXi) { peakXi = xi; }
      } else {
        inC = 0u;
      }
      outb[NSTEPS + k] = fM;
      wgF = fM;
      wgHA = -(uHp - 2.0 * uH0 + uHm);
      wgHB = -(HT * 0.5) * (uHp - uHm);
      wgHC = -HT * HT * (PS * uH0 + ((1.0 - PS) * 0.5) * (uHp + uHm));
      if (phase == 0u) {
        if (strokeIdx < KS && f32(k) >= ops[OP_KGO + strokeIdx]) {
          phase = 1u; vtar = ops[OP_VSTR + strokeIdx];
        }
      } else if (phase == 1u) {
        vH = max(vH - AH * HT, -vtar);
        if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
        else if (vH <= -vtar) { phase = 2u; }
      } else if (phase == 2u) {
        if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
        else if (uHp <= -MISSD) { phase = 4u; }
      } else if (phase == 3u) {
        holdLeft = holdLeft - 1;
        if (holdLeft <= 0) { phase = 4u; }
      } else if (phase == 4u) {
        vH = min(vH + AHR * HT, 0.5 * vtar);
        if (vH > 0.0 && uHp >= -vH * vH / (2.0 * AHR)) { phase = 5u; }
      } else {
        vH = vH - AHR * HT;
        if (vH <= 0.0 || uHp >= 0.0) { vH = 0.0; phase = 0u; strokeIdx = strokeIdx + 1u; }
      }
      uHm = uH0; uH0 = uHp; uHp = uHp + vH * HT;
      st[1u] = f32(shellA); st[2u] = f32(inC); st[3u] = f32(phase);
      st[4u] = f32(strokeIdx); st[5u] = f32(holdLeft); st[6u] = vtar;
      st[7u] = vH; st[8u] = uHm; st[9u] = uH0; st[10u] = uHp;
      st[12u] = peakXi; st[13u] = peakDisp; st[14u] = fM;
    }
    if (lid < NF) {
      shD[lid] = st[ST_D + ((k + 2u) % 3u) * NF + lid];
      shD1[lid] = st[ST_D + ((k + 1u) % 3u) * NF + lid];
    }
    workgroupBarrier();
    if (lid < NF) {
      let r = lid;
      var sum: f32 = ops[OP_VM + r] * wgHA + ops[OP_VR + r] * wgHB + ops[OP_VK + r] * wgHC
        + (HT * HT * wgF) * ops[OP_TIPCOL + r];
      for (var p: u32 = 0u; p < NF; p = p + 1u) {
        sum = sum + ops[OP_ASIBT + p + r * NF] * shD[p]
                  + ops[OP_ASICT + p + r * NF] * shD1[p];
      }
      st[ST_D + (k % 3u) * NF + r] = sum;
    }
  }
  let i = gid.x;
  if (i >= N) { return; }
  if (st[1u] == 0.0) { return; }
  let l = i % NRT; let j = i / NRT;
  var ur: f32 = 0.0; var urr: f32 = 0.0;
  if (l >= 1u && l <= NRT - 2u) {
    let um = st[uS + (l - 1u) + j * NRT];
    let u0 = st[uS + l + j * NRT];
    let up = st[uS + (l + 1u) + j * NRT];
    ur = (up - um) * C1R;
    urr = (um - 2.0 * u0 + up) * C2R;
  }
  let ujm = st[uS + l + thm(j) * NRT];
  let uj0 = st[uS + l + j * NRT];
  let ujp = st[uS + l + thp(j) * NRT];
  let v1 = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T);
  let urt = (u_dr(uS, l, thp(j)) - u_dr(uS, l, thm(j))) * C1T;
  let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * ((ujp - ujm) * C1T);
  st[A_URR + i] = urr;
  st[A_V1 + i] = v1;
  st[A_V2 + i] = v2;
  var rv: f32 = 0.0;
  if (l >= 3u && l <= NRU - 2u) {
    rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2);
  }
  st[A_RHSA + i] = rv;
}
`;

  const gate = `
  let k = u32(st[0u]) + ctl.x;
  if (k >= NSTEPS) { return; }
  let i = gid.x;
  if (i >= N) { return; }
  if (st[1u] == 0.0) { return; }
  let l = i % NRT; let j = i / NRT;`;

  const transform = (src, dst, mat) => `${header}
@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
${gate}
  var sum: f32 = 0.0;
  for (var q: u32 = 0u; q < NTH; q = q + 1u) {
    sum = sum + st[${src} + l + q * NRT] * ops[${mat} + q + j * NTH];
  }
  st[${dst} + i] = sum;
}
`;

  stages.C = transform('A_RHSA', 'A_RHSAF', 'OP_F');

  stages.D = `${header}
@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
${gate}
  let base = OP_FI + 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[A_RHSAF + p + j * NRT];
  }
  st[A_PH3 + i] = sum;
}
`;

  stages.E = transform('A_PH3', 'A_PHI', 'OP_GI');

  stages.FG = `${header}
fn p_dr(l: u32, j: u32) -> f32 {
  if (l < 1u || l > NRT - 2u) { return 0.0; }
  return (st[A_PHI + (l + 1u) + j * NRT] - st[A_PHI + (l - 1u) + j * NRT]) * C1R;
}
@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
${gate}
  var nl: f32 = 0.0;
  if (l >= NLR0 && l < NRU) {
    let pm = st[A_PHI + (l - 1u) + j * NRT];
    let p0 = st[A_PHI + l + j * NRT];
    let pp = st[A_PHI + (l + 1u) + j * NRT];
    let pr = (pp - pm) * C1R;
    let prr = (pm - 2.0 * p0 + pp) * C2R;
    let pjm = st[A_PHI + l + thm(j) * NRT];
    let pjp = st[A_PHI + l + thp(j) * NRT];
    let w1 = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * p0 + pjp) * C2T);
    let prt = (p_dr(l, thp(j)) - p_dr(l, thm(j))) * C1T;
    let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * ((pjp - pjm) * C1T);
    nl = prr * st[A_V1 + i] + st[A_URR + i] * w1 - 2.0 * w2 * st[A_V2 + i]
       - (prr + w1) * RSI;
    if (l == LS0 && j == JS0) { nl = nl - st[14u] * FSC; }
    nl = nl * SNL;
  }
  st[A_NL + i] = nl;
}
`;

  stages.H = transform('A_NL', 'A_NLF', 'OP_F');

  stages.I = `${header}
@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
${gate}
  let h3S = ST_UH + (k % 3u) * N;
  let h13S = ST_UH + ((k + 2u) % 3u) * N;
  let h3nS = ST_UH + ((k + 1u) % 3u) * N;
  let base = j * NRT * NRT;
  var sum: f32 = 0.0;
  for (var p: u32 = 0u; p < NRT; p = p + 1u) {
    let pj = p + j * NRT;
    sum = sum + ops[OP_AIB + base + l + p * NRT] * st[h3S + pj]
              + ops[OP_AIC + base + l + p * NRT] * st[h13S + pj]
              + ops[OP_AI + base + l + p * NRT] * st[A_NLF + pj];
  }
  st[h3nS + i] = sum;
}
`;

  stages.J = `${header}
@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
${gate}
  let h3nS = ST_UH + ((k + 1u) % 3u) * N;
  let uS = ST_U + (k % 3u) * N;
  let u1S = ST_U + ((k + 2u) % 3u) * N;
  let unS = ST_U + ((k + 1u) % 3u) * N;
  var sum: f32 = 0.0;
  for (var q: u32 = 0u; q < NTH; q = q + 1u) {
    sum = sum + st[h3nS + l + q * NRT] * ops[OP_GI + q + j * NTH];
  }
  st[unS + i] = sum;
  if (l == LP0 && j == JP0) {
    outb[k] = (sum - 2.0 * st[uS + i] + st[u1S + i]) * HT2I;
  }
}
`;

  const nwgMain = Math.ceil(N / wg);
  return {
    order: ['AB', 'C', 'D', 'E', 'FG', 'H', 'I', 'J'],
    sources: stages,
    wg,
    nwg: nwgMain,
  };
}