/**
 * The room recording: every pattern-to-mic impulse response, from one run.
 *
 * The fast render path (src/studio/modal.ts) needs, per coupling pattern —
 * a membrane mode's prof .* rad(phi_mn), a string or bar mode's
 * prof .* rad(phi_m), a string's bridge patch bprof — the air's response at
 * the microphone to an impulse of that pattern. Recording them one run per
 * pattern would cost thousands of runs; discrete reciprocity collapses them
 * to one. The air update's transfer matrix is symmetric — the masked
 * Laplacian in divergence form has w_i·w_j/h² off-diagonal entries and
 * everything else is diagonal — so the response at the mic to pattern g
 * equals the projection onto g of the field one impulse *at the mic*
 * produces. (This needs the uniform sound speed every generated scene has;
 * scripts/modal-poc.ts checks the symmetry to machine precision.) So: run
 * the air alone, kicked at the microphone, and after every step project the
 * field onto every pattern at once.
 *
 * The air runs through the same pipeline as everything else — an air-only .m
 * compiled by the same compiler against the same lapw op, so the fast path's
 * numerics are the real solver's, not a reimplementation. The impulse is the
 * init: kicking the update with source s at step 0 produces
 * p[1] = dt²·wall·s/(1 + sig·dt), which init writes directly.
 *
 * The projections ride each step's submission, the way the Recorder's mic
 * tap does, organized as *pattern groups*: each group contracts the field
 * against one 3D profile over z (the sheet, the tube, the bridge patch),
 * then against per-axis row tables sx and sy — rad of a product of sines
 * factorizes per axis, a line object's y factor is a row of ones, and a
 * fixed pattern like the bridge is a single (ones, ones) pair. Kernels
 * accumulate on the GPU in chunks and come back to the CPU between chunks.
 */
import type { AirGrid } from '../grids.ts';
import { GpuModel } from '../mgpu/model.ts';
import { makeLapw, OpPlan } from '../mgpu/ops.ts';
import { WORKGROUP_SIZE } from '../mgpu/wgsl.ts';

/** The air-only model the recording runs: the generated step's air lines,
 *  with the reciprocal impulse as the initial condition. */
const roomModel = (): string => `% Generated by soundbox for the room recording: the air alone, kicked at
% the microphone. The step is the generated model's air update, verbatim.
function [p, pm] = init(npts, wall, sig, imp, dt)
  p = (dt*dt) * ((wall .* imp) ./ (1 + sig*dt));
  pm = zeros(npts, 1);
end

function [pn, po] = step(p, pm, c, sig, wall, dt)
  lp = lapw(p, wall);
  sd = sig * dt;
  pn = wall .* ((2*p - (1 - sd) .* pm + (c*dt).^2 .* lp) ./ (1 + sd));
  po = p;
end
`;

/** Thrown when a recording is cancelled through `RoomOptions.cancelled`. */
export class RoomCancelled extends Error {
  constructor() {
    super('room recording cancelled');
  }
}

/** One family of patterns sharing a 3D profile: each pair (xi, yi) is the
 *  pattern prof(x,y,z) · sx[xi](ix) · sy[yi](iy), and produces one kernel
 *  row. */
export interface RoomPatternGroup {
  label: string;
  /** The 3D profile, air-sized; contracted over z first. */
  prof: Float32Array;
  /** Row tables: sx is sxRows × nx, sy is syRows × ny. */
  sx: Float32Array;
  sxRows: number;
  sy: Float32Array;
  syRows: number;
  /** The (sx row, sy row) of each kernel row this group records. */
  pairs: { xi: number; yi: number }[];
}

export interface RoomOptions {
  device: GPUDevice;
  air: AirGrid;
  /** The scene's evaluated medium. */
  fields: { c: Float32Array; sig: Float32Array; wall: Float32Array };
  dt: number;
  /** Flat air index of the microphone cell. */
  micIndex: number;
  groups: RoomPatternGroup[];
  /** Kernel length, steps. */
  steps: number;
  onProgress?: (done: number, total: number) => void;
  /** Polled between batches; returning true abandons the recording (the
   *  buffers are torn down and RoomCancelled is thrown). */
  cancelled?: () => boolean;
}

export interface Room {
  /** One row of `steps` samples per pattern, row-major: the groups'
   *  `pairs`, concatenated in group order. */
  kernels: Float32Array;
  steps: number;
  rows: number;
}

const planeShader = (air: AirGrid): string => `
@group(0) @binding(0) var<storage, read_write> plane: array<f32>;
@group(0) @binding(1) var<storage, read> field: array<f32>;
@group(0) @binding(2) var<storage, read> prof: array<f32>;
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let c = gid.x;
  if (c >= ${air.nx * air.ny}u) { return; }
  var s = 0.0;
  for (var iz = 0u; iz < ${air.nz}u; iz++) {
    let i = c + ${air.nx * air.ny}u * iz;
    s += prof[i] * field[i];
  }
  plane[c] = s;
}
`;

const xcShader = (air: AirGrid, rows: number): string => `
@group(0) @binding(0) var<storage, read_write> tmp: array<f32>;
@group(0) @binding(1) var<storage, read> plane: array<f32>;
@group(0) @binding(2) var<storage, read> sx: array<f32>;
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = gid.x;
  if (k >= ${rows * air.ny}u) { return; }
  let m = k % ${rows}u;
  let iy = k / ${rows}u;
  var s = 0.0;
  for (var ix = 0u; ix < ${air.nx}u; ix++) {
    s += sx[m * ${air.nx}u + ix] * plane[ix + ${air.nx}u * iy];
  }
  tmp[k] = s;
}
`;

const projShader = (
  air: AirGrid,
  sxRows: number,
  nPairs: number,
  rowOffset: number,
  totalRows: number,
  chunk: number,
): string => `
@group(0) @binding(0) var<storage, read_write> chunkBuf: array<f32>;
@group(0) @binding(1) var<storage, read> tmp: array<f32>;
@group(0) @binding(2) var<storage, read> sy: array<f32>;
@group(0) @binding(3) var<storage, read> pairs: array<u32>; // xi | yi<<16
@group(0) @binding(4) var<storage, read> head: array<u32>;
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let q = gid.x;
  if (q >= ${nPairs}u) { return; }
  let mn = pairs[q];
  let xi = mn & 0xffffu;
  let yi = mn >> 16u;
  var s = 0.0;
  for (var iy = 0u; iy < ${air.ny}u; iy++) {
    s += sy[yi * ${air.ny}u + iy] * tmp[xi + ${sxRows}u * iy];
  }
  let i = head[0];
  if (i < ${chunk}u) {
    chunkBuf[i * ${totalRows}u + ${rowOffset}u + q] = s;
  }
}
`;

const bumpShader = `
@group(0) @binding(0) var<storage, read_write> head: array<u32>;
@compute @workgroup_size(1)
fn main() {
  head[0] = head[0] + 1u;
}
`;

/** Run the reciprocal recording and return every pattern's kernel. */
export async function recordRoom(opts: RoomOptions): Promise<Room> {
  const { device, air, fields, dt, micIndex, groups, steps } = opts;
  const rows = groups.reduce((a, g) => a + g.pairs.length, 0);
  if (rows === 0) return { kernels: new Float32Array(0), steps, rows };
  for (const g of groups) {
    for (const p of g.pairs) {
      if (p.xi < 0 || p.xi >= g.sxRows || p.yi < 0 || p.yi >= g.syRows) {
        throw new Error(`room group '${g.label}': pair (${p.xi},${p.yi}) outside its tables`);
      }
    }
  }

  // The air-only model, through the ordinary pipeline.
  const imp = new Float32Array(air.npts);
  imp[micIndex] = 1;
  const gpu = await GpuModel.create({
    device,
    ops: new OpPlan(device, [makeLapw('lapw', air)]),
    consts: { npts: air.npts, nx: air.nx, ny: air.ny, nz: air.nz },
    fields: { c: fields.c, sig: fields.sig, wall: fields.wall, imp },
    source: roomModel(),
    paramNames: [],
    state: [
      { name: 'p', count: air.npts },
      { name: 'pm', count: air.npts },
    ],
  });
  gpu.setDt(dt);
  gpu.setParams({});

  // Kernels come back in chunks of whole samples; each sample is one f32 per
  // pattern row. Cap a chunk at 64 MB.
  const chunk = Math.max(1, Math.min(steps, Math.floor((64 << 20) / (4 * rows))));
  const owned: GPUBuffer[] = [];
  const mk = (label: string, size: number): GPUBuffer => {
    const b = device.createBuffer({
      label: `room-${label}`,
      size,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
    });
    owned.push(b);
    return b;
  };
  try {
    const headBuf = mk('head', 4);
    const chunkBuf = mk('chunk', 4 * rows * chunk);
    const readback = device.createBuffer({
      label: 'room-readback',
      size: 4 * rows * chunk,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });
    owned.push(readback);
    device.queue.writeBuffer(headBuf, 0, new Uint32Array([0]));

    const pipe = async (code: string, label: string): Promise<GPUComputePipeline> =>
      device.createComputePipelineAsync({
        label: `room-${label}`,
        layout: 'auto',
        compute: {
          module: device.createShaderModule({ code, label: `room-${label}` }),
          entryPoint: 'main',
        },
      });
    const bind = (p: GPUComputePipeline, buffers: GPUBuffer[]): GPUBindGroup =>
      device.createBindGroup({
        layout: p.getBindGroupLayout(0),
        entries: buffers.map((buffer, i) => ({ binding: i, resource: { buffer } })),
      });
    const upload = (label: string, data: Float32Array | Uint32Array): GPUBuffer => {
      const b = mk(label, data.byteLength);
      device.queue.writeBuffer(b, 0, data as Float32Array<ArrayBuffer>);
      return b;
    };

    const pField = gpu.stateBuffer('p')!;
    const planePipe = await pipe(planeShader(air), 'plane');
    const bumpPipe = await pipe(bumpShader, 'bump');
    const bumpBind = bind(bumpPipe, [headBuf]);

    interface GroupRun {
      xcPipe: GPUComputePipeline;
      projPipe: GPUComputePipeline;
      planeBind: GPUBindGroup;
      xcBind: GPUBindGroup;
      projBind: GPUBindGroup;
      xcThreads: number;
      nPairs: number;
    }
    const runs: GroupRun[] = [];
    let rowOffset = 0;
    for (const g of groups) {
      const planeBuf = mk(`plane-${g.label}`, 4 * air.nx * air.ny);
      const tmpBuf = mk(`tmp-${g.label}`, 4 * g.sxRows * air.ny);
      const profBuf = upload(`prof-${g.label}`, g.prof);
      const sxBuf = upload(`sx-${g.label}`, g.sx);
      const syBuf = upload(`sy-${g.label}`, g.sy);
      const pairsBuf = upload(
        `pairs-${g.label}`,
        new Uint32Array(g.pairs.map((p) => p.xi | (p.yi << 16))),
      );
      const xcPipe = await pipe(xcShader(air, g.sxRows), `xc-${g.label}`);
      const projPipe = await pipe(
        projShader(air, g.sxRows, g.pairs.length, rowOffset, rows, chunk),
        `proj-${g.label}`,
      );
      runs.push({
        xcPipe,
        projPipe,
        planeBind: bind(planePipe, [planeBuf, pField, profBuf]),
        xcBind: bind(xcPipe, [tmpBuf, planeBuf, sxBuf]),
        projBind: bind(projPipe, [chunkBuf, tmpBuf, syBuf, pairsBuf, headBuf]),
        xcThreads: g.sxRows * air.ny,
        nPairs: g.pairs.length,
      });
      rowOffset += g.pairs.length;
    }

    const encodeProjection = (enc: GPUCommandEncoder): void => {
      const pass = enc.beginComputePass({ label: 'room-projection' });
      const go = (p: GPUComputePipeline, b: GPUBindGroup, threads: number): void => {
        pass.setPipeline(p);
        pass.setBindGroup(0, b);
        pass.dispatchWorkgroups(Math.ceil(threads / WORKGROUP_SIZE));
      };
      for (const r of runs) {
        go(planePipe, r.planeBind, air.nx * air.ny);
        go(r.xcPipe, r.xcBind, r.xcThreads);
        go(r.projPipe, r.projBind, r.nPairs);
      }
      go(bumpPipe, bumpBind, 1);
      pass.end();
    };

    // Sample 0 is the projection of the post-init field (the kicked air, one
    // step in); every later sample rides its step's own submission.
    const kernels = new Float32Array(rows * steps);
    gpu.init();
    {
      const enc = device.createCommandEncoder({ label: 'room-first-sample' });
      encodeProjection(enc);
      device.queue.submit([enc.finish()]);
    }
    let got = 0; // samples already on the CPU
    let pending = 1; // samples sitting in the chunk buffer
    const flush = async (): Promise<void> => {
      const enc = device.createCommandEncoder({ label: 'room-flush' });
      enc.copyBufferToBuffer(chunkBuf, 0, readback, 0, 4 * rows * pending);
      device.queue.submit([enc.finish()]);
      await readback.mapAsync(GPUMapMode.READ, 0, 4 * rows * pending);
      const data = new Float32Array(readback.getMappedRange(0, 4 * rows * pending));
      // Chunk samples are row-minor; kernels are row-major.
      for (let s = 0; s < pending; s++)
        for (let q = 0; q < rows; q++) kernels[q * steps + got + s] = data[s * rows + q];
      readback.unmap();
      device.queue.writeBuffer(headBuf, 0, new Uint32Array([0]));
      got += pending;
      pending = 0;
    };
    // Steps per submission, as the plain renderer batches; the sync between
    // batches keeps the queue from running unboundedly ahead, and is what
    // paces the progress reports and the cancellation checks.
    const batch = 2048;
    while (got + pending < steps) {
      if (opts.cancelled?.()) throw new RoomCancelled();
      const n = Math.min(batch, chunk - pending, steps - got - pending);
      gpu.step(n, encodeProjection);
      pending += n;
      await device.queue.onSubmittedWorkDone();
      opts.onProgress?.(got + pending, steps);
      if (pending === chunk) await flush();
    }
    if (pending > 0) await flush();
    return { kernels, steps, rows };
  } finally {
    for (const b of owned) b.destroy();
    gpu.destroy();
  }
}