/**
 * The Laplacian stencils, as GPU dispatches.
 *
 * A grid field is an npts x 1 column vector in the buffer, laid out x-fastest:
 * the point (ix, iy) is element `ix + nx*iy`. That is the only place in the
 * project where the layout matters to anything but the renderer, which is the
 * reason the stencils are host-provided operations rather than something a .m
 * expresses with array slicing.
 *
 * Outside the grid the field is taken to be zero, which makes the outer
 * boundary sound-soft (a pressure-release wall). Every scene puts an absorbing
 * layer in front of it (tools/sponge.m), so in a well set-up run almost nothing
 * reaches the wall to be reflected; what does is attenuated on the way out and
 * again on the way back. See the README on how good an open boundary that is.
 *
 * Both stencils divide by h^2, so a .m reads `lap2(p)` as the Laplacian in the
 * physical units the scene's coordinates are in. The grid spacing is compiled
 * into the shader — the grid is fixed when a model is compiled, and changing it
 * recompiles anyway.
 */
import { UnsupportedOnGpu, WORKGROUP_SIZE } from './wgsl.ts';

export type StencilKind = 'lap2' | 'lap4';

export interface StencilGrid {
  nx: number;
  ny: number;
  /** Grid spacing, the same in x and y. */
  h: number;
}

/**
 * The 5-point second-order Laplacian, and the 9-point fourth-order one.
 *
 * The fourth-order stencil is the standard (-1/12, 4/3, -5/2, 4/3, -1/12)/h^2
 * one-dimensional second derivative applied along each axis, so it is a
 * five-point line in x plus a five-point line in y (nine points in all, not a
 * 3x3 block). It costs twice the reads of `lap2` and buys two orders of
 * accuracy, which for a wave problem shows up as much less grid dispersion:
 * a pulse travelling many wavelengths stays a pulse instead of trailing
 * numerical ripples. See models/leapfrog4.m.
 */
function stencilWGSL(kind: StencilKind, grid: StencilGrid, npts: number): string {
  const { nx, ny, h } = grid;
  const inv = 1 / (h * h);
  // A read of a point outside the grid returns zero.
  const at = `
fn at(ix: i32, iy: i32) -> f32 {
  if (ix < 0 || ix >= ${nx} || iy < 0 || iy >= ${ny}) { return 0.0; }
  return src[u32(ix + ${nx} * iy)];
}
`;
  const body =
    kind === 'lap2'
      ? `  let s = at(ix - 1, iy) + at(ix + 1, iy) + at(ix, iy - 1) + at(ix, iy + 1)
          - 4.0 * at(ix, iy);`
      : `  let s = (-1.0 / 12.0) * (at(ix - 2, iy) + at(ix + 2, iy) + at(ix, iy - 2) + at(ix, iy + 2))
          + (4.0 / 3.0) * (at(ix - 1, iy) + at(ix + 1, iy) + at(ix, iy - 1) + at(ix, iy + 1))
          - 5.0 * at(ix, iy);`;
  return `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;
${at}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${npts}u) { return; }
  let ix = i32(i % ${nx}u);
  let iy = i32(i / ${nx}u);
${body}
  dst[i] = ${inv}f * s;
}
`;
}

/** Compiled stencil pipelines for one grid. Shared by every plan on it. */
export class StencilPlan {
  readonly grid: StencilGrid;
  readonly npts: number;

  #device: GPUDevice;
  #layout: GPUBindGroupLayout;
  #pipelines = new Map<StencilKind, GPUComputePipeline>();

  private constructor(device: GPUDevice, grid: StencilGrid, layout: GPUBindGroupLayout) {
    this.#device = device;
    this.grid = grid;
    this.npts = grid.nx * grid.ny;
    this.#layout = layout;
  }

  static create(device: GPUDevice, grid: StencilGrid): StencilPlan {
    const layout = device.createBindGroupLayout({
      label: 'stencil',
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          buffer: { type: 'storage' },
        },
        {
          binding: 1,
          visibility: GPUShaderStage.COMPUTE,
          buffer: { type: 'read-only-storage' },
        },
      ],
    });
    return new StencilPlan(device, grid, layout);
  }

  /**
   * The pipeline for one stencil, compiled on first use. Both are cheap, but
   * a model uses one of them, and compiling only what is asked for keeps the
   * op sequence honest about what the .m actually costs.
   */
  async pipeline(kind: StencilKind): Promise<GPUComputePipeline> {
    const existing = this.#pipelines.get(kind);
    if (existing) return existing;
    const code = stencilWGSL(kind, this.grid, this.npts);
    const module = this.#device.createShaderModule({ code, label: kind });
    let pipeline: GPUComputePipeline;
    try {
      pipeline = await this.#device.createComputePipelineAsync({
        layout: this.#device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
        compute: { module, entryPoint: 'main' },
        label: kind,
      });
    } catch (e) {
      // No error scope here either; see makePipeline in plan.ts.
      throw new UnsupportedOnGpu(
        `stencil '${kind}': ${e instanceof Error ? e.message : String(e)}`,
      );
    }
    this.#pipelines.set(kind, pipeline);
    return pipeline;
  }

  createBinding(src: GPUBuffer, dst: GPUBuffer): GPUBindGroup {
    return this.#device.createBindGroup({
      layout: this.#layout,
      entries: [
        { binding: 0, resource: { buffer: dst } },
        { binding: 1, resource: { buffer: src } },
      ],
    });
  }

  get workgroups(): number {
    return Math.ceil(this.npts / WORKGROUP_SIZE);
  }
}