// WebGPU time-stepping driver for the von Karman plate.
//
// The simulation advances in batches of one movie frame (1/fps seconds). A batch
// encodes one compute dispatch per time step (with a per-step uniform addressed
// by dynamic offset), a rigid-body projection every RB_INTERVAL steps, one
// retarded-time audio sum per audio sample, and one movie-frame extraction at
// the end. Audio, frame, and statistics are read back after every batch.

import {
  type PhysParams,
  type Grid,
  type Material,
  makeGrid,
  makeMaterial,
  nodeX,
  nodeY,
  trapWeight,
  HIST_LEN,
  RB_INTERVAL,
  RB_INTERVAL_HELD,
  RHO_AIR,
  SPEED_OF_SOUND,
} from "../params.ts";
import { maxAngularFrequency } from "./cpuRef.ts";
import commonWgsl from "./shaders/common.wgsl?raw";
import stepWgsl from "./shaders/step.wgsl?raw";
import rigidWgsl from "./shaders/rigid.wgsl?raw";
import audioWgsl from "./shaders/audio.wgsl?raw";
import extractWgsl from "./shaders/extract.wgsl?raw";

export type Phase = "push" | "hold" | "recording" | "done";

export interface SimStatus {
  phase: Phase;
  /** Simulation time relative to the release [s]. */
  t: number;
  totalSteps: number;
  ke: number;
  wmax: number;
  samples: number;
  frames: number;
  peakPressure: number;
  delayClamped: boolean;
}

const STATS_LEN = 16;

export class Simulation {
  readonly grid: Grid;
  readonly mat: Material;
  readonly dt: number;
  /** Time steps per audio sample. */
  readonly stepsPerSample: number;
  readonly samplesPerBatch: number;
  readonly stepsPerBatch: number;
  readonly omegaMax: number;

  /** Raw listener pressure [Pa] at fs, from the release. */
  audio: Float32Array = new Float32Array(0);
  samples = 0;
  /** Movie frames of w [m]; frame k is at t = k / fps. */
  frames: Float32Array<ArrayBuffer>[] = [];

  phase: Phase = "push";
  /** Steps since the start of the current phase clock (prep or recording). */
  private stepCount = 0;
  private prepTime = 0;
  totalSteps = 0;
  ke = 0;
  wmax = 0;
  peakPressure = 0;
  delayClamped = false;
  private hpState = { y: 0, x: 0 };

  private device: GPUDevice;
  private params: PhysParams;
  private paramsBuf: GPUBuffer;
  private stepBuf: GPUBuffer;
  private stepStride: number;
  private stepData: ArrayBuffer;
  private pos: GPUBuffer[];
  private vel: GPUBuffer[];
  private hist: GPUBuffer;
  private react: GPUBuffer;
  private lw: GPUBuffer;
  private audioOut: GPUBuffer;
  private frameBuf: GPUBuffer;
  private statsBuf: GPUBuffer;
  private staging: GPUBuffer[];
  private stagingIdx = 0;
  /** Index of the buffer pair holding the current state. */
  private cur = 0;
  /** The `displayW` buffer is a plain f32 array of w for the renderer. */
  readonly displayW: GPUBuffer;

  private stepPipe: GPUComputePipeline;
  private sumsPipe: GPUComputePipeline;
  private applyPipe: GPUComputePipeline;
  private audioPipe: GPUComputePipeline;
  private extractPipe: GPUComputePipeline;
  private stepBG: GPUBindGroup[];
  private rbBG: GPUBindGroup[];
  private audioBG: GPUBindGroup;
  private extractBG: GPUBindGroup[];
  private displayBG: GPUBindGroup[];
  private stagingSize: number;

  constructor(device: GPUDevice, params: PhysParams) {
    this.device = device;
    this.params = { ...params };
    const g = makeGrid(params);
    const m = makeMaterial(params);
    this.grid = g;
    this.mat = m;

    // Stable time step, aligned so that an integer number of steps make one audio sample.
    this.omegaMax = maxAngularFrequency(g, m);
    const dtMax = (0.85 * 2) / this.omegaMax;
    this.stepsPerSample = Math.ceil(1 / (params.fs * dtMax));
    this.dt = 1 / (params.fs * this.stepsPerSample);
    this.samplesPerBatch = Math.max(1, Math.round(params.fs / params.fps));
    this.stepsPerBatch = this.samplesPerBatch * this.stepsPerSample;

    const n = g.n;
    const f32 = (size: number, usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST) =>
      device.createBuffer({ size: Math.max(16, size * 4), usage });

    this.pos = [f32(4 * n), f32(4 * n)]; // [u | v | w | wLo]
    this.vel = [f32(3 * n), f32(3 * n)]; // [du | dv | dw]
    this.hist = f32(HIST_LEN * n);
    this.react = f32(n);
    this.lw = f32(2 * n);
    this.audioOut = f32(this.samplesPerBatch);
    this.frameBuf = f32(n);
    this.displayW = f32(n);
    this.statsBuf = f32(STATS_LEN);
    this.stagingSize = (this.samplesPerBatch + n + STATS_LEN) * 4;
    this.staging = [0, 1].map(() =>
      device.createBuffer({ size: this.stagingSize, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }),
    );

    // Constant parameters.
    let mass = 0, Ix = 0, Iy = 0;
    for (let j = 0; j < g.ny; j++) {
      for (let i = 0; i < g.nx; i++) {
        const mk = m.rhoh * g.hx * g.hy * trapWeight(g, i, j);
        mass += mk;
        Ix += mk * nodeY(g, j) ** 2;
        Iy += mk * nodeX(g, i) ** 2;
      }
    }
    const pf = new Float32Array(20);
    const pu = new Uint32Array(pf.buffer);
    pu[0] = g.nx; pu[1] = g.ny; pu[2] = n; pu[3] = HIST_LEN;
    pf[4] = g.hx; pf[5] = g.hy; pf[6] = this.dt; pf[7] = m.nu;
    pf[8] = m.D; pf[9] = m.C; pf[10] = m.G; pf[11] = m.rhoh;
    pf[12] = params.Lx / 2; pf[13] = -params.Ly / 2; pf[14] = params.fingerRadius ** 2; pf[15] = mass;
    pf[16] = Ix; pf[17] = Iy; pf[18] = Ix + Iy; pf[19] = 0;
    this.paramsBuf = device.createBuffer({ size: pf.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(this.paramsBuf, 0, pf);

    // Per-step parameters, one 48-byte struct per step at a dynamic offset.
    this.stepStride = Math.max(48, device.limits.minUniformBufferOffsetAlignment);
    this.stepData = new ArrayBuffer(this.stepStride * this.stepsPerBatch);
    this.stepBuf = device.createBuffer({ size: this.stepData.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });

    const mk = (code: string, entry: string, layout: GPUPipelineLayout | "auto" = "auto") =>
      device.createComputePipeline({
        layout,
        compute: { module: device.createShaderModule({ code: commonWgsl + code }), entryPoint: entry },
      });
    const entry = (binding: number, type: GPUBufferBindingType, hasDynamicOffset = false): GPUBindGroupLayoutEntry => ({
      binding,
      visibility: GPUShaderStage.COMPUTE,
      buffer: { type, hasDynamicOffset },
    });
    const stepLayout = device.createBindGroupLayout({
      entries: [
        entry(0, "uniform"),
        entry(1, "uniform", true),
        entry(2, "read-only-storage"),
        entry(3, "read-only-storage"),
        entry(4, "storage"),
        entry(5, "storage"),
        entry(6, "storage"),
        entry(7, "storage"),
      ],
    });
    this.stepPipe = mk(stepWgsl, "main", device.createPipelineLayout({ bindGroupLayouts: [stepLayout] }));
    const audioLayout = device.createBindGroupLayout({
      entries: [entry(0, "uniform"), entry(1, "uniform", true), entry(2, "read-only-storage"), entry(3, "read-only-storage"), entry(4, "storage")],
    });
    const rigidModule = device.createShaderModule({ code: commonWgsl + rigidWgsl });
    const rbLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
        { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
        { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
        { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
      ],
    });
    const rbPipeLayout = device.createPipelineLayout({ bindGroupLayouts: [rbLayout] });
    this.sumsPipe = device.createComputePipeline({ layout: rbPipeLayout, compute: { module: rigidModule, entryPoint: "sums" } });
    this.applyPipe = device.createComputePipeline({ layout: rbPipeLayout, compute: { module: rigidModule, entryPoint: "apply" } });
    this.audioPipe = mk(audioWgsl, "main", device.createPipelineLayout({ bindGroupLayouts: [audioLayout] }));
    this.extractPipe = mk(extractWgsl, "main");

    const sp = { buffer: this.stepBuf, size: 48 };
    this.stepBG = [0, 1].map((p) =>
      device.createBindGroup({
        layout: stepLayout,
        entries: [
          { binding: 0, resource: { buffer: this.paramsBuf } },
          { binding: 1, resource: sp },
          { binding: 2, resource: { buffer: this.pos[p] } },
          { binding: 3, resource: { buffer: this.vel[p] } },
          { binding: 4, resource: { buffer: this.pos[1 - p] } },
          { binding: 5, resource: { buffer: this.vel[1 - p] } },
          { binding: 6, resource: { buffer: this.hist } },
          { binding: 7, resource: { buffer: this.react } },
        ],
      }),
    );
    this.rbBG = [0, 1].map((p) =>
      device.createBindGroup({
        layout: rbLayout,
        entries: [
          { binding: 0, resource: { buffer: this.paramsBuf } },
          { binding: 1, resource: { buffer: this.pos[p] } },
          { binding: 2, resource: { buffer: this.vel[p] } },
          { binding: 3, resource: { buffer: this.statsBuf } },
        ],
      }),
    );
    this.audioBG = device.createBindGroup({
      layout: audioLayout,
      entries: [
        { binding: 0, resource: { buffer: this.paramsBuf } },
        { binding: 1, resource: sp },
        { binding: 2, resource: { buffer: this.hist } },
        { binding: 3, resource: { buffer: this.lw } },
        { binding: 4, resource: { buffer: this.audioOut } },
      ],
    });
    const extractBG = (p: number, out: GPUBuffer) =>
      device.createBindGroup({
        layout: this.extractPipe.getBindGroupLayout(0),
        entries: [
          { binding: 0, resource: { buffer: this.paramsBuf } },
          { binding: 1, resource: { buffer: this.pos[p] } },
          { binding: 2, resource: { buffer: out } },
        ],
      });
    this.extractBG = [0, 1].map((p) => extractBG(p, this.frameBuf));
    this.displayBG = [0, 1].map((p) => extractBG(p, this.displayW));

    this.setListener(params.listenerDist, params.listenerAzimuth, params.listenerElevation, params.radiation);
  }

  /** Recomputes the per-node weights and delays of the listener sum. Takes effect on later samples. */
  setListener(dist: number, azimuthDeg: number, elevationDeg: number, radiation: "rayleigh" | "dipole"): void {
    const g = this.grid;
    const az = (azimuthDeg * Math.PI) / 180;
    const el = (elevationDeg * Math.PI) / 180;
    const L = [dist * Math.sin(az) * Math.cos(el), dist * Math.sin(el), dist * Math.cos(az) * Math.cos(el)];
    const data = new Float32Array(2 * g.n);
    let dmin = Infinity;
    for (let j = 0; j < g.ny; j++) {
      for (let i = 0; i < g.nx; i++) {
        const k = j * g.nx + i;
        const rx = L[0] - nodeX(g, i), ry = L[1] - nodeY(g, j), rz = L[2];
        const r = Math.sqrt(rx * rx + ry * ry + rz * rz);
        const dA = g.hx * g.hy * trapWeight(g, i, j);
        let wgt = (RHO_AIR * dA) / (2 * Math.PI * r);
        if (radiation === "dipole") wgt *= Math.abs(rz) / r;
        data[2 * k] = wgt;
        data[2 * k + 1] = (r / SPEED_OF_SOUND) * this.params.fs;
        dmin = Math.min(dmin, data[2 * k + 1]);
      }
    }
    this.delayClamped = false;
    for (let k = 0; k < g.n; k++) {
      let d = data[2 * k + 1] - dmin;
      if (d > HIST_LEN - 2) {
        d = HIST_LEN - 2;
        this.delayClamped = true;
      }
      data[2 * k + 1] = d;
    }
    this.device.queue.writeBuffer(this.lw, 0, data);
    this.params.listenerDist = dist;
    this.params.listenerAzimuth = azimuthDeg;
    this.params.listenerElevation = elevationDeg;
    this.params.radiation = radiation;
  }

  setDamping(sigma0: number, sigma1: number): void {
    this.params.sigma0 = sigma0;
    this.params.sigma1 = sigma1;
  }

  setDuration(duration: number): void {
    this.params.duration = duration;
    if (this.phase === "done" && this.samples < this.targetSamples()) this.phase = "recording";
  }

  targetSamples(): number {
    return Math.round(this.params.duration * this.params.fs);
  }

  status(): SimStatus {
    const t = this.phase === "push" || this.phase === "hold" ? this.prepTime - this.prepDuration() : this.stepCount * this.dt;
    return {
      phase: this.phase,
      t,
      totalSteps: this.totalSteps,
      ke: this.ke,
      wmax: this.wmax,
      samples: this.samples,
      frames: this.frames.length,
      peakPressure: this.peakPressure,
      delayClamped: this.delayClamped,
    };
  }

  private prepDuration(): number {
    return this.params.rampTime + this.params.holdMax;
  }

  /** Enqueues an extraction of the current w field into `displayW` (for the live view). */
  encodeDisplay(encoder: GPUCommandEncoder): void {
    const pass = encoder.beginComputePass();
    pass.setPipeline(this.extractPipe);
    pass.setBindGroup(0, this.displayBG[this.cur]);
    pass.dispatchWorkgroups(Math.ceil(this.grid.n / 256));
    pass.end();
  }

  /** Listener audio with the radiation post-processing applied [Pa]. */
  processedAudio(): Float32Array {
    return this.processed.subarray(0, this.samples);
  }
  private processed: Float32Array = new Float32Array(0);

  private appendAudio(raw: Float32Array): void {
    const need = this.samples + raw.length;
    if (this.audio.length < need) {
      const grow = Math.max(need, 2 * this.audio.length, this.params.fs);
      const a = new Float32Array(grow);
      a.set(this.audio.subarray(0, this.samples));
      this.audio = a;
      const b = new Float32Array(grow);
      b.set(this.processed.subarray(0, this.samples));
      this.processed = b;
    }
    this.audio.set(raw, this.samples);
    // Dipole option: first-order high-pass at the baffle cutoff c / (2 sqrt(Lx Ly)).
    const dipole = this.params.radiation === "dipole";
    const fc = SPEED_OF_SOUND / (2 * Math.sqrt(this.params.Lx * this.params.Ly));
    const alpha = 1 / (1 + (2 * Math.PI * fc) / this.params.fs);
    for (let k = 0; k < raw.length; k++) {
      const x = raw[k];
      let y = x;
      if (dipole) {
        y = alpha * (this.hpState.y + x - this.hpState.x);
      }
      this.hpState = { y, x };
      this.processed[this.samples + k] = y;
      this.peakPressure = Math.max(this.peakPressure, Math.abs(y));
    }
    this.samples = need;
  }

  /** Runs one batch (one movie frame of time). Resolves when its results have been read back. */
  async runBatch(): Promise<void> {
    if (this.phase === "done") return;
    const p = this.params;
    const recording = this.phase === "recording";
    const nSteps = this.stepsPerBatch;
    const stepsLeft = recording ? (this.targetSamples() - this.samples) * this.stepsPerSample : nSteps;
    const steps = Math.min(nSteps, Math.max(0, stepsLeft));
    if (steps === 0) {
      this.phase = "done";
      return;
    }

    // Fill the per-step parameters.
    const u32 = new Uint32Array(this.stepData);
    const i32 = new Int32Array(this.stepData);
    const f32 = new Float32Array(this.stepData);
    const stride = this.stepStride / 4;
    const audioStepsInBatch: number[] = [];
    for (let s = 0; s < steps; s++) {
      const o = s * stride;
      const local = this.stepCount + s;
      u32[o] = this.totalSteps + s;
      let flags = 0, pushW = 0, pushV = 0, sigma0 = p.sigma0, slot = -1, outIdx = -1, releaseGain = 0;
      if (!recording) {
        const t = this.prepTime + (s + 1) * this.dt; // state after this step
        const tau = Math.min(1, t / p.rampTime);
        const smooth = tau * tau * (3 - 2 * tau);
        const dsmooth = tau < 1 ? (6 * tau * (1 - tau)) / p.rampTime : 0;
        flags = 1 | (tau >= 1 ? 2 : 0); // record the reaction once the ramp is complete
        pushW = p.pushAmount * smooth;
        pushV = p.pushAmount * dsmooth;
        sigma0 = p.settleSigma0;
      } else {
        const t = local * this.dt;
        if (p.releaseTime > 0 && t < p.releaseTime) {
          const tau = t / p.releaseTime;
          releaseGain = 1 - tau * tau * (3 - 2 * tau);
          flags = 4;
        }
        if (local % this.stepsPerSample === 0) {
          const sampleIdx = local / this.stepsPerSample;
          slot = sampleIdx % HIST_LEN;
          outIdx = audioStepsInBatch.length;
          audioStepsInBatch.push(s);
        }
      }
      u32[o + 1] = flags;
      i32[o + 2] = slot;
      i32[o + 3] = outIdx;
      f32[o + 4] = pushW;
      f32[o + 5] = pushV;
      f32[o + 6] = sigma0;
      f32[o + 7] = p.sigma1;
      f32[o + 8] = releaseGain;
    }
    this.device.queue.writeBuffer(this.stepBuf, 0, this.stepData, 0, steps * this.stepStride);

    // Encode the batch in chunks so that the GPU starts while JS is still encoding.
    const CHUNK = 1024;
    let enc = this.device.createCommandEncoder();
    let pass = enc.beginComputePass();
    const wgX = Math.ceil(this.grid.nx / 16), wgY = Math.ceil(this.grid.ny / 16);
    const wgN = Math.ceil(this.grid.n / 256);
    let cur = this.cur;
    let audioPtr = 0;
    const rbInterval = recording ? RB_INTERVAL : RB_INTERVAL_HELD;
    let stepBound = false;
    for (let s = 0; s < steps; s++) {
      if (s > 0 && s % CHUNK === 0) {
        pass.end();
        this.device.queue.submit([enc.finish()]);
        enc = this.device.createCommandEncoder();
        pass = enc.beginComputePass();
        stepBound = false;
      }
      if (!stepBound) {
        pass.setPipeline(this.stepPipe);
        stepBound = true;
      }
      pass.setBindGroup(0, this.stepBG[cur], [s * this.stepStride]);
      pass.dispatchWorkgroups(wgX, wgY);
      cur = 1 - cur;
      const global = this.totalSteps + s + 1;
      if (global % rbInterval === 0 || s === steps - 1) {
        pass.setPipeline(this.sumsPipe);
        pass.setBindGroup(0, this.rbBG[cur]);
        pass.dispatchWorkgroups(1);
        pass.setPipeline(this.applyPipe);
        pass.setBindGroup(0, this.rbBG[cur]);
        pass.dispatchWorkgroups(wgN);
        stepBound = false;
      }
      if (audioPtr < audioStepsInBatch.length && audioStepsInBatch[audioPtr] === s) {
        pass.setPipeline(this.audioPipe);
        pass.setBindGroup(0, this.audioBG, [s * this.stepStride]);
        pass.dispatchWorkgroups(1);
        audioPtr++;
        stepBound = false;
      }
    }
    pass.setPipeline(this.extractPipe);
    pass.setBindGroup(0, this.extractBG[cur]);
    pass.dispatchWorkgroups(wgN);
    pass.end();
    const staging = this.staging[this.stagingIdx];
    this.stagingIdx = 1 - this.stagingIdx;
    const nAudio = audioStepsInBatch.length;
    enc.copyBufferToBuffer(this.audioOut, 0, staging, 0, Math.max(4, nAudio * 4));
    enc.copyBufferToBuffer(this.frameBuf, 0, staging, this.samplesPerBatch * 4, this.grid.n * 4);
    enc.copyBufferToBuffer(this.statsBuf, 0, staging, (this.samplesPerBatch + this.grid.n) * 4, STATS_LEN * 4);
    this.device.queue.submit([enc.finish()]);
    this.cur = cur;
    this.totalSteps += steps;
    this.stepCount += steps;
    if (!recording) this.prepTime += steps * this.dt;

    await staging.mapAsync(GPUMapMode.READ);
    const view = new Float32Array(staging.getMappedRange());
    const stats = view.subarray(this.samplesPerBatch + this.grid.n, this.samplesPerBatch + this.grid.n + STATS_LEN);
    this.ke = stats[12];
    this.wmax = stats[13];
    if (recording) {
      this.appendAudio(view.slice(0, nAudio));
      this.frames.push(view.slice(this.samplesPerBatch, this.samplesPerBatch + this.grid.n));
    } else {
      this.lastPrepFrame = view.slice(this.samplesPerBatch, this.samplesPerBatch + this.grid.n);
    }
    staging.unmap();

    if (!recording) {
      if (this.phase === "push" && this.prepTime >= p.rampTime) this.phase = "hold";
      const settled = this.phase === "hold" && this.ke < p.settleKE;
      if (settled || this.prepTime >= this.prepDuration()) this.release();
    } else if (this.samples >= this.targetSamples()) {
      this.phase = "done";
    }
  }
  private lastPrepFrame: Float32Array<ArrayBuffer> | null = null;

  private release(): void {
    this.phase = "recording";
    this.stepCount = 0;
    this.totalSteps = 0;
    // Frame 0 is the held, settled shape.
    if (this.lastPrepFrame) this.frames.push(this.lastPrepFrame);
  }

  /** Reads back the full state (debugging and tests). */
  async readState(): Promise<{ pos: Float32Array; vel: Float32Array }> {
    const n = this.grid.n;
    const stg = this.device.createBuffer({ size: 7 * n * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
    const enc = this.device.createCommandEncoder();
    enc.copyBufferToBuffer(this.pos[this.cur], 0, stg, 0, 4 * n * 4);
    enc.copyBufferToBuffer(this.vel[this.cur], 0, stg, 4 * n * 4, 3 * n * 4);
    this.device.queue.submit([enc.finish()]);
    await stg.mapAsync(GPUMapMode.READ);
    const all = new Float32Array(stg.getMappedRange()).slice();
    stg.unmap();
    stg.destroy();
    // Returned as [u | v | w | wLo] and [du | dv | dw].
    return { pos: all.subarray(0, 4 * n), vel: all.subarray(4 * n) };
  }

  destroy(): void {
    for (const b of [...this.pos, ...this.vel, this.hist, this.react, this.lw, this.audioOut, this.frameBuf, this.displayW, this.statsBuf, this.paramsBuf, this.stepBuf, ...this.staging]) b.destroy();
  }

  get fs(): number {
    return this.params.fs;
  }
  get fps(): number {
    return this.params.fps;
  }
  get physParams(): PhysParams {
    return this.params;
  }
}