// WebGPU time-stepping driver of the full (geometrically nonlinear)
// simulation. The run advances in batches of one movie frame (1/fps s). A
// batch encodes one compute dispatch per time step (with a per-step uniform
// addressed by dynamic offset), one retarded-time listener sum over the
// batch's audio samples, and one extraction of the transverse displacement
// at the end. Audio, frame, and stick state are read back after every batch;
// one batch is kept encoded ahead of the one being read back so that the
// GPU never waits for JavaScript. The recording starts at the impact, with
// the plate at rest in its stress-free shape.

import {
  type Plate,
  type SimParams,
  CFL,
  HIST_MARGIN,
  massProperties,
  strikePatch,
  strikePoint,
  thicknessRange,
  nodeWeights,
  listenerTerms,
} from "../params.ts";
import { type FullLoss, maxAngularFrequency } from "../model/reference.ts";
import commonWgsl from "./shaders/common.wgsl?raw";
import stepWgsl from "./shaders/step.wgsl?raw";
import audioWgsl from "./shaders/audio.wgsl?raw";
import extractWgsl from "./shaders/extract.wgsl?raw";

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

export interface StickState {
  x: number;
  v: number;
  force: number;
  maxForce: number;
  withdrawn: boolean;
}

export interface SimStatus {
  phase: Phase;
  /** Simulation time since the impact of the state read back so far [s]. */
  t: number;
  totalSteps: number;
  samples: number;
  frames: number;
  peakPressure: number;
  /** Largest transverse displacement seen in the frames [m]. */
  wmax: number;
  delayClamped: boolean;
  stick: StickState;
}

const PARAMS_LEN = 40;
const STEP_LEN = 12;
const STICK_LEN = 16;
/** Nodes per workgroup along each axis in step.wgsl. */
export const TILE = 12;

interface Pending {
  staging: GPUBuffer;
  steps: number;
  nAudio: number;
  stickParity: number;
}

export class Simulation {
  readonly plate: Plate;
  readonly dt: number;
  readonly stepsPerSample: number;
  readonly samplesPerBatch: number;
  readonly stepsPerBatch: number;
  readonly omegaMax: number;
  readonly histLen: number;
  readonly mass: number;

  audio: Float32Array = new Float32Array(0);
  samples = 0;
  /** Movie frames of w [m]; frame k is at t = k / fps. */
  frames: Float32Array[] = [];
  phase: Phase = "recording";
  totalSteps = 0;
  peakPressure = 0;
  wmax = 0;
  delayClamped = false;

  private stick: StickState;
  private inFlight: Pending | null = null;
  private device: GPUDevice;
  private params: SimParams;
  private loss: FullLoss;
  private paramsBuf: GPUBuffer;
  private stepBuf: GPUBuffer;
  private stepStride: number;
  private stepData: ArrayBuffer;
  private pos: GPUBuffer[];
  private vel: GPUBuffer[];
  private hist: GPUBuffer;
  private stickBuf: GPUBuffer;
  private nd: GPUBuffer;
  private lw: GPUBuffer;
  private audioOut: GPUBuffer;
  private frameBuf: GPUBuffer;
  private staging: GPUBuffer[];
  private stagingIdx = 0;
  private cur = 0;
  private stepPipe: GPUComputePipeline;
  private audioPipe: GPUComputePipeline;
  private extractPipe: GPUComputePipeline;
  private stepBG: GPUBindGroup[];
  private audioBG: GPUBindGroup;
  private extractBG: GPUBindGroup[];

  constructor(device: GPUDevice, plate: Plate, params: SimParams, loss: FullLoss) {
    this.device = device;
    this.plate = plate;
    this.params = { ...params };
    this.loss = loss;
    const g = plate.grid, m = plate.mat;

    // Stable time step, rounded so that an integer number of steps make one audio sample.
    this.omegaMax = maxAngularFrequency(g, m, thicknessRange(plate).max);
    const dtMax = (CFL * 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;
    this.histLen = this.samplesPerBatch + HIST_MARGIN;

    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)];
    this.vel = [f32(3 * n), f32(3 * n)];
    this.hist = f32(this.histLen * n);
    this.stickBuf = f32(STICK_LEN);
    const weights = nodeWeights(g);
    const packed = new Float32Array(4 * n);
    for (let k = 0; k < n; k++) {
      packed[4 * k] = plate.w0[k]; packed[4 * k + 1] = weights[k]; packed[4 * k + 2] = plate.h[k];
    }
    this.nd = f32(4 * n);
    device.queue.writeBuffer(this.nd, 0, packed);
    this.lw = f32(2 * n);
    this.audioOut = f32(this.samplesPerBatch);
    this.frameBuf = f32(n);
    const stagingSize = (this.samplesPerBatch + n + STICK_LEN) * 4;
    this.staging = [0, 1].map(() => device.createBuffer({ size: stagingSize, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }));

    const { mass, Ix, Iy } = massProperties(plate);
    this.mass = mass;
    const cp = strikePatch(plate, params);
    let i0 = g.nx, i1 = 0, j0 = g.ny, j1 = 0;
    for (const k of cp.nodes) {
      const i = k % g.nx, j = Math.floor(k / g.nx);
      i0 = Math.min(i0, i); i1 = Math.max(i1, i); j0 = Math.min(j0, j); j1 = Math.max(j1, j);
    }
    const stick = new Float32Array(STICK_LEN);
    stick[1] = -params.strikeVelocity; stick[9] = -params.strikeVelocity;
    device.queue.writeBuffer(this.stickBuf, 0, stick);
    this.stick = { x: 0, v: -params.strikeVelocity, force: 0, maxForce: 0, withdrawn: false };

    const sp = strikePoint(plate, params);
    const pf = new Float32Array(PARAMS_LEN);
    const pu = new Uint32Array(pf.buffer);
    pu[0] = g.nx; pu[1] = g.ny; pu[2] = n; pu[3] = this.histLen;
    pf[4] = g.hx; pf[5] = g.hy; pf[6] = this.dt; pf[7] = m.nu;
    pf[8] = m.E / (12 * (1 - m.nu * m.nu)); pf[9] = m.E / (1 - m.nu * m.nu); pf[10] = m.E / (2 * (1 + m.nu)); pf[11] = m.rho;
    pf[12] = mass; pf[13] = Ix; pf[14] = Iy; pf[15] = params.fs;
    pf[16] = sp.x; pf[17] = sp.y; pf[18] = params.contactRadius ** 2; pf[19] = cp.mass;
    pu[20] = cp.nearest; pf[21] = params.contactStiffness; pf[22] = params.contactExponent; pf[23] = params.stickMass;
    pu[24] = i0; pu[25] = i1; pu[26] = j0; pu[27] = j1;
    pf[28] = cp.cx; pf[29] = cp.cy;
    this.paramsBuf = device.createBuffer({ size: pf.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(this.paramsBuf, 0, pf);

    this.stepStride = Math.max(STEP_LEN * 4, 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 shader = (code: string) => device.createShaderModule({ code: commonWgsl + code });
    const entry = (binding: number, type: GPUBufferBindingType, hasDynamicOffset = false): GPUBindGroupLayoutEntry => ({
      binding, visibility: GPUShaderStage.COMPUTE, buffer: { type, hasDynamicOffset },
    });
    const pipeline = (module: GPUShaderModule, layout: GPUBindGroupLayout) =>
      device.createComputePipeline({ layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }), compute: { module, entryPoint: "main" } });
    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"), entry(8, "read-only-storage")],
    });
    this.stepPipe = pipeline(shader(stepWgsl), 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"), entry(8, "read-only-storage")],
    });
    this.audioPipe = pipeline(shader(audioWgsl), audioLayout);
    const extractLayout = device.createBindGroupLayout({ entries: [entry(0, "uniform"), entry(1, "read-only-storage"), entry(2, "storage"), entry(8, "read-only-storage")] });
    this.extractPipe = pipeline(shader(extractWgsl), extractLayout);

    const spb = { buffer: this.stepBuf, size: STEP_LEN * 4 };
    this.stepBG = [0, 1].map((p) =>
      device.createBindGroup({
        layout: stepLayout,
        entries: [
          { binding: 0, resource: { buffer: this.paramsBuf } },
          { binding: 1, resource: spb },
          { 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.stickBuf } },
          { binding: 8, resource: { buffer: this.nd } },
        ],
      }),
    );
    this.audioBG = device.createBindGroup({
      layout: audioLayout,
      entries: [
        { binding: 0, resource: { buffer: this.paramsBuf } },
        { binding: 1, resource: spb },
        { binding: 2, resource: { buffer: this.hist } },
        { binding: 3, resource: { buffer: this.lw } },
        { binding: 4, resource: { buffer: this.audioOut } },
        { binding: 8, resource: { buffer: this.nd } },
      ],
    });
    this.extractBG = [0, 1].map((p) =>
      device.createBindGroup({
        layout: extractLayout,
        entries: [
          { binding: 0, resource: { buffer: this.paramsBuf } },
          { binding: 1, resource: { buffer: this.pos[p] } },
          { binding: 2, resource: { buffer: this.frameBuf } },
          { binding: 8, resource: { buffer: this.nd } },
        ],
      }),
    );
    this.setListener(params);
    this.frames.push(new Float32Array(n));
  }

  /** Recomputes the per-node weights and delays of the listener sum; takes effect on later samples. */
  setListener(p: SimParams): void {
    const g = this.plate.grid;
    this.params.listenerDist = p.listenerDist; this.params.listenerAngle = p.listenerAngle; this.params.listenerAzimuth = p.listenerAzimuth;
    const { weight, delay } = listenerTerms(this.plate, this.params);
    const data = new Float32Array(2 * g.n);
    let dmin = Infinity;
    for (let k = 0; k < g.n; k++) if (weight[k] > 0) dmin = Math.min(dmin, delay[k]);
    this.delayClamped = false;
    for (let k = 0; k < g.n; k++) {
      data[2 * k] = weight[k];
      let d = weight[k] > 0 ? (delay[k] - dmin) * this.params.fs : 0;
      if (d > HIST_MARGIN - 3) { d = HIST_MARGIN - 3; this.delayClamped = true; }
      data[2 * k + 1] = d;
    }
    this.device.queue.writeBuffer(this.lw, 0, data);
  }

  setDuration(duration: number): void {
    this.params.duration = duration;
    if (this.canContinue()) this.phase = "recording";
  }
  canContinue(): boolean {
    return this.phase === "done" && this.samples < this.targetSamples();
  }
  stop(): void {
    this.phase = "done";
  }
  targetSamples(): number {
    return Math.round(this.params.duration * this.params.fs);
  }
  status(): SimStatus {
    const ahead = this.inFlight ? this.inFlight.steps * this.dt : 0;
    return { phase: this.phase, t: this.totalSteps * this.dt - ahead, totalSteps: this.totalSteps, samples: this.samples, frames: this.frames.length, peakPressure: this.peakPressure, wmax: this.wmax, delayClamped: this.delayClamped, stick: this.stick };
  }
  recordedAudio(): Float32Array {
    return this.audio.subarray(0, this.samples);
  }
  private stepWorkgroups(): [number, number] {
    return [Math.ceil(this.plate.grid.nx / TILE), Math.ceil(this.plate.grid.ny / TILE)];
  }
  private appendAudio(raw: Float32Array): void {
    const need = this.samples + raw.length;
    if (this.audio.length < need) {
      const a = new Float32Array(Math.max(need, 2 * this.audio.length, this.params.fs));
      a.set(this.audio.subarray(0, this.samples));
      this.audio = a;
    }
    this.audio.set(raw, this.samples);
    for (let k = 0; k < raw.length; k++) this.peakPressure = Math.max(this.peakPressure, Math.abs(raw[k]));
    this.samples = need;
  }
  private writeStepParams(s: number, fields: { flags: number; slot: number; outIdx: number }): void {
    const o = (s * this.stepStride) / 4;
    const u32 = new Uint32Array(this.stepData), i32 = new Int32Array(this.stepData), f32 = new Float32Array(this.stepData);
    u32[o] = this.totalSteps + s;
    u32[o + 1] = fields.flags;
    i32[o + 2] = fields.slot;
    i32[o + 3] = fields.outIdx;
    f32[o + 4] = this.loss.sigma0;
    f32[o + 5] = this.loss.sigma1;
    f32[o + 6] = this.loss.sigmaIn;
    u32[o + 7] = (this.totalSteps + s) & 1;
  }

  /** 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 cur = this.inFlight ?? this.encodeBatch();
    this.inFlight = null;
    if (!cur) { this.phase = "done"; return; }
    if (this.stepsLeft() > 0) this.inFlight = this.encodeBatch();
    await this.collect(cur);
  }
  private stepsLeft(): number {
    return this.targetSamples() * this.stepsPerSample - this.totalSteps;
  }
  private encodeBatch(): Pending | null {
    const steps = Math.min(this.stepsPerBatch, Math.max(0, this.stepsLeft()));
    if (steps === 0) return null;
    const audioSteps: number[] = [];
    for (let s = 0; s < steps; s++) {
      const local = this.totalSteps + s;
      let slot = -1, outIdx = -1;
      if (local % this.stepsPerSample === 0) {
        slot = (local / this.stepsPerSample) % this.histLen;
        outIdx = audioSteps.length;
        audioSteps.push(s);
      }
      this.writeStepParams(s, { flags: 1, slot, outIdx });
    }
    this.device.queue.writeBuffer(this.stepBuf, 0, this.stepData, 0, steps * this.stepStride);
    const CHUNK = 1024;
    let enc = this.device.createCommandEncoder();
    let pass = enc.beginComputePass();
    const [wgX, wgY] = this.stepWorkgroups();
    const n = this.plate.grid.n;
    let cur = this.cur;
    pass.setPipeline(this.stepPipe);
    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();
        pass.setPipeline(this.stepPipe);
      }
      pass.setBindGroup(0, this.stepBG[cur], [s * this.stepStride]);
      pass.dispatchWorkgroups(wgX, wgY);
      cur = 1 - cur;
    }
    const nAudio = audioSteps.length;
    if (nAudio > 0) {
      pass.setPipeline(this.audioPipe);
      pass.setBindGroup(0, this.audioBG, [audioSteps[0] * this.stepStride]);
      pass.dispatchWorkgroups(nAudio);
    }
    pass.setPipeline(this.extractPipe);
    pass.setBindGroup(0, this.extractBG[cur]);
    pass.dispatchWorkgroups(Math.ceil(n / 256));
    pass.end();
    const staging = this.staging[this.stagingIdx];
    this.stagingIdx = 1 - this.stagingIdx;
    enc.copyBufferToBuffer(this.audioOut, 0, staging, 0, Math.max(4, nAudio * 4));
    enc.copyBufferToBuffer(this.frameBuf, 0, staging, this.samplesPerBatch * 4, n * 4);
    enc.copyBufferToBuffer(this.stickBuf, 0, staging, (this.samplesPerBatch + n) * 4, STICK_LEN * 4);
    this.device.queue.submit([enc.finish()]);
    this.cur = cur;
    this.totalSteps += steps;
    return { staging, steps, nAudio, stickParity: this.totalSteps & 1 };
  }
  private async collect(b: Pending): Promise<void> {
    const n = this.plate.grid.n;
    await b.staging.mapAsync(GPUMapMode.READ);
    const view = new Float32Array(b.staging.getMappedRange());
    const base = this.samplesPerBatch + n + 8 * b.stickParity;
    this.stick = { x: view[base], v: view[base + 1], force: view[base + 2], maxForce: view[base + 3], withdrawn: view[base + 5] !== 0 };
    this.appendAudio(view.slice(0, b.nAudio));
    const frame = view.slice(this.samplesPerBatch, this.samplesPerBatch + n);
    for (let k = 0; k < n; k++) this.wmax = Math.max(this.wmax, Math.abs(frame[k]));
    this.frames.push(frame);
    b.staging.unmap();
    if (this.samples >= this.targetSamples()) this.phase = "done";
  }

  // ---------- Testing support ----------

  /** Advances the state by exactly one free step (no drumstick, no audio) with the losses. */
  async stepOnce(): Promise<void> {
    this.writeStepParams(0, { flags: 0, slot: -1, outIdx: -1 });
    this.device.queue.writeBuffer(this.stepBuf, 0, this.stepData, 0, this.stepStride);
    const enc = this.device.createCommandEncoder();
    const pass = enc.beginComputePass();
    pass.setPipeline(this.stepPipe);
    pass.setBindGroup(0, this.stepBG[this.cur], [0]);
    pass.dispatchWorkgroups(...this.stepWorkgroups());
    pass.end();
    this.device.queue.submit([enc.finish()]);
    this.cur = 1 - this.cur;
    this.totalSteps++;
    await this.device.queue.onSubmittedWorkDone();
  }
  /** Overwrites the state: pos as [u | v | w | wLo], vel as [du | dv | dw]. */
  writeState(pos: Float32Array<ArrayBuffer>, vel: Float32Array<ArrayBuffer>): void {
    const n = this.plate.grid.n;
    if (pos.length !== 4 * n || vel.length !== 3 * n) throw new Error("state arrays have the wrong length");
    this.device.queue.writeBuffer(this.pos[this.cur], 0, pos);
    this.device.queue.writeBuffer(this.vel[this.cur], 0, vel);
  }
  /** Reads back the full state as [u | v | w | wLo] and [du | dv | dw]. */
  async readState(): Promise<{ pos: Float32Array; vel: Float32Array }> {
    const n = this.plate.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();
    return { pos: all.subarray(0, 4 * n), vel: all.subarray(4 * n) };
  }

  destroy(): void {
    this.inFlight = null;
    for (const b of [...this.pos, ...this.vel, this.hist, this.stickBuf, this.nd, this.lw, this.audioOut, this.frameBuf, this.paramsBuf, this.stepBuf, ...this.staging]) b.destroy();
  }
  get fs(): number { return this.params.fs; }
  get fps(): number { return this.params.fps; }
  get simParams(): SimParams { return this.params; }
}