// WGSL kernels of the time loop. Three entry points: the element pass (forces, sigma1 term,
// contact forces of the hits, decimation accumulators of the radiated acceleration), the node
// pass (gather, contact, lagged rigid-mode projection, leapfrog update, partial sums, exciter
// update) and the gather pass (retarded-time sum at the listener once per audio sample). Sizes
// and constants are baked into the source. See THEORY.md for the buffer layouts and the ordering
// argument behind the two step counters.

export interface ShaderConstants {
  N: number; E: number; K: number; P: number; M: number; W: number; D: number; G: number;
  dt: number; sigma0: number; nl: number;
}

export const ELEM_STRIDE = 96;     // Gx[3] Gy[3] A z0x z0y C11 C12 C33 cd Ke[81] pad[2]
export const NODE_STRIDE = 40;     // Minv[5] M[5] Rm[6][5]
export const HIT_HEADER = 12;      // kPoint alpha mass gap velocity zs vs F Fmax pad[3]
export const HITU_HEADER = 4;      // releaseStep ncont pad[2]
export const SYNC_HEADER = 4;      // stepA stepB pad[2], then partials[2][G][12]
export const WORKGROUP_NODES = 256;

function f(x: number): string {
  const s = x.toPrecision(9);
  return /[.e]/.test(s) ? s : s + ".0";
}

export function shaderSource(c: ShaderConstants): string {
  const HF = HIT_HEADER + 5 * c.P, HU = HITU_HEADER + 3 * c.P;
  return /* wgsl */ `
const N: u32 = ${c.N}u; const E: u32 = ${c.E}u; const K: u32 = ${c.K}u; const P: u32 = ${c.P}u;
const M: u32 = ${c.M}u; const W: u32 = ${c.W}u; const D: u32 = ${c.D}u; const G: u32 = ${c.G}u;
const DT: f32 = ${f(c.dt)}; const DAMP0: f32 = ${f(1 / (1 + 2 * c.sigma0 * c.dt))}; const NL: f32 = ${f(c.nl)};
const EF: u32 = ${ELEM_STRIDE}u; const NF: u32 = ${NODE_STRIDE}u; const HF: u32 = ${HF}u; const HU: u32 = ${HU}u;
const OFF_V: u32 = 5u * N; const OFF_ACC: u32 = 10u * N; const OFF_VWP: u32 = 11u * N;
const OFF_HIST: u32 = E * W; const OFF_TAPS: u32 = EF * E; const OFF_PART: u32 = ${SYNC_HEADER}u;

struct Chunk { base: i32, size: u32, frameBase: u32, every: u32 }

// ------------------------------------------------------------------ element pass
@group(0) @binding(0) var<storage, read> elemF: array<f32>;
@group(0) @binding(1) var<storage, read> elemU: array<u32>;
@group(0) @binding(2) var<storage, read> stateR: array<f32>;
@group(0) @binding(3) var<storage, read_write> elemForce: array<f32>;
@group(0) @binding(4) var<storage, read_write> hitsF: array<f32>;
@group(0) @binding(5) var<storage, read> hitsUR: array<u32>;
@group(0) @binding(6) var<storage, read_write> syncE: array<f32>;
@group(0) @binding(7) var<storage, read_write> rad: array<f32>;

@compute @workgroup_size(64)
fn elementPass(@builtin(global_invocation_id) gid: vec3<u32>) {
  let id = gid.x;
  let step = bitcast<u32>(syncE[0]);
  if (id == 0u) { syncE[1] = syncE[0]; }
  if (id < E) {
    let e = id;
    let b = e * EF;
    let n0 = elemU[3u * e]; let n1 = elemU[3u * e + 1u]; let n2 = elemU[3u * e + 2u];
    var qe: array<f32, 15>;
    for (var d = 0u; d < 5u; d++) { qe[d] = stateR[5u * n0 + d]; qe[5u + d] = stateR[5u * n1 + d]; qe[10u + d] = stateR[5u * n2 + d]; }
    let gx = vec3<f32>(elemF[b], elemF[b + 1u], elemF[b + 2u]);
    let gy = vec3<f32>(elemF[b + 3u], elemF[b + 4u], elemF[b + 5u]);
    let A = elemF[b + 6u]; let z0x = elemF[b + 7u]; let z0y = elemF[b + 8u];
    let C11 = elemF[b + 9u]; let C12 = elemF[b + 10u]; let C33 = elemF[b + 11u]; let cd = elemF[b + 12u];
    let uu = vec3<f32>(qe[0], qe[5], qe[10]); let vv = vec3<f32>(qe[1], qe[6], qe[11]); let ww = vec3<f32>(qe[2], qe[7], qe[12]);
    let ux = dot(gx, uu); let uy = dot(gy, uu); let vx = dot(gx, vv); let vy = dot(gy, vv); let wx = dot(gx, ww); let wy = dot(gy, ww);
    let ax = NL * wx + z0x; let ay = NL * wy + z0y;
    let exx = ux + NL * 0.5 * wx * wx + z0x * wx;
    let eyy = vy + NL * 0.5 * wy * wy + z0y * wy;
    let gxy = uy + vx + NL * wx * wy + z0x * wy + z0y * wx;
    let Nxx = C11 * exx + C12 * eyy; let Nyy = C12 * exx + C11 * eyy; let Nxy = C33 * gxy;
    let vw = vec3<f32>(stateR[OFF_V + 5u * n0 + 2u], stateR[OFF_V + 5u * n1 + 2u], stateR[OFF_V + 5u * n2 + 2u]);
    let vwx = dot(gx, vw); let vwy = dot(gy, vw);
    var out: array<f32, 15>;
    for (var i = 0u; i < 3u; i++) {
      let g_x = gx[i]; let g_y = gy[i];
      out[5u * i] = -A * (Nxx * g_x + Nxy * g_y);
      out[5u * i + 1u] = -A * (Nyy * g_y + Nxy * g_x);
      out[5u * i + 2u] = -A * (Nxx * ax * g_x + Nyy * ay * g_y + Nxy * (ay * g_x + ax * g_y)) - cd * (g_x * vwx + g_y * vwy);
      out[5u * i + 3u] = 0.0; out[5u * i + 4u] = 0.0;
    }
    var U: array<f32, 9>;
    for (var i = 0u; i < 3u; i++) { U[3u * i] = qe[5u * i + 2u]; U[3u * i + 1u] = qe[5u * i + 3u]; U[3u * i + 2u] = qe[5u * i + 4u]; }
    for (var i = 0u; i < 9u; i++) {
      var s = 0.0;
      for (var j = 0u; j < 9u; j++) { s += elemF[b + 13u + 9u * i + j] * U[j]; }
      let slot = 5u * (i / 3u) + 2u + (i % 3u);
      out[slot] -= s;
    }
    for (var i = 0u; i < 15u; i++) { elemForce[15u * e + i] = out[i]; }
    // decimation accumulators of the radiated acceleration (of the previous step)
    if (step >= 1u) {
      let s = step - 1u; let j0 = s / M; let kk = s - j0 * M;
      if (kk == 0u && j0 > 0u) {
        let j = i32(j0) - i32(W / 2u);
        if (j >= 0) { rad[OFF_HIST + e * D + (u32(j) % D)] = rad[e * W]; }
        for (var w = 0u; w < W - 1u; w++) { rad[e * W + w] = rad[e * W + w + 1u]; }
        rad[e * W + W - 1u] = 0.0;
      }
      let a = (stateR[OFF_ACC + n0] + stateR[OFF_ACC + n1] + stateR[OFF_ACC + n2]) / 3.0;
      for (var w = 0u; w < W; w++) { rad[e * W + w] += a * elemF[OFF_TAPS + w * M + kk]; }
    }
  } else if (id < E + K) {
    // contact forces of hit k at this step
    let k = id - E; let hf = k * HF; let hu = k * HU;
    var wp: array<f32, ${c.P}>;
    var wmax = -1.0e30;
    for (var p = 0u; p < P; p++) {
      var w = hitsF[hf + ${HIT_HEADER}u + 3u * P + p];
      for (var i = 0u; i < 3u; i++) { w += hitsF[hf + ${HIT_HEADER}u + 3u * p + i] * stateR[5u * hitsUR[hu + ${HITU_HEADER}u + 3u * p + i] + 2u]; }
      wp[p] = w; wmax = max(wmax, w);
    }
    var zs = hitsF[hf + 5u];
    if (step == hitsUR[hu]) { zs = wmax + hitsF[hf + 3u]; hitsF[hf + 5u] = zs; hitsF[hf + 6u] = -hitsF[hf + 4u]; }
    let isActive = step >= hitsUR[hu];
    let kPoint = hitsF[hf]; let alpha = hitsF[hf + 1u];
    var F = 0.0;
    for (var p = 0u; p < P; p++) {
      let delta = max(wp[p] - zs, 0.0);
      var fp = 0.0;
      if (isActive && delta > 0.0) { fp = kPoint * pow(delta, alpha); }
      hitsF[hf + ${HIT_HEADER}u + 4u * P + p] = fp; F += fp;
    }
    hitsF[hf + 7u] = F;
  }
}

// ------------------------------------------------------------------ node pass
@group(1) @binding(0) var<storage, read_write> state: array<f32>;
@group(1) @binding(1) var<storage, read> elemForceR: array<f32>;
@group(1) @binding(2) var<storage, read> nodeF: array<f32>;
@group(1) @binding(3) var<storage, read> adj: array<u32>;
@group(1) @binding(4) var<storage, read_write> hitsFN: array<f32>;
@group(1) @binding(5) var<storage, read_write> hitsU: array<u32>;
@group(1) @binding(6) var<storage, read_write> sync: array<f32>;

var<workgroup> red: array<f32, ${12 * WORKGROUP_NODES}>;

@compute @workgroup_size(${WORKGROUP_NODES})
fn nodePass(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
  let nd = gid.x; let t = lid.x;
  let step = bitcast<u32>(sync[1]);
  let parity = step & 1u; let prevp = (step + 1u) & 1u;
  var c: array<f32, 12>;
  for (var k = 0u; k < 12u; k++) {
    var s = 0.0;
    for (var g = 0u; g < G; g++) { s += sync[OFF_PART + (prevp * G + g) * 12u + k]; }
    c[k] = s;
  }
  var part: array<f32, 12>;
  for (var k = 0u; k < 12u; k++) { part[k] = 0.0; }
  if (nd < N) {
    var fo: array<f32, 5>;
    for (var d = 0u; d < 5u; d++) { fo[d] = 0.0; }
    let a0 = adj[nd]; let a1 = adj[nd + 1u];
    for (var a = a0; a < a1; a++) {
      let ent = adj[N + 1u + a]; let e = ent / 3u; let i = ent - 3u * e;
      for (var d = 0u; d < 5u; d++) { fo[d] += elemForceR[15u * e + 5u * i + d]; }
    }
    for (var k = 0u; k < K; k++) {
      for (var p = 0u; p < P; p++) {
        for (var i = 0u; i < 3u; i++) {
          if (hitsU[k * HU + ${HITU_HEADER}u + 3u * p + i] == nd) {
            fo[2] -= hitsFN[k * HF + ${HIT_HEADER}u + 4u * P + p] * hitsFN[k * HF + ${HIT_HEADER}u + 3u * p + i];
          }
        }
      }
    }
    let nb = nd * NF;
    for (var d = 0u; d < 5u; d++) {
      let idx = 5u * nd + d;
      var vd = state[OFF_V + idx]; var qd = state[idx];
      for (var k = 0u; k < 6u; k++) { let r = nodeF[nb + 10u + 5u * k + d]; vd -= c[k] * r; qd -= c[6u + k] * r; }
      let v1 = (vd + DT * fo[d] * nodeF[nb + d]) * DAMP0;
      if (d == 2u) { state[OFF_ACC + nd] = (vd - state[OFF_VWP + nd]) / DT; state[OFF_VWP + nd] = vd; }
      let q1 = qd + DT * v1;
      state[idx] = q1; state[OFF_V + idx] = v1;
      let mm = nodeF[nb + 5u + d];
      for (var k = 0u; k < 6u; k++) { let r = nodeF[nb + 10u + 5u * k + d] * mm; part[k] += r * v1; part[6u + k] += r * q1; }
    }
  }
  for (var k = 0u; k < 12u; k++) { red[t * 12u + k] = part[k]; }
  workgroupBarrier();
  for (var s = ${WORKGROUP_NODES / 2}u; s > 0u; s = s >> 1u) {
    if (t < s) { for (var k = 0u; k < 12u; k++) { red[t * 12u + k] += red[(t + s) * 12u + k]; } }
    workgroupBarrier();
  }
  if (t == 0u) { for (var k = 0u; k < 12u; k++) { sync[OFF_PART + (parity * G + wid.x) * 12u + k] = red[k]; } }
  if (nd == 0u) {
    for (var k = 0u; k < K; k++) {
      let hf = k * HF; let hu = k * HU;
      let F = hitsFN[hf + 7u];
      let vs = hitsFN[hf + 6u] + DT * F / hitsFN[hf + 2u];
      hitsFN[hf + 6u] = vs; hitsFN[hf + 5u] = hitsFN[hf + 5u] + DT * vs;
      hitsFN[hf + 8u] = max(hitsFN[hf + 8u], F);
      if (F > 0.0) { hitsU[hu + 1u] = hitsU[hu + 1u] + 1u; }
    }
    sync[0] = bitcast<f32>(step + 1u);
  }
}

// ------------------------------------------------------------------ gather pass (listener)
@group(2) @binding(0) var<storage, read> syncG: array<f32>;
@group(2) @binding(1) var<storage, read> radR: array<f32>;
@group(2) @binding(2) var<storage, read> radCoef: array<f32>;
@group(2) @binding(3) var<storage, read_write> pressure: array<f32>;
@group(2) @binding(4) var<uniform> chunk: Chunk;

var<workgroup> redF: array<f32, 256>;
var<workgroup> redB: array<f32, 256>;

@compute @workgroup_size(256)
fn gatherPass(@builtin(local_invocation_id) lid: vec3<u32>) {
  let t = lid.x;
  let step = bitcast<u32>(syncG[0]);
  let j = i32((step - 1u) / M) - i32(W / 2u);
  var fs = 0.0; var bs = 0.0;
  if (j >= 0) {
    for (var e = t; e < E; e += 256u) {
      let rc = 6u * e;
      let jf = j - i32(radCoef[rc + 2u]); let jb = j - i32(radCoef[rc + 4u]);
      if (jf >= 1) {
        let fr = radCoef[rc + 3u];
        fs += radCoef[rc] * (radR[OFF_HIST + e * D + (u32(jf) % D)] * (1.0 - fr) + radR[OFF_HIST + e * D + (u32(jf - 1) % D)] * fr);
      }
      if (jb >= 1) {
        let fr = radCoef[rc + 5u];
        bs += radCoef[rc + 1u] * (radR[OFF_HIST + e * D + (u32(jb) % D)] * (1.0 - fr) + radR[OFF_HIST + e * D + (u32(jb - 1) % D)] * fr);
      }
    }
  }
  redF[t] = fs; redB[t] = bs;
  workgroupBarrier();
  for (var s = 128u; s > 0u; s = s >> 1u) {
    if (t < s) { redF[t] += redF[t + s]; redB[t] += redB[t + s]; }
    workgroupBarrier();
  }
  if (t == 0u && j >= 0) {
    let idx = j - chunk.base;
    if (idx >= 0 && u32(idx) < chunk.size) { pressure[u32(idx)] = redF[0]; pressure[chunk.size + u32(idx)] = redB[0]; }
  }
}

// ------------------------------------------------------------------ capture pass (displacement movie)
@group(3) @binding(0) var<storage, read> stateC: array<f32>;
@group(3) @binding(1) var<storage, read_write> movie: array<f32>;
@group(3) @binding(2) var<uniform> chunkC: Chunk;
@group(3) @binding(3) var<storage, read> syncC: array<f32>;

@compute @workgroup_size(256)
fn capturePass(@builtin(global_invocation_id) gid: vec3<u32>) {
  let nd = gid.x;
  if (nd >= N) { return; }
  let step = bitcast<u32>(syncC[0]);
  let frame = (step / M) / chunkC.every - chunkC.frameBase;
  movie[frame * N + nd] = stateC[5u * nd + 2u];
}
`;
}