/**
 * The solver: a second-order leapfrog for the acoustic wave equation on a
 * cubic grid,
 *
 *   p_tt + 2*sig*p_t = c(x)^2 * lap(p) + s(x, t)
 *
 * Pressure only, at constant density, so the medium is two fields: the sound
 * speed c and the absorption sig, both supplied by the scene. Centring the
 * second time derivative and the damping on step n gives an explicit update,
 * which is one compute dispatch per timestep.
 *
 * Two things are worth explaining.
 *
 * **The update is in place.** A step reads p at its six neighbours but reads
 * the previous field pm only at its own index, so a thread may overwrite
 * pm[i] with the new value: no other thread will read it. That leaves two
 * pressure buffers instead of three, which matters at 192^3 where each one is
 * 28 MB, and it removes the buffer-to-buffer copy per step. The roles swap
 * every step, so which buffer holds the current field depends on the parity of
 * the step count; `pressure` reports it and the renderer follows.
 *
 * **Time is a uniform read at a dynamic offset.** A frame's worth of steps is
 * recorded into one command encoder, so nothing written between submits can
 * change inside it, and a clock uploaded per frame would stand still for the
 * whole batch. Rather than carrying time as a grid field (the flat sibling's
 * answer, forced there by its compiler), each step reads its own 64-byte slice
 * of one parameter buffer through a dynamic offset. The whole batch's
 * parameters are written in one call before the pass.
 */
import type { Grid } from './grid.ts';

/** Bytes of parameters per step. Padded to the uniform dynamic-offset
 *  alignment, which WebGPU guarantees to be at most 256. */
const PARAM_STRIDE = 256;

/** Steps that fit in the parameter buffer, and so in one submit. */
export const MAX_STEPS_PER_FRAME = 64;

/** Workgroups are reduced into this many partial maxima before readback. */
const REDUCE_GROUPS = 64;

/** Samples the microphone trace can hold: 4 MB of f32, a few seconds of
 *  audio at the rates the default timestep implies. */
export const MIC_CAPACITY = 1 << 20;

const STEP_SHADER = `
struct Params {
  n: u32,
  h: f32,
  dt: f32,
  t: f32,
  L: f32,
  f: f32,
  t0: f32,
  tw: f32,
  cw: f32,
  x0: f32,
  y0: f32,
  z0: f32,
  w: f32,
  point: f32,
};

@group(0) @binding(0) var<uniform> P: Params;
@group(1) @binding(0) var<storage, read> p: array<f32>;
@group(1) @binding(1) var<storage, read_write> pm: array<f32>;
@group(1) @binding(2) var<storage, read> cs: array<f32>;
@group(1) @binding(3) var<storage, read> sg: array<f32>;

// Outside the domain the field is zero. The absorbing layer is meant to have
// swallowed the wave long before it reaches here.
fn at(ix: i32, iy: i32, iz: i32) -> f32 {
  let n = i32(P.n);
  if (ix < 0 || iy < 0 || iz < 0 || ix >= n || iy >= n || iz >= n) { return 0.0; }
  return p[u32(ix + n * iy + n * n * iz)];
}

@compute @workgroup_size(8, 8, 4)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let n = P.n;
  if (gid.x >= n || gid.y >= n || gid.z >= n) { return; }
  let i = gid.x + n * gid.y + n * n * gid.z;
  let ix = i32(gid.x);
  let iy = i32(gid.y);
  let iz = i32(gid.z);

  let pc = p[i];
  let lap = (at(ix + 1, iy, iz) + at(ix - 1, iy, iz)
           + at(ix, iy + 1, iz) + at(ix, iy - 1, iz)
           + at(ix, iy, iz + 1) + at(ix, iy, iz - 1)
           - 6.0 * pc) / (P.h * P.h);

  let x = -0.5 * P.L + (f32(ix) + 0.5) * P.h;
  let y = -0.5 * P.L + (f32(iy) + 0.5) * P.h;
  let z = -0.5 * P.L + (f32(iz) + 0.5) * P.h;

  // The source. \`cw\` blends between a Gaussian pulse (0) and a wave that
  // turns on smoothly and stays on (1); \`point\` blends between a planar
  // source spanning the grid in y and z, whose far field is a plane wave, and
  // a point source at (x0, y0, z0).
  //
  // The om^2 is a choice of units, not a physical amplitude: a body force of
  // fixed strength drives a response falling off as 1/om^2, so without it the
  // field would shrink every time the frequency slider went up.
  let u = (P.t - P.t0) / P.tw;
  let env = (1.0 - P.cw) * exp(-u * u) + P.cw * 0.5 * (1.0 + tanh(u));
  let gx = (x - P.x0) / P.w;
  let gy = P.point * (y - P.y0) / P.w;
  let gz = P.point * (z - P.z0) / P.w;
  let om = 6.283185307179586 * P.f;
  let s = om * om * env * sin(om * (P.t - P.t0)) * exp(-(gx * gx + gy * gy + gz * gz));

  // One step. The damping is what the absorbing layer acts through: sig is
  // zero over the interior, so there this is the plain leapfrog update.
  let sd = sg[i] * P.dt;
  let cd = cs[i] * P.dt;
  pm[i] = (2.0 * pc - (1.0 - sd) * pm[i] + cd * cd * lap + P.dt * P.dt * s) / (1.0 + sd);
}
`;

// The microphone: the pressure at one grid point, appended to a trace by a
// one-thread dispatch that runs after each step inside the same command
// stream. The obvious implementation — reading the field back and picking out
// one number — costs a GPU-to-CPU round trip per step and would be slower
// than the step; this way the whole trace comes back once, when there is
// something to play. One thread and in-order dispatches mean the plain
// read-modify-write on the counter is safe.
const MIC_SHADER = `
struct MicInfo { probe: u32, cap: u32 };
@group(0) @binding(0) var<uniform> M: MicInfo;
@group(0) @binding(1) var<storage, read> p: array<f32>;
@group(0) @binding(2) var<storage, read_write> trace: array<f32>;
@group(0) @binding(3) var<storage, read_write> count: array<u32>;

@compute @workgroup_size(1)
fn main() {
  let k = count[0];
  if (k < M.cap) {
    trace[k] = p[M.probe];
    count[0] = k + 1u;
  }
}
`;

const REDUCE_SHADER = `
@group(0) @binding(0) var<uniform> npts: u32;
@group(0) @binding(1) var<storage, read> p: array<f32>;
@group(0) @binding(2) var<storage, read_write> out: array<f32>;

var<workgroup> sh: array<f32, 256>;

// Both loops are given uniform trip counts on purpose: a workgroupBarrier may
// only be reached in uniform control flow, and a loop whose exit depends on
// the thread index taints everything after it.
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3u,
        @builtin(local_invocation_id) lid: vec3u,
        @builtin(workgroup_id) wid: vec3u) {
  let stride = 256u * ${REDUCE_GROUPS}u;
  let per = (npts + stride - 1u) / stride;
  var m = 0.0;
  for (var k = 0u; k < per; k = k + 1u) {
    let i = gid.x + k * stride;
    if (i < npts) { m = max(m, abs(p[i])); }
  }
  sh[lid.x] = m;
  workgroupBarrier();
  for (var s = 128u; s > 0u; s = s >> 1u) {
    if (lid.x < s) { sh[lid.x] = max(sh[lid.x], sh[lid.x + s]); }
    workgroupBarrier();
  }
  if (lid.x == 0u) { out[wid.x] = sh[0]; }
}
`;

/** Everything the source term needs, in SI. */
export interface SourceParams {
  f: number;
  cycles: number;
  cw: number;
  point: number;
  x0: number;
  y0: number;
  z0: number;
  w: number;
}

export class Sim {
  readonly grid: Grid;
  /** Model time, seconds. */
  t = 0;
  steps = 0;
  dt: number;
  source: SourceParams;

  #device: GPUDevice;
  #pa: GPUBuffer;
  #pb: GPUBuffer;
  #c: GPUBuffer;
  #sig: GPUBuffer;
  #params: GPUBuffer;
  #paramsHost = new ArrayBuffer(MAX_STEPS_PER_FRAME * PARAM_STRIDE);
  #paramsBG: GPUBindGroup;
  #fieldBG: [GPUBindGroup, GPUBindGroup];
  #pipeline: GPUComputePipeline;
  /** Index of the field bind group to use for the next step. It is also which
   *  buffer currently holds the field: 0 means `pa`. */
  #next = 0;

  #reducePipeline: GPUComputePipeline;
  #reduceBG: [GPUBindGroup, GPUBindGroup];
  #partials: GPUBuffer;
  #readback: GPUBuffer;
  #reading = false;

  /** Samples recorded so far. A host-side mirror of the GPU counter: both
   *  add one per step until the capacity, so they agree exactly. */
  recorded = 0;
  #micPipeline: GPUComputePipeline;
  #micBG: [GPUBindGroup, GPUBindGroup];
  #micUniform: GPUBuffer;
  #trace: GPUBuffer;
  #micCount: GPUBuffer;
  #micRB: GPUBuffer;
  #micReading = false;

  constructor(
    device: GPUDevice,
    grid: Grid,
    medium: { c: Float32Array<ArrayBuffer>; sig: Float32Array<ArrayBuffer> },
    source: SourceParams,
    dt: number,
  ) {
    this.#device = device;
    this.grid = grid;
    this.source = source;
    this.dt = dt;

    const bytes = grid.npts * 4;
    const field = (label: string) =>
      device.createBuffer({
        label,
        size: bytes,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
      });
    this.#pa = field('p-a');
    this.#pb = field('p-b');
    this.#c = field('speed');
    this.#sig = field('absorption');
    device.queue.writeBuffer(this.#c, 0, medium.c);
    device.queue.writeBuffer(this.#sig, 0, medium.sig);

    this.#params = device.createBuffer({
      label: 'step-params',
      size: this.#paramsHost.byteLength,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    const module = device.createShaderModule({ code: STEP_SHADER, label: 'leapfrog3d' });
    const paramsLayout = device.createBindGroupLayout({
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          buffer: { type: 'uniform', hasDynamicOffset: true, minBindingSize: 64 },
        },
      ],
    });
    const fieldLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
      ],
    });
    this.#pipeline = device.createComputePipeline({
      label: 'leapfrog3d',
      layout: device.createPipelineLayout({ bindGroupLayouts: [paramsLayout, fieldLayout] }),
      compute: { module, entryPoint: 'main' },
    });
    this.#paramsBG = device.createBindGroup({
      layout: paramsLayout,
      entries: [{ binding: 0, resource: { buffer: this.#params, size: 64 } }],
    });
    const pair = (read: GPUBuffer, write: GPUBuffer) =>
      device.createBindGroup({
        layout: fieldLayout,
        entries: [
          { binding: 0, resource: { buffer: read } },
          { binding: 1, resource: { buffer: write } },
          { binding: 2, resource: { buffer: this.#c } },
          { binding: 3, resource: { buffer: this.#sig } },
        ],
      });
    this.#fieldBG = [pair(this.#pa, this.#pb), pair(this.#pb, this.#pa)];

    // The colour scale wants max |p| over the whole field, which at 128^3 is
    // 8 MB to read back. A workgroup reduction turns it into 64 floats first.
    const nptsBuf = device.createBuffer({
      size: 4,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    device.queue.writeBuffer(nptsBuf, 0, new Uint32Array([grid.npts]));
    this.#partials = device.createBuffer({
      size: REDUCE_GROUPS * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    this.#readback = device.createBuffer({
      size: REDUCE_GROUPS * 4,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    const reduceLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      ],
    });
    this.#reducePipeline = device.createComputePipeline({
      label: 'maxabs',
      layout: device.createPipelineLayout({ bindGroupLayouts: [reduceLayout] }),
      compute: {
        module: device.createShaderModule({ code: REDUCE_SHADER, label: 'maxabs' }),
        entryPoint: 'main',
      },
    });
    const red = (p: GPUBuffer) =>
      device.createBindGroup({
        layout: reduceLayout,
        entries: [
          { binding: 0, resource: { buffer: nptsBuf } },
          { binding: 1, resource: { buffer: p } },
          { binding: 2, resource: { buffer: this.#partials } },
        ],
      });
    this.#reduceBG = [red(this.#pa), red(this.#pb)];

    // The microphone. All buffers start zeroed, so the counter needs no init.
    this.#micUniform = device.createBuffer({
      label: 'mic-info',
      size: 8,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    this.#trace = device.createBuffer({
      label: 'mic-trace',
      size: MIC_CAPACITY * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    this.#micCount = device.createBuffer({
      label: 'mic-count',
      size: 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });
    this.#micRB = device.createBuffer({
      label: 'mic-readback',
      size: MIC_CAPACITY * 4,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    const micLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
        { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      ],
    });
    this.#micPipeline = device.createComputePipeline({
      label: 'microphone',
      layout: device.createPipelineLayout({ bindGroupLayouts: [micLayout] }),
      compute: {
        module: device.createShaderModule({ code: MIC_SHADER, label: 'microphone' }),
        entryPoint: 'main',
      },
    });
    const micOn = (p: GPUBuffer) =>
      device.createBindGroup({
        layout: micLayout,
        entries: [
          { binding: 0, resource: { buffer: this.#micUniform } },
          { binding: 1, resource: { buffer: p } },
          { binding: 2, resource: { buffer: this.#trace } },
          { binding: 3, resource: { buffer: this.#micCount } },
        ],
      });
    this.#micBG = [micOn(this.#pa), micOn(this.#pb)];
    this.setProbe(0, 0, 0);
  }

  /** Both pressure buffers; `pressureIndex` says which holds the field now. */
  get pressures(): [GPUBuffer, GPUBuffer] {
    return [this.#pa, this.#pb];
  }
  get pressureIndex(): number {
    return this.#next;
  }
  get speed(): GPUBuffer {
    return this.#c;
  }

  /** Point the microphone at a grid cell. Does not restart the trace: moving
   *  the microphone during a run is a microphone that moved. */
  setProbe(ix: number, iy: number, iz: number): void {
    const n = this.grid.n;
    const cl = (i: number) => Math.min(n - 1, Math.max(0, i));
    const probe = cl(ix) + n * cl(iy) + n * n * cl(iz);
    this.#device.queue.writeBuffer(this.#micUniform, 0, new Uint32Array([probe, MIC_CAPACITY]));
  }

  /** Start the recording over. Called on restart, and whenever dt changes:
   *  the trace is one sample per step, and two timesteps would be two sample
   *  rates in one buffer. */
  resetTrace(): void {
    this.#device.queue.writeBuffer(this.#micCount, 0, new Uint32Array([0]));
    this.recorded = 0;
  }

  /** The recorded trace, back from the GPU. Null if a read is in flight. */
  async readTrace(): Promise<Float32Array | null> {
    if (this.#micReading) return null;
    if (this.recorded === 0) return new Float32Array(0);
    this.#micReading = true;
    try {
      const bytes = this.recorded * 4;
      const enc = this.#device.createCommandEncoder({ label: 'mic-read' });
      enc.copyBufferToBuffer(this.#trace, 0, this.#micRB, 0, bytes);
      this.#device.queue.submit([enc.finish()]);
      await this.#micRB.mapAsync(GPUMapMode.READ, 0, bytes);
      const out = new Float32Array(this.#micRB.getMappedRange(0, bytes).slice(0));
      this.#micRB.unmap();
      return out;
    } finally {
      this.#micReading = false;
    }
  }

  /** Back to a silent grid at t = 0. */
  restart(): void {
    const enc = this.#device.createCommandEncoder();
    enc.clearBuffer(this.#pa);
    enc.clearBuffer(this.#pb);
    this.#device.queue.submit([enc.finish()]);
    this.t = 0;
    this.steps = 0;
    this.#next = 0;
    this.resetTrace();
  }

  /** Take `n` timesteps, all in one submit. */
  run(n: number): void {
    const count = Math.min(n, MAX_STEPS_PER_FRAME);
    const g = this.grid;
    const s = this.source;
    // A pulse `cycles` long, delayed enough that it starts near zero.
    const tw = s.cycles / Math.max(s.f, 1e-6);
    const t0 = 2.5 * tw;
    for (let k = 0; k < count; k++) {
      const off = k * PARAM_STRIDE;
      const u32 = new Uint32Array(this.#paramsHost, off, 1);
      const f32 = new Float32Array(this.#paramsHost, off, 16);
      u32[0] = g.n;
      f32[1] = g.h;
      f32[2] = this.dt;
      f32[3] = this.t + k * this.dt;
      f32[4] = g.L;
      f32[5] = s.f;
      f32[6] = t0;
      f32[7] = tw;
      f32[8] = s.cw;
      f32[9] = s.x0;
      f32[10] = s.y0;
      f32[11] = s.z0;
      f32[12] = s.w;
      f32[13] = s.point;
    }
    this.#device.queue.writeBuffer(this.#params, 0, this.#paramsHost, 0, count * PARAM_STRIDE);

    const wg = [g.n / 8, g.n / 8, g.n / 4] as const;
    const enc = this.#device.createCommandEncoder({ label: 'steps' });
    const pass = enc.beginComputePass();
    for (let k = 0; k < count; k++) {
      pass.setPipeline(this.#pipeline);
      pass.setBindGroup(0, this.#paramsBG, [k * PARAM_STRIDE]);
      pass.setBindGroup(1, this.#fieldBG[this.#next]);
      pass.dispatchWorkgroups(wg[0], wg[1], wg[2]);
      this.#next ^= 1;
      // Record the field this step just wrote, which the toggle now points at.
      pass.setPipeline(this.#micPipeline);
      pass.setBindGroup(0, this.#micBG[this.#next]);
      pass.dispatchWorkgroups(1);
    }
    pass.end();
    this.#device.queue.submit([enc.finish()]);
    this.t += count * this.dt;
    this.steps += count;
    this.recorded = Math.min(this.recorded + count, MIC_CAPACITY);
  }

  /**
   * Largest |p| anywhere, for the colour scale. Asynchronous and
   * self-throttling: while one request is in flight further ones return null
   * rather than queueing up.
   */
  async peak(): Promise<number | null> {
    if (this.#reading) return null;
    this.#reading = true;
    try {
      const enc = this.#device.createCommandEncoder({ label: 'peak' });
      const pass = enc.beginComputePass();
      pass.setPipeline(this.#reducePipeline);
      pass.setBindGroup(0, this.#reduceBG[this.pressureIndex]);
      pass.dispatchWorkgroups(REDUCE_GROUPS);
      pass.end();
      enc.copyBufferToBuffer(this.#partials, 0, this.#readback, 0, REDUCE_GROUPS * 4);
      this.#device.queue.submit([enc.finish()]);
      await this.#readback.mapAsync(GPUMapMode.READ);
      const v = new Float32Array(this.#readback.getMappedRange().slice(0));
      this.#readback.unmap();
      let m = 0;
      for (const x of v) m = Math.max(m, x);
      return m;
    } finally {
      this.#reading = false;
    }
  }

  destroy(): void {
    const own = [this.#pa, this.#pb, this.#c, this.#sig, this.#params, this.#partials];
    for (const b of [...own, this.#micUniform, this.#trace, this.#micCount]) b.destroy();
    // Mapping is asynchronous; destroying a buffer with a pending map is an
    // error, so leave the readback buffers to be collected.
  }
}