// Bead solver on the GPU: the same scheme as solver.ts (velocity Verlet or
// its Yoshida composition, with exact damping applied once per audio sample),
// in single precision.
//
// Each phase of a step is its own dispatch over beads or elements; the host
// encodes the phases for a chunk of audio samples into one command buffer.
// Per-phase constants (kick and drift lengths, Taylor coefficients) come from
// a table selected with a dynamic uniform offset, and the current sample
// index lives in a GPU buffer, which the record phase advances.
//
// To keep single precision accurate, the state is the displacement u from
// the rest positions, never the absolute position, and every element works
// from its rest geometry (computed in double precision on the CPU) plus
// displacement differences. Element forces are written to per-element slots
// and gathered per bead through an incidence list, since there are no atomic
// float additions.

const WG: u32 = 64u;

struct Hdr {
  n: u32, ns: u32, nb: u32, nh: u32,
  nd: u32, nl: u32, quantity: u32, npulse: u32,
  sample0: u32, frameStride: u32, frameSize: u32, nImp: u32,
  impOff: u32, oFrames: u32, pad0: u32, pad1: u32,
  // offsets into P (f32 parameters)
  oInvm: u32, oMass: u32, oFric: u32, oS: u32,
  oB: u32, oH: u32, oD: u32, oPulse: u32,
  // offsets into I (u32 indices)
  iS: u32, iB: u32, iH: u32, iD: u32,
  iFr: u32, iFe: u32, iDr: u32, iDe: u32,
  iPr: u32, iPe: u32, iL: u32, iImp: u32,
  // offsets into E for the bending and hinge slots (forces / damping)
  eFb: u32, eFh: u32, eDb: u32, eDh: u32,
  Delta: f32, dirx: f32, diry: f32, dirz: f32,
}

// per-phase constants
struct Ph {
  kick: f32, drift: f32, tOff: f32, gather: u32,
  c: f32, saveAv: u32, last: u32, pad: u32,
}

@group(0) @binding(0) var<uniform> H: Hdr;
@group(0) @binding(1) var<storage, read> P: array<f32>;
@group(0) @binding(2) var<storage, read> I: array<u32>;
@group(0) @binding(3) var<storage, read_write> S: array<f32>; // u | v | F | T | Av | KE0 | diss | theta0
@group(0) @binding(4) var<storage, read_write> E: array<f32>; // element slots
@group(0) @binding(5) var<storage, read_write> O: array<f32>; // audio | frames
@group(0) @binding(6) var<storage, read_write> C: array<u32>; // [current sample]
@group(1) @binding(0) var<uniform> Q: Ph;

fn offV() -> u32 { return 3u * H.n; }
fn offF() -> u32 { return 6u * H.n; }
fn offT() -> u32 { return 9u * H.n; }
fn offAv() -> u32 { return 12u * H.n; }
fn offKE() -> u32 { return 15u * H.n; }
fn offDiss() -> u32 { return 16u * H.n; }
fn offTh0() -> u32 { return 17u * H.n; }

fn S3(o: u32) -> vec3<f32> { return vec3<f32>(S[o], S[o + 1u], S[o + 2u]); }
fn setS3(o: u32, w: vec3<f32>) { S[o] = w.x; S[o + 1u] = w.y; S[o + 2u] = w.z; }
fn setE3(o: u32, w: vec3<f32>) { E[o] = w.x; E[o + 1u] = w.y; E[o + 2u] = w.z; }
fn E3(o: u32) -> vec3<f32> { return vec3<f32>(E[o], E[o + 1u], E[o + 2u]); }
fn P3(o: u32) -> vec3<f32> { return vec3<f32>(P[o], P[o + 1u], P[o + 2u]); }
fn U(b: u32) -> vec3<f32> { return S3(3u * b); }
fn V(b: u32) -> vec3<f32> { return S3(offV() + 3u * b); }

// Hinge h: angle and gradient from the rest vectors relative to bead j plus
// displacement differences.
struct HingeG { theta: f32, g: array<vec3<f32>, 4> }

fn hingeAt(h: u32) -> HingeG {
  let bi = I[H.iH + 4u * h];
  let bj = I[H.iH + 4u * h + 1u];
  let bk = I[H.iH + 4u * h + 2u];
  let bl = I[H.iH + 4u * h + 3u];
  let o = H.oH + 16u * h;
  let uj = U(bj);
  let pa = P3(o) + (U(bi) - uj); // bead i
  let pk = P3(o + 3u) + (U(bk) - uj); // bead k
  let pl = P3(o + 6u) + (U(bl) - uj); // bead l; bead j is the origin
  let e = pk;
  let eLen = length(e);
  let b = pa - pk;
  let c = pl - pk;
  let n1 = cross(pa, b);
  let n2 = cross(c, pl);
  let n1sq = dot(n1, n1);
  let n2sq = dot(n2, n2);
  var r: HingeG;
  r.theta = -atan2(dot(cross(n1, n2), e) / eLen, dot(n1, n2));
  let p1 = dot(pa, e) / eLen;
  let q1 = dot(b, e) / eLen;
  let p2 = dot(pl, e) / eLen;
  let q2 = dot(c, e) / eLen;
  let m1 = 1.0 / n1sq;
  let m2 = 1.0 / n2sq;
  r.g[0] = (eLen * m1) * n1;
  r.g[1] = q1 * m1 * n1 + q2 * m2 * n2;
  r.g[2] = -p1 * m1 * n1 - p2 * m2 * n2;
  r.g[3] = (eLen * m2) * n2;
  return r;
}

fn wrapAngle(d: f32) -> f32 {
  var x = d;
  if (x > 3.14159265) { x -= 6.28318531; }
  if (x <= -3.14159265) { x += 6.28318531; }
  return x;
}

// ---- forces ---------------------------------------------------------------------

// one invocation per element: springs, then bending triplets, then hinges
@compute @workgroup_size(WG)
fn elements(@builtin(global_invocation_id) gid: vec3<u32>) {
  var q = gid.x;
  if (q < H.ns) {
    let o = H.oS + 8u * q;
    let r0 = P3(o);
    let du = U(I[H.iS + 2u * q + 1u]) - U(I[H.iS + 2u * q]);
    let r = r0 + du;
    let L = length(r);
    // L - |r0| without cancellation, plus the rest stretch |r0| - L0
    let e = P[o + 4u] + (2.0 * dot(r0, du) + dot(du, du)) / (L + P[o + 3u]);
    let f = P[o + 5u] * e + P[o + 6u] * e * e * e;
    setE3(3u * q, (f / L) * r);
    return;
  }
  q -= H.ns;
  if (q < H.nb) {
    let o = H.oB + 8u * q;
    let D = P3(o) + U(I[H.iB + 3u * q]) - 2.0 * U(I[H.iB + 3u * q + 1u]) + U(I[H.iB + 3u * q + 2u]);
    setE3(H.eFb + 3u * q, P[o + 3u] * D);
    return;
  }
  q -= H.nb;
  if (q < H.nh) {
    let hg = hingeAt(q);
    let s = P[H.oH + 16u * q + 9u] * wrapAngle(hg.theta - S[offTh0() + q]);
    for (var p = 0u; p < 4u; p++) { setE3(H.eFh + 12u * q + 3u * p, s * hg.g[p]); }
  }
}

fn pulseForce(b: u32, t: f32) -> vec3<f32> {
  var f = vec3<f32>(0.0);
  for (var q = I[H.iPr + b]; q < I[H.iPr + b + 1u]; q++) {
    let o = H.oPulse + 8u * I[H.iPe + q];
    let x = (t - P[o]) / P[o + 1u];
    if (x > 0.0 && x < 1.0) { f += (0.5 * (1.0 - cos(6.28318531 * x))) * P3(o + 2u); }
  }
  return f;
}

// gather (optional), kick and drift; time t = (sample index) * Delta + tOff
@compute @workgroup_size(WG)
fn beads(@builtin(global_invocation_id) gid: vec3<u32>) {
  let b = gid.x;
  if (b >= H.n) { return; }
  var F = S3(offF() + 3u * b);
  if (Q.gather == 1u) {
    F = vec3<f32>(0.0);
    for (var q = I[H.iFr + b]; q < I[H.iFr + b + 1u]; q++) {
      F += bitcast<f32>(I[H.iFe + 2u * q + 1u]) * E3(I[H.iFe + 2u * q]);
    }
    if (H.npulse > 0u) { F += pulseForce(b, f32(C[0]) * H.Delta + Q.tOff); }
    setS3(offF() + 3u * b, F);
  }
  let v = V(b) + (Q.kick * P[H.oInvm + b]) * F;
  setS3(offV() + 3u * b, v);
  setS3(3u * b, U(b) + Q.drift * v);
}

// ---- sample boundary ----------------------------------------------------------------

// frame copy (if this sample is a frame), T = v, kinetic energy before damping
@compute @workgroup_size(WG)
fn pre(@builtin(global_invocation_id) gid: vec3<u32>) {
  let b = gid.x;
  if (b >= H.n) { return; }
  let sample = C[0];
  let v = V(b);
  if (sample % H.frameStride == 0u) {
    let first = ((H.sample0 + H.frameStride - 1u) / H.frameStride) * H.frameStride;
    let o = H.oFrames + ((sample - first) / H.frameStride) * H.frameSize;
    let u = U(b);
    O[o + 3u * b] = u.x;
    O[o + 3u * b + 1u] = u.y;
    O[o + 3u * b + 2u] = u.z;
    O[o + 3u * H.n + 3u * b] = v.x;
    O[o + 3u * H.n + 3u * b + 1u] = v.y;
    O[o + 3u * H.n + 3u * b + 2u] = v.z;
    O[o + 6u * H.n + b] = S[offDiss() + b];
  }
  setS3(offT() + 3u * b, v);
  S[offKE() + b] = 0.5 * P[H.oMass + b] * dot(v, v);
}

// T = v, for each further piece of a damping step split into pieces
@compute @workgroup_size(WG)
fn resetT(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x < H.n) { setS3(offT() + 3u * gid.x, V(gid.x)); }
}

// one Taylor term of exp(-tau A): the damping forces of the elements from T
@compute @workgroup_size(WG)
fn dampElements(@builtin(global_invocation_id) gid: vec3<u32>) {
  var q = gid.x;
  if (q < H.nd) {
    let bi = I[H.iD + 2u * q];
    let bj = I[H.iD + 2u * q + 1u];
    let o = H.oD + 8u * q;
    let rw = S3(offT() + 3u * bi) - S3(offT() + 3u * bj);
    var f = P[o + 3u] * rw;
    if (P[o + 4u] != 0.0) {
      let dir = normalize(P3(o) + U(bj) - U(bi));
      f += (P[o + 4u] * dot(dir, rw)) * dir;
    }
    setE3(3u * q, f);
    return;
  }
  q -= H.nd;
  if (q < H.nb) {
    let cb = P[H.oB + 8u * q + 4u];
    let w = S3(offT() + 3u * I[H.iB + 3u * q]) - 2.0 * S3(offT() + 3u * I[H.iB + 3u * q + 1u]) + S3(offT() + 3u * I[H.iB + 3u * q + 2u]);
    setE3(H.eDb + 3u * q, cb * w);
    return;
  }
  q -= H.nb;
  if (q < H.nh) {
    let ch = P[H.oH + 16u * q + 10u];
    if (ch == 0.0) { return; } // no damping entries reference these slots
    let hg = hingeAt(q);
    var s = 0.0;
    for (var p = 0u; p < 4u; p++) { s += dot(hg.g[p], S3(offT() + 3u * I[H.iH + 4u * q + p])); }
    for (var p = 0u; p < 4u; p++) { setE3(H.eDh + 12u * q + 3u * p, (ch * s) * hg.g[p]); }
  }
}

// A T per bead, then T <- c A T and v <- v + T (c = -tau/k)
@compute @workgroup_size(WG)
fn dampBeads(@builtin(global_invocation_id) gid: vec3<u32>) {
  let b = gid.x;
  if (b >= H.n) { return; }
  var f = P[H.oFric + b] * S3(offT() + 3u * b);
  for (var q = I[H.iDr + b]; q < I[H.iDr + b + 1u]; q++) {
    f += bitcast<f32>(I[H.iDe + 2u * q + 1u]) * E3(I[H.iDe + 2u * q]);
  }
  let a = P[H.oInvm + b] * f;
  if (Q.saveAv == 1u) { setS3(offAv() + 3u * b, a); }
  let t = Q.c * a;
  setS3(offT() + 3u * b, t);
  let v = V(b) + t;
  setS3(offV() + 3u * b, v);
  if (Q.last == 1u) {
    S[offDiss() + b] += S[offKE() + b] - 0.5 * P[H.oMass + b] * dot(v, v);
  }
}

// the listener's sample; advances the sample index
@compute @workgroup_size(1)
fn record() {
  let dir = vec3<f32>(H.dirx, H.diry, H.dirz);
  var val = 0.0;
  for (var q = 0u; q < H.nl; q++) {
    let b = I[H.iL + q];
    if (H.quantity == 0u) { val += dot(U(b), dir); }
    else if (H.quantity == 1u) { val += dot(S3(offT() + 3u * b), dir); } // T holds v at the boundary
    else { val += dot(P[H.oInvm + b] * S3(offF() + 3u * b) - S3(offAv() + 3u * b), dir); }
  }
  O[C[0] - H.sample0] = val / f32(H.nl);
  C[0] = C[0] + 1u;
}

// velocity impulses at the start of a chunk: [bead, vx, vy, vz] per entry
@compute @workgroup_size(WG)
fn impulse(@builtin(global_invocation_id) gid: vec3<u32>) {
  let q = gid.x;
  if (q >= H.nImp) { return; }
  let o = H.iImp + 4u * (H.impOff + q);
  let b = I[o];
  if (P[H.oInvm + b] > 0.0) {
    setS3(offV() + 3u * b, V(b) + vec3<f32>(bitcast<f32>(I[o + 1u]), bitcast<f32>(I[o + 2u]), bitcast<f32>(I[o + 3u])));
  }
}

// rest angles of the hinges (u = 0)
@compute @workgroup_size(WG)
fn initHinges(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x < H.nh) { S[offTh0() + gid.x] = hingeAt(gid.x).theta; }
}