// WGSL megakernel generator for the cymbal2 time loop. One workgroup loops
// over CHUNK time steps inside a single dispatch, with barriers between the
// stages of each step; every dispatch is identical, so a whole render is a
// queue of identical dispatches with no CPU in the loop. All constants
// (grid sizes, node indices, physics scalars, the stroke schedule thresholds)
// are baked into the shader source; the operators live in one read-only
// storage buffer, the evolving state and scratch in a read-write one, the
// acceleration and force histories in a third.
//
// The stage structure mirrors cymbal2_sim.m line for line, with the sparse
// radial (D1, D2) and azimuthal (T1t, T2) difference matrices applied as
// stencils and the dense transforms (F, Gi), the per-wavenumber page solves
// (Fi, Ai, AiB, AiC) and the beam matrices applied as stored matrices.

export const fmt = (x) => {
  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';
};

// Shared between the megakernel and the multi-dispatch backend: the packed
// operator buffer, the state-buffer layout, and the baked constants.
export function buildLayout(setup) {
  const nrt = setup.nrt | 0;
  const nth = setup.nth | 0;
  const nf = setup.nfree | 0;
  const N = nrt * nth;
  const K = setup.K | 0;

  const o = setup.ops;
  const parts = [
    ['D1', o.D1, nrt * nrt],
    ['D2', o.D2, nrt * nrt],
    ['F', o.F, nth * nth],
    ['GI', o.Gi, nth * nth],
    ['RD', o.Rd, nrt],
    ['R2D', o.R2d, nrt],
    ['AIB', o.AiB, nrt * nrt * nth],
    ['AIC', o.AiC, nrt * nrt * nth],
    ['AI', o.Ai, nrt * nrt * nth],
    ['FI', o.Fi, nrt * nrt * nth],
    ['ASIB', o.ASiB, nf * nf],
    ['ASIC', o.ASiC, nf * nf],
    ['VM', o.vM, nf],
    ['VR', o.vR, nf],
    ['VK', o.vK, nf],
    ['TIPCOL', o.tipcol, nf],
  ];
  const ht = setup.ht;
  const kgo = setup.tlaunch.map((tl) => {
    let g = Math.max(0, Math.ceil(tl / ht) - 1);
    while (g > 0 && g * ht >= tl) g -= 1;
    while ((g + 1) * ht < tl) g += 1;
    return g;
  });
  parts.push(['KGO', Float32Array.from(kgo), K]);
  parts.push(['VSTR', Float32Array.from(setup.vstr), K]);
  // Transposed beam matrices: each beam row thread streams its own row
  // sequentially instead of striding across columns.
  const tr = (m) => {
    const t = new Float32Array(nf * nf);
    for (let r = 0; r < nf; r++) for (let p = 0; p < nf; p++) t[p + r * nf] = m[r + p * nf];
    return t;
  };
  parts.push(['ASIBT', tr(o.ASiB), nf * nf]);
  parts.push(['ASICT', tr(o.ASiC), nf * nf]);

  const off = {};
  let cursor = 0;
  for (const [name, arr, len] of parts) {
    if (arr.length !== len) throw new Error(`ops part ${name}: ${arr.length} != ${len}`);
    off[name] = cursor;
    cursor += len;
  }
  const opsData = new Float32Array(cursor);
  for (const [name, arr] of parts) opsData.set(arr, off[name]);

  const SC = 0;
  const ST_CVEL = 16;
  const ST_D = ST_CVEL + K;
  const ST_U = ST_D + 3 * nf;
  const ST_UH = ST_U + 3 * N;
  const scratchNames = [
    'UR', 'URR', 'V1', 'UT1', 'V2', 'RHSA', 'RHSAF', 'PH3',
    'PHI', 'PR', 'PRR', 'W1', 'PHT1', 'NL', 'NLF',
  ];
  const A = {};
  let scur = ST_UH + 3 * N;
  for (const n of scratchNames) {
    A[n] = scur;
    scur += N;
  }
  const stLen = scur;

  const c = {
    HT: ht,
    HT2I: 1 / (ht * ht),
    GAP: setup.gap,
    KM: setup.KM,
    ALM: setup.alM,
    PS: setup.ps,
    AH: setup.aH,
    AHR: setup.aHr,
    MISSD: setup.missDepth,
    EHRS: setup.Eh / setup.Rs,
    EH2: setup.Eh / 2,
    RSI: 1 / setup.Rs,
    SNL: (ht * ht) / setup.rhoh,
    FSC: setup.fsc,
    C1R: 1 / (2 * setup.hr),
    C2R: 1 / (setup.hr * setup.hr),
    C1T: 1 / (2 * setup.hth),
    C2T: 1 / (setup.hth * setup.hth),
  };

  return {
    nrt, nth, Nr: setup.Nr | 0, nf, N, K,
    nsteps: setup.nsteps | 0,
    opsData, off, stLen,
    layout: { SC, ST_CVEL, ST_D, ST_U, ST_UH },
    A, scratchNames, c,
    LS0: setup.ls - 1, JS0: setup.js - 1,
    LP0: setup.lp - 1, JP0: setup.jp - 1,
    TIPI0: setup.tipi - 1,
    NLR0: setup.nlr0 - 1,
    HOLDSTEPS: setup.holdSteps | 0,
  };
}

export function buildKernel(setup, opts = {}) {
  const ws = opts.ws ?? 256;
  const chunk = opts.chunk ?? 2048;

  const nrt = setup.nrt | 0;
  const nth = setup.nth | 0;
  const Nr = setup.Nr | 0;
  const nf = setup.nfree | 0;
  const N = nrt * nth;
  const K = setup.K | 0;
  const nsteps = setup.nsteps | 0;

  // ---- ops buffer packing ----
  const o = setup.ops;
  const parts = [
    ['D1', o.D1, nrt * nrt],
    ['D2', o.D2, nrt * nrt],
    ['F', o.F, nth * nth],
    ['GI', o.Gi, nth * nth],
    ['RD', o.Rd, nrt],
    ['R2D', o.R2d, nrt],
    ['AIB', o.AiB, nrt * nrt * nth],
    ['AIC', o.AiC, nrt * nrt * nth],
    ['AI', o.Ai, nrt * nrt * nth],
    ['FI', o.Fi, nrt * nrt * nth],
    ['ASIB', o.ASiB, nf * nf],
    ['ASIC', o.ASiC, nf * nf],
    ['VM', o.vM, nf],
    ['VR', o.vR, nf],
    ['VK', o.vK, nf],
    ['TIPCOL', o.tipcol, nf],
  ];
  // Launch-step thresholds: smallest 0-based kk with (kk+1)*ht >= tlaunch,
  // decided in f64 exactly as MATLAB's k*ht >= tlaunch decides it.
  const ht = setup.ht;
  const kgo = setup.tlaunch.map((tl) => {
    let g = Math.max(0, Math.ceil(tl / ht) - 1);
    while (g > 0 && g * ht >= tl) g -= 1;
    while ((g + 1) * ht < tl) g += 1;
    return g;
  });
  parts.push(['KGO', Float32Array.from(kgo), K]);
  parts.push(['VSTR', Float32Array.from(setup.vstr), K]);

  const off = {};
  let cursor = 0;
  for (const [name, arr, len] of parts) {
    if (arr.length !== len) throw new Error(`ops part ${name}: ${arr.length} != ${len}`);
    off[name] = cursor;
    cursor += len;
  }
  const opsData = new Float32Array(cursor);
  for (const [name, arr] of parts) opsData.set(arr, off[name]);

  // ---- state buffer layout ----
  const SC = 0; // 16 scalars
  const ST_CVEL = 16;
  const ST_D = ST_CVEL + K;
  const ST_U = ST_D + 3 * nf;
  const ST_UH = ST_U + 3 * N;
  const scratchNames = [
    'UR', 'URR', 'V1', 'UT1', 'V2', 'RHSA', 'RHSAF', 'PH3',
    'PHI', 'PR', 'PRR', 'W1', 'PHT1', 'NL', 'NLF',
  ];
  const A = {};
  let scur = ST_UH + 3 * N;
  for (const n of scratchNames) {
    A[n] = scur;
    scur += N;
  }
  const stLen = scur;

  // ---- baked constants ----
  const LS0 = setup.ls - 1, JS0 = setup.js - 1;
  const LP0 = setup.lp - 1, JP0 = setup.jp - 1;
  const TIPI0 = setup.tipi - 1;
  const NLR0 = setup.nlr0 - 1; // 0-based first kept NL row
  const c = {
    HT: ht,
    HT2I: 1 / (ht * ht),
    GAP: setup.gap,
    KM: setup.KM,
    ALM: setup.alM,
    PS: setup.ps,
    AH: setup.aH,
    AHR: setup.aHr,
    MISSD: setup.missDepth,
    EHRS: setup.Eh / setup.Rs,
    EH2: setup.Eh / 2,
    RSI: 1 / setup.Rs,
    SNL: (ht * ht) / setup.rhoh,
    FSC: setup.fsc,
    C1R: 1 / (2 * setup.hr),
    C2R: 1 / (setup.hr * setup.hr),
    C1T: 1 / (2 * setup.hth),
    C2T: 1 / (setup.hth * setup.hth),
  };
  const constLines = Object.entries(c)
    .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
    .join('\n');

  const code = /* wgsl */ `
const NRT = ${nrt}u;
const NTH = ${nth}u;
const N = ${N}u;
const NF = ${nf}u;
const WS = ${ws}u;
const CHUNK = ${chunk}u;
const NSTEPS = ${nsteps}u;
const KS = ${K}u;
const NRU = ${Nr}u;
const LS0 = ${LS0}u; const JS0 = ${JS0}u;
const LP0 = ${LP0}u; const JP0 = ${JP0}u;
const TIPI = ${TIPI0}u;
const NLR0 = ${NLR0}u;
const HOLDSTEPS = ${setup.holdSteps | 0};
${constLines}
${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')}
const SC = ${SC}u;
const ST_CVEL = ${ST_CVEL}u;
const ST_D = ${ST_D}u;
const ST_U = ${ST_U}u;
const ST_UH = ${ST_UH}u;
${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>;

var<workgroup> wgKbase: u32;
var<workgroup> wgF: f32;
var<workgroup> wgHA: f32;
var<workgroup> wgHB: f32;
var<workgroup> wgHC: f32;
var<workgroup> wgShell: u32;
var<workgroup> wgRot: u32;

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

@compute @workgroup_size(${ws})
fn main(@builtin(local_invocation_index) lid: u32) {
  var shellA: u32 = 0u; var inC: u32 = 0u; var phase: u32 = 0u;
  var strokeIdx: u32 = 0u; var holdLeft: i32 = 0; var vtar: f32 = 0.0;
  var vH: f32 = 0.0; var uHm: f32 = 0.0; var uH0: f32 = 0.0; var uHp: f32 = 0.0;
  var peakXi: f32 = 0.0; var peakDisp: f32 = 0.0;
  if (lid == 0u) {
    wgKbase = u32(st[SC + 0u]);
    shellA = u32(st[SC + 1u]); inC = u32(st[SC + 2u]); phase = u32(st[SC + 3u]);
    strokeIdx = u32(st[SC + 4u]); holdLeft = i32(st[SC + 5u]); vtar = st[SC + 6u];
    vH = st[SC + 7u]; uHm = st[SC + 8u]; uH0 = st[SC + 9u]; uHp = st[SC + 10u];
    peakXi = st[SC + 12u]; peakDisp = st[SC + 13u];
    wgRot = u32(st[SC + 11u]);
    wgShell = shellA;
  }
  let kbase = workgroupUniformLoad(&wgKbase);

  for (var s: u32 = 0u; s < CHUNK; s = s + 1u) {
    let k = kbase + s;
    if (k >= NSTEPS) { break; }
    let dNew = ST_D + (k % 3u) * NF;
    let dCur = ST_D + ((k + 2u) % 3u) * NF;
    let dOld = ST_D + ((k + 1u) % 3u) * NF;

    // ---- stage 1: contact force + hand state machine (lane 0) ----
    if (lid == 0u) {
      let xi = st[ST_U + wgRot * N + 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;
      wgShell = shellA;
    }
    workgroupBarrier();
    storageBarrier();
    let fM1 = wgF;
    let hA = wgHA; let hB = wgHB; let hC = wgHC;

    // ---- stage 2: beam update ----
    for (var r: u32 = lid; r < NF; r = r + WS) {
      var sum: f32 = ops[OP_VM + r] * hA + ops[OP_VR + r] * hB + ops[OP_VK + r] * hC
        + (HT * HT * fM1) * ops[OP_TIPCOL + r];
      for (var p: u32 = 0u; p < NF; p = p + 1u) {
        sum = sum + ops[OP_ASIB + r + p * NF] * st[dCur + p]
                  + ops[OP_ASIC + r + p * NF] * st[dOld + p];
      }
      st[dNew + r] = sum;
    }

    let sh = workgroupUniformLoad(&wgShell);
    if (sh != 0u) {
      let rot = workgroupUniformLoad(&wgRot);
      let uS = ST_U + rot * N;
      let u1S = ST_U + ((rot + 2u) % 3u) * N;
      let unS = ST_U + ((rot + 1u) % 3u) * N;
      let h3S = ST_UH + rot * N;
      let h13S = ST_UH + ((rot + 2u) % 3u) * N;
      let h3nS = ST_UH + ((rot + 1u) % 3u) * N;

      // ---- A: radial/azimuthal derivatives of U ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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;
        }
        st[A_UR + i] = ur;
        st[A_URR + i] = urr;
        let ujm = st[uS + l + thm(j) * NRT];
        let uj0 = st[uS + l + j * NRT];
        let ujp = st[uS + l + thp(j) * NRT];
        st[A_V1 + i] = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T);
        st[A_UT1 + i] = (ujp - ujm) * C1T;
      }
      storageBarrier();

      // ---- B: bracket pieces and Airy right-hand side ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        let l = i % NRT; let j = i / NRT;
        let urt = (st[A_UR + l + thp(j) * NRT] - st[A_UR + l + thm(j) * NRT]) * C1T;
        let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * st[A_UT1 + i];
        st[A_V2 + i] = v2;
        var rv: f32 = 0.0;
        if (l >= 3u && l <= NRU - 2u) {
          let urr = st[A_URR + i]; let v1 = st[A_V1 + i];
          rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2);
        }
        st[A_RHSA + i] = rv;
      }
      storageBarrier();

      // ---- C: transform the Airy rhs ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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[A_RHSA + l + q * NRT] * ops[OP_F + q + j * NTH];
        }
        st[A_RHSAF + i] = sum;
      }
      storageBarrier();

      // ---- D: Airy solve per wavenumber column ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        let l = i % NRT; let j = i / NRT;
        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;
      }
      storageBarrier();

      // ---- E: stress function back to physical space ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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[A_PH3 + l + q * NRT] * ops[OP_GI + q + j * NTH];
        }
        st[A_PHI + i] = sum;
      }
      storageBarrier();

      // ---- F: derivatives of the stress function ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        let l = i % NRT; let j = i / NRT;
        var pr: f32 = 0.0; var prr: f32 = 0.0;
        if (l >= 1u && l <= NRT - 2u) {
          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];
          pr = (pp - pm) * C1R;
          prr = (pm - 2.0 * p0 + pp) * C2R;
        }
        st[A_PR + i] = pr;
        st[A_PRR + i] = prr;
        let pjm = st[A_PHI + l + thm(j) * NRT];
        let pj0 = st[A_PHI + l + j * NRT];
        let pjp = st[A_PHI + l + thp(j) * NRT];
        st[A_W1 + i] = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * pj0 + pjp) * C2T);
        st[A_PHT1 + i] = (pjp - pjm) * C1T;
      }
      storageBarrier();

      // ---- G: von Karman bracket, forcing, membrane term ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        let l = i % NRT; let j = i / NRT;
        var nl: f32 = 0.0;
        if (l >= NLR0 && l < NRU) {
          let prt = (st[A_PR + l + thp(j) * NRT] - st[A_PR + l + thm(j) * NRT]) * C1T;
          let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * st[A_PHT1 + i];
          let prr = st[A_PRR + i]; let w1 = st[A_W1 + i];
          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 - fM1 * FSC; }
          nl = nl * SNL;
        }
        st[A_NL + i] = nl;
      }
      storageBarrier();

      // ---- H: transform the nonlinear load ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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[A_NL + l + q * NRT] * ops[OP_F + q + j * NTH];
        }
        st[A_NLF + i] = sum;
      }
      storageBarrier();

      // ---- I: implicit shell update per wavenumber column ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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) {
          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;
      }
      storageBarrier();

      // ---- J: back to physical space; pickup acceleration ----
      for (var i: u32 = lid; i < N; i = i + WS) {
        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[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;
        }
      }
      storageBarrier();
      if (lid == 0u) {
        wgRot = (rot + 1u) % 3u;
        if ((k + 1u) % 512u == 0u) {
          var pd: f32 = 0.0;
          for (var l2: u32 = 0u; l2 < NRT; l2 = l2 + 1u) {
            pd = max(pd, abs(st[unS + l2 + JS0 * NRT]));
          }
          if (pd > peakDisp) { peakDisp = pd; }
        }
      }
    }
    workgroupBarrier();
    storageBarrier();
  }

  if (lid == 0u) {
    st[SC + 0u] = f32(min(kbase + CHUNK, NSTEPS));
    st[SC + 1u] = f32(shellA); st[SC + 2u] = f32(inC); st[SC + 3u] = f32(phase);
    st[SC + 4u] = f32(strokeIdx); st[SC + 5u] = f32(holdLeft); st[SC + 6u] = vtar;
    st[SC + 7u] = vH; st[SC + 8u] = uHm; st[SC + 9u] = uH0; st[SC + 10u] = uHp;
    st[SC + 11u] = f32(wgRot); st[SC + 12u] = peakXi; st[SC + 13u] = peakDisp;
  }
}
`;

  return { code, opsData, stLen, N, chunk, ws, off, layout: { SC, ST_CVEL, ST_D, ST_U, ST_UH } };
}