/**
 * A microphone: the pressure at one grid point, sampled every timestep.
 *
 * The obvious implementation — read the field back and pick out one number —
 * costs a GPU-to-CPU round trip per step, which is more than the step itself.
 * So the trace is written on the GPU instead, by a one-thread dispatch that
 * runs after each step and appends `p` at the probe point to a buffer. The
 * whole trace comes back to the CPU once, when there is something to listen
 * to.
 *
 * One sample per timestep is the natural rate: it is every value the
 * simulation has, and nothing is being resampled or interpolated on the way
 * in. What that means in seconds is decided at playback (src/audio/play.ts),
 * because the simulation has no seconds in it — only model time.
 */

/** Samples the trace holds: about 7 seconds of audio at the default
 *  timestep, and 4 MB of GPU memory. Recording stops when it is full rather
 *  than wrapping, so what you hear always starts where the pluck did. */
export const TRACE_CAPACITY = 1 << 20;

const SHADER = `
struct Probe {
  index: u32,
  capacity: u32,
};

@group(0) @binding(0) var<storage, read_write> trace: array<f32>;
@group(0) @binding(1) var<storage, read_write> head: array<u32>;
@group(0) @binding(2) var<storage, read> field: array<f32>;
@group(0) @binding(3) var<uniform> probe: Probe;

// One invocation, so the read-modify-write of the head needs no atomic.
@compute @workgroup_size(1)
fn main() {
  let i = head[0];
  if (i < probe.capacity) {
    trace[i] = field[probe.index];
    head[0] = i + 1u;
  }
}
`;

export interface RecorderOptions {
  device: GPUDevice;
  /** The pressure buffer to sample — the host-owned one, which is what both
   *  `init` and `step` leave their result in. */
  field: GPUBuffer;
  nx: number;
  ny: number;
  nz: number;
}

export class Recorder {
  readonly capacity = TRACE_CAPACITY;

  #device: GPUDevice;
  #pipeline: GPUComputePipeline | null = null;
  #layout: GPUBindGroupLayout;
  #bindGroup: GPUBindGroup | null = null;
  #trace: GPUBuffer;
  #head: GPUBuffer;
  #probe: GPUBuffer;
  #readback: GPUBuffer;
  #field: GPUBuffer;
  #nx: number;
  #ny: number;
  #nz: number;
  /** Samples written since the last clear, as far as the host knows. Counted
   *  here rather than read back from the GPU: the dispatch runs once per step
   *  and the host knows exactly how many steps it asked for. */
  #count = 0;
  #reading = false;

  private constructor(init: {
    device: GPUDevice;
    layout: GPUBindGroupLayout;
    trace: GPUBuffer;
    head: GPUBuffer;
    probe: GPUBuffer;
    readback: GPUBuffer;
    field: GPUBuffer;
    nx: number;
    ny: number;
    nz: number;
  }) {
    this.#device = init.device;
    this.#layout = init.layout;
    this.#trace = init.trace;
    this.#head = init.head;
    this.#probe = init.probe;
    this.#readback = init.readback;
    this.#field = init.field;
    this.#nx = init.nx;
    this.#ny = init.ny;
    this.#nz = init.nz;
  }

  static async create(opts: RecorderOptions): Promise<Recorder> {
    const { device } = opts;
    const layout = device.createBindGroupLayout({
      label: 'recorder',
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: '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: 'uniform' } },
      ],
    });
    const trace = device.createBuffer({
      label: 'recorder-trace',
      size: 4 * TRACE_CAPACITY,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    const head = device.createBuffer({
      label: 'recorder-head',
      size: 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });
    const probe = device.createBuffer({
      label: 'recorder-probe',
      size: 8,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    const readback = device.createBuffer({
      label: 'recorder-readback',
      size: 4 * TRACE_CAPACITY,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });

    const rec = new Recorder({
      device, layout, trace, head, probe, readback,
      field: opts.field, nx: opts.nx, ny: opts.ny, nz: opts.nz,
    });
    rec.#pipeline = await device.createComputePipelineAsync({
      label: 'recorder',
      layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
      compute: {
        module: device.createShaderModule({ code: SHADER, label: 'recorder' }),
        entryPoint: 'main',
      },
    });
    rec.#bindGroup = device.createBindGroup({
      layout,
      entries: [
        { binding: 0, resource: { buffer: trace } },
        { binding: 1, resource: { buffer: head } },
        { binding: 2, resource: { buffer: opts.field } },
        { binding: 3, resource: { buffer: probe } },
      ],
    });
    rec.clear();
    return rec;
  }

  /** Samples recorded so far. Stops rising once the trace is full. */
  get count(): number {
    return Math.min(this.#count, this.capacity);
  }

  get full(): boolean {
    return this.#count >= this.capacity;
  }

  /** Put the microphone at the grid point nearest (ix, iy, iz). Free: the
   *  probe index is a uniform, so moving it disturbs neither the run nor the
   *  recording already made. */
  setProbe(ix: number, iy: number, iz: number): void {
    const cx = Math.max(0, Math.min(this.#nx - 1, Math.round(ix)));
    const cy = Math.max(0, Math.min(this.#ny - 1, Math.round(iy)));
    const cz = Math.max(0, Math.min(this.#nz - 1, Math.round(iz)));
    this.#device.queue.writeBuffer(
      this.#probe,
      0,
      new Uint32Array([cx + this.#nx * (cy + this.#ny * cz), this.capacity]),
    );
  }

  /** Start again from an empty trace. */
  clear(): void {
    this.#count = 0;
    this.#device.queue.writeBuffer(this.#head, 0, new Uint32Array([0]));
  }

  /** Record one sample. Called once per timestep, inside the step's own
   *  submission, so no extra work crosses to the host. */
  encode(encoder: GPUCommandEncoder): void {
    if (!this.#pipeline || !this.#bindGroup || this.full) return;
    const pass = encoder.beginComputePass({ label: 'recorder' });
    pass.setPipeline(this.#pipeline);
    pass.setBindGroup(0, this.#bindGroup);
    pass.dispatchWorkgroups(1);
    pass.end();
    this.#count++;
  }

  /** The recorded trace. The only readback the microphone ever does. */
  async read(): Promise<Float32Array> {
    const n = this.count;
    if (n === 0) return new Float32Array(0);
    if (this.#reading) throw new Error('a trace readback is already in flight');
    this.#reading = true;
    try {
      const enc = this.#device.createCommandEncoder({ label: 'recorder-read' });
      enc.copyBufferToBuffer(this.#trace, 0, this.#readback, 0, 4 * n);
      this.#device.queue.submit([enc.finish()]);
      await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * n);
      const out = new Float32Array(this.#readback.getMappedRange(0, 4 * n).slice(0));
      this.#readback.unmap();
      return out;
    } finally {
      this.#reading = false;
    }
  }

  /** The pressure buffer this microphone listens to. */
  get field(): GPUBuffer {
    return this.#field;
  }

  destroy(): void {
    this.#trace.destroy();
    this.#head.destroy();
    this.#probe.destroy();
    this.#readback.destroy();
  }
}