// Rigid-body projection: (1) mass-weighted sums of the six rigid modes for
// displacement and velocity, kinetic energy and max |w| in one workgroup;
// (2) subtraction of the rigid components from every node.

@group(0) @binding(0) var<uniform> P: Params;
@group(0) @binding(1) var<storage, read_write> pos: array<f32>; // [u | v | w | wLo]
@group(0) @binding(2) var<storage, read_write> vel: array<f32>; // [du | dv | dw]
@group(0) @binding(3) var<storage, read_write> stats: array<f32>; // 0..11 coefficients, 12 KE, 13 max|w|

const NS = 14u;
var<workgroup> sh: array<array<f32, NS>, 256>;

@compute @workgroup_size(256)
fn sums(@builtin(local_invocation_id) lid: vec3u) {
  let t = lid.x;
  var s: array<f32, NS>;
  for (var q = 0u; q < NS; q++) { s[q] = 0.0; }
  for (var k = t; k < P.n; k += 256u) {
    let i = i32(k % P.nx); let j = i32(k / P.nx);
    let x = nodeX(i); let y = nodeY(j);
    let m = P.rhoh * P.hx * P.hy * trapWeight(i, j);
    let n = P.n;
    let p = vec3f(pos[k], pos[n + k], pos[2u * n + k]);
    let v = vec3f(vel[k], vel[n + k], vel[2u * n + k]);
    s[0] += m * p.x; s[1] += m * p.y; s[2] += m * p.z;
    s[3] += m * (x * p.y - y * p.x); s[4] += m * y * p.z; s[5] += m * (-x * p.z);
    s[6] += m * v.x; s[7] += m * v.y; s[8] += m * v.z;
    s[9] += m * (x * v.y - y * v.x); s[10] += m * y * v.z; s[11] += m * (-x * v.z);
    s[12] += 0.5 * m * dot(v, v);
    s[13] = max(s[13], abs(p.z));
  }
  for (var q = 0u; q < NS; q++) { sh[t][q] = s[q]; }
  workgroupBarrier();
  for (var stride = 128u; stride > 0u; stride >>= 1u) {
    if (t < stride) {
      for (var q = 0u; q < 13u; q++) { sh[t][q] += sh[t + stride][q]; }
      sh[t][13] = max(sh[t][13], sh[t + stride][13]);
    }
    workgroupBarrier();
  }
  if (t == 0u) {
    let inv = array<f32, 6>(1.0 / P.mass, 1.0 / P.mass, 1.0 / P.mass, 1.0 / P.Iz, 1.0 / P.Ix, 1.0 / P.Iy);
    for (var q = 0u; q < 6u; q++) {
      stats[q] = sh[0][q] * inv[q];
      stats[q + 6u] = sh[0][q + 6u] * inv[q];
    }
    // Kinetic energy of the non-rigid motion (the six rigid modes are mass-orthogonal).
    var keRigid = 0.0;
    for (var q = 6u; q < 12u; q++) { keRigid += 0.5 * sh[0][q] * sh[0][q] * inv[q - 6u]; }
    stats[12] = max(sh[0][12] - keRigid, 0.0);
    stats[13] = sh[0][13];
  }
}

@compute @workgroup_size(256)
fn apply(@builtin(global_invocation_id) gid: vec3u) {
  let k = gid.x;
  if (k >= P.n) { return; }
  let i = i32(k % P.nx); let j = i32(k / P.nx);
  let x = nodeX(i); let y = nodeY(j);
  let n = P.n;
  pos[k] -= stats[0] - stats[3] * y;
  pos[n + k] -= stats[1] + stats[3] * x;
  pos[2u * n + k] -= stats[2] + stats[4] * y - stats[5] * x;
  vel[k] -= stats[6] - stats[9] * y;
  vel[n + k] -= stats[7] + stats[9] * x;
  vel[2u * n + k] -= stats[8] + stats[10] * y - stats[11] * x;
}

// Removes the three rigid components (translation and two rotations) from a
// per-node acceleration field so that it is self-equilibrated. One workgroup.
@group(0) @binding(4) var<storage, read_write> field: array<f32>;
var<workgroup> sh3: array<vec3f, 256>;

@compute @workgroup_size(256)
fn projectField(@builtin(local_invocation_id) lid: vec3u) {
  let t = lid.x;
  var s = vec3f(0.0);
  for (var k = t; k < P.n; k += 256u) {
    let i = i32(k % P.nx); let j = i32(k / P.nx);
    let m = P.rhoh * P.hx * P.hy * trapWeight(i, j);
    let f = field[k];
    s += vec3f(m * f, m * nodeY(j) * f, -m * nodeX(i) * f);
  }
  sh3[t] = s;
  workgroupBarrier();
  for (var stride = 128u; stride > 0u; stride >>= 1u) {
    if (t < stride) { sh3[t] += sh3[t + stride]; }
    workgroupBarrier();
  }
  let c = vec3f(sh3[0].x / P.mass, sh3[0].y / P.Ix, sh3[0].z / P.Iy);
  for (var k = t; k < P.n; k += 256u) {
    let i = i32(k % P.nx); let j = i32(k / P.nx);
    field[k] -= c.x + c.y * nodeY(j) - c.z * nodeX(i);
  }
}