// The WGSL for the WebGPU method of fundamental solutions.
//
// Everything here is f32, because that is what WebGPU has: WGSL's floating
// point types are f32 and f16, there is no f64, and no extension in the
// standard adds one. That single fact is the most interesting thing about
// this solver, because the MFS is conditioning-limited: its accuracy stops
// improving where rounding in the collocation solve overtakes the
// approximation error, and in f32 that happens six to nine orders short of
// where it happens in double. See the solver's manifest entry for the
// measured numbers.
//
// Note that software emulation of double precision (carrying each value as
// an unevaluated sum of two f32, the "double-single" trick) is not a way
// out. It depends on error-free transformations such as
// s = a + b; err = b - (s - a), which are only error-free if the compiler
// evaluates them exactly as written; WGSL permits an implementation to
// compute with greater precision or to reassociate, and the Mesa driver
// this was tried on does exactly that, returning f32-accurate results for
// values loaded from a buffer. So f32 is what a WebGPU solver gets.

/** Workgroup sizes: the two-dimensional kernels (assembly and the trailing
 * update), the one-dimensional ones, and the evaluation. */
export const TILE = 16;
export const LANES = 256;
export const EVAL_GROUP = 64;

/**
 * The whole solver as one shader module.
 *
 * The matrix is the augmented system [A | b], n rows of ld = n + 1 entries,
 * row-major, so that the elimination applies to the right-hand side as it
 * goes and no separate forward substitution is needed. The multipliers are
 * never stored: nothing reads them again.
 *
 * One read-only buffer holds every input, because the default WebGPU limit
 * is eight storage buffers per stage. Its layout, in f32 elements, is
 * derived from n rather than passed: collocation points (x, y interleaved)
 * at 0, charges at 2n, the Dirichlet data at 4n, the targets (interleaved)
 * at 5n.
 */
export function mfsShader(): string {
  return `// generated by src/solvers/mfs-gpu/wgsl.ts

struct Dims {
  n: u32,   // charges = collocation points
  ld: u32,  // row stride of the augmented matrix, n + 1
  m: u32,   // evaluation targets
  pad: u32,
};

struct Step { k: u32, pad0: u32, pad1: u32, pad2: u32 };

@group(0) @binding(0) var<uniform> dims: Dims;
@group(0) @binding(1) var<uniform> step: Step;
@group(0) @binding(2) var<storage, read> data: array<f32>;
@group(0) @binding(3) var<storage, read_write> mat: array<f32>;
@group(0) @binding(4) var<storage, read_write> lcol: array<f32>;
@group(0) @binding(5) var<storage, read_write> sol: array<f32>;
@group(0) @binding(6) var<storage, read_write> out: array<f32>;

fn pt_x(i: u32) -> f32 { return data[2u * i]; }
fn pt_y(i: u32) -> f32 { return data[2u * i + 1u]; }
fn ch_x(j: u32) -> f32 { return data[2u * dims.n + 2u * j]; }
fn ch_y(j: u32) -> f32 { return data[2u * dims.n + 2u * j + 1u]; }
fn bdata(i: u32) -> f32 { return data[4u * dims.n + i]; }
fn tg_x(t: u32) -> f32 { return data[5u * dims.n + 2u * t]; }
fn tg_y(t: u32) -> f32 { return data[5u * dims.n + 2u * t + 1u]; }

/** The MFS kernel log|p - q|, written as half the log of the squared
 * distance so that no square root is taken. */
fn logdist(ax: f32, ay: f32, bx: f32, by: f32) -> f32 {
  let dx = ax - bx;
  let dy = ay - by;
  return 0.5 * log(dx * dx + dy * dy);
}

@compute @workgroup_size(${TILE}, ${TILE})
fn assemble(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.y;
  let j = gid.x;
  if (i >= dims.n || j >= dims.ld) { return; }
  if (j == dims.n) {
    mat[i * dims.ld + j] = bdata(i);
    return;
  }
  mat[i * dims.ld + j] = logdist(pt_x(i), pt_y(i), ch_x(j), ch_y(j));
}

var<workgroup> best_val: array<f32, ${LANES}>;
var<workgroup> best_idx: array<u32, ${LANES}>;

/** Partial pivoting for step k: find the largest |A(i,k)| over i >= k and
 * swap that row with row k. One workgroup, so the search reduces in shared
 * memory and the swap needs no second dispatch. */
@compute @workgroup_size(${LANES})
fn pivot(@builtin(local_invocation_id) lid: vec3<u32>) {
  let k = step.k;
  let n = dims.n;
  let ld = dims.ld;
  var bv = -1.0;
  var bi = k;
  for (var i = k + lid.x; i < n; i = i + ${LANES}u) {
    let v = abs(mat[i * ld + k]);
    if (v > bv) { bv = v; bi = i; }
  }
  best_val[lid.x] = bv;
  best_idx[lid.x] = bi;
  workgroupBarrier();
  for (var s = ${LANES / 2}u; s > 0u; s = s >> 1u) {
    if (lid.x < s && best_val[lid.x + s] > best_val[lid.x]) {
      best_val[lid.x] = best_val[lid.x + s];
      best_idx[lid.x] = best_idx[lid.x + s];
    }
    workgroupBarrier();
  }
  let p = best_idx[0];
  if (p != k) {
    for (var j = k + lid.x; j < ld; j = j + ${LANES}u) {
      let a = mat[k * ld + j];
      mat[k * ld + j] = mat[p * ld + j];
      mat[p * ld + j] = a;
    }
  }
}

/** The multipliers of step k, computed once rather than once per element of
 * the trailing update. */
@compute @workgroup_size(${LANES})
fn multipliers(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = step.k;
  let i = k + 1u + gid.x;
  if (i >= dims.n) { return; }
  lcol[i] = mat[i * dims.ld + k] / mat[k * dims.ld + k];
}

/** The rank-one update of the trailing submatrix, columns k+1 .. ld-1, so
 * the augmented right-hand column is eliminated along with the rest. */
@compute @workgroup_size(${TILE}, ${TILE})
fn update(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = step.k;
  let i = k + 1u + gid.y;
  let j = k + 1u + gid.x;
  if (i >= dims.n || j >= dims.ld) { return; }
  let ld = dims.ld;
  mat[i * ld + j] = mat[i * ld + j] - lcol[i] * mat[k * ld + j];
}

/** Back substitution, one step per dispatch: x(k) from row k of U and the
 * running right-hand column, then that column updated above row k. Every
 * lane recomputes x(k) so that no barrier is needed. */
@compute @workgroup_size(${LANES})
fn backsub(@builtin(local_invocation_id) lid: vec3<u32>) {
  let k = step.k;
  let ld = dims.ld;
  let xk = mat[k * ld + dims.n] / mat[k * ld + k];
  if (lid.x == 0u) { sol[k] = xk; }
  for (var i = lid.x; i < k; i = i + ${LANES}u) {
    mat[i * ld + dims.n] = mat[i * ld + dims.n] - mat[i * ld + k] * xk;
  }
}

/** The potential at the targets: one thread per target, summing the
 * charges. */
@compute @workgroup_size(${EVAL_GROUP})
fn evaluate(@builtin(global_invocation_id) gid: vec3<u32>) {
  let t = gid.x;
  if (t >= dims.m) { return; }
  let ax = tg_x(t);
  let ay = tg_y(t);
  var acc = 0.0;
  for (var j = 0u; j < dims.n; j = j + 1u) {
    acc = acc + sol[j] * logdist(ax, ay, ch_x(j), ch_y(j));
  }
  out[t] = acc;
}
`;
}