// Re-render the listener signals at any position from the acceleration record, with the same
// free-field model as cymsim's built-in listener (cymsimweb src/model/radiation.ts): a Rayleigh
// integral over element-mean normal accelerations for the front path and a diffracted path over
// the rim for the back term, each element delayed by its distance at a 16-tap windowed-sinc
// fractional delay. The record already holds the audio-rate node accelerations, so this is a
// gather over elements per output sample; cymsim's test suite validates the same reconstruction
// against the built-in listener to 1.5e-4 relative rms. The gather runs on WebGPU when available
// (one workgroup per output sample, the record streamed through a sliding window) and on the CPU
// otherwise. besselI0 and fractionalDelayTaps are copied from cymsimweb.
import type { CymsimRun } from "./io/cymsimFile.ts";

export const DELAY_TAPS = 16;

/** Modified Bessel function I0 by its power series (converges quickly for x < 30). */
function besselI0(x: number): number {
  let sum = 1, term = 1;
  const y = (x * x) / 4;
  for (let k = 1; k < 60; k++) { term *= y / (k * k); sum += term; if (term < 1e-17 * sum) break; }
  return sum;
}

/** Windowed-sinc taps that read x(j - i0 - fr), lags i0 - T/2 + 1 ... i0 + T/2 (Kaiser, beta 6, unit sum). */
function fractionalDelayTaps(fr: number, T = DELAY_TAPS): Float64Array {
  const taps = new Float64Array(T), half = T / 2, i0 = besselI0(6);
  let sum = 0;
  for (let t = 0; t < T; t++) {
    const k = t - half + 1, x = k - fr;
    const sinc = x === 0 ? 1 : Math.sin(Math.PI * x) / (Math.PI * x);
    const u = x / half, win = besselI0(6 * Math.sqrt(Math.max(1 - u * u, 0))) / i0;
    taps[t] = sinc * win; sum += taps[t];
  }
  for (let t = 0; t < T; t++) taps[t] /= sum;
  return taps;
}

interface Geometry {
  E: number;
  coefs: Float32Array<ArrayBuffer>;   // [2e] front rho A / (2 pi r), [2e+1] back with the diffraction gain
  delays: Uint32Array<ArrayBuffer>;   // [2e], [2e+1]: integer delays in audio samples
  taps: Float32Array<ArrayBuffer>;    // per element 32 floats: 16 front taps then 16 back taps
  maxLag: number;        // largest lag any tap reaches back, in samples
}

/** Per-element coefficients, delays and interpolator taps for a listener at ``position``. */
function buildGeometry(run: CymsimRun, position: [number, number, number]): Geometry {
  const m = run.mesh, E = m.nElements, [Lx, Ly, Lz] = position;
  const { c_air: c, rho_air: rho, back_diffraction_gain: g } = run.config.listener;
  // Diffraction ring: the boundary loop, subdivided in polar coordinates to about 720 points so
  // that the min-over-ring path length matches cymsimweb's analytic rim (which uses nRim = 720).
  const nb = m.boundary.length, sub = Math.max(1, Math.ceil(720 / nb));
  const rim: number[][] = [];
  for (let k = 0; k < nb; k++) {
    const a = m.boundary[k], b = m.boundary[(k + 1) % nb];
    const ra = Math.hypot(m.xy[2 * a], m.xy[2 * a + 1]), rb = Math.hypot(m.xy[2 * b], m.xy[2 * b + 1]);
    const pa = Math.atan2(m.xy[2 * a + 1], m.xy[2 * a]);
    let pb = Math.atan2(m.xy[2 * b + 1], m.xy[2 * b]);
    if (pb - pa > Math.PI) pb -= 2 * Math.PI;
    if (pa - pb > Math.PI) pb += 2 * Math.PI;
    for (let s = 0; s < sub; s++) {
      const t = s / sub, r = ra + t * (rb - ra), phi = pa + t * (pb - pa), z = m.z0[a] + t * (m.z0[b] - m.z0[a]);
      rim.push([r * Math.cos(phi), r * Math.sin(phi), z]);
    }
  }
  const rimToL = rim.map(([x, y, z]) => Math.hypot(x - Lx, y - Ly, z - Lz));
  const coefs = new Float32Array(2 * E), delays = new Uint32Array(2 * E), taps = new Float32Array(32 * E);
  let maxLag = 0;
  for (let e = 0; e < E; e++) {
    const a = m.triangles[3 * e], b = m.triangles[3 * e + 1], d = m.triangles[3 * e + 2];
    const cx = (m.xy[2 * a] + m.xy[2 * b] + m.xy[2 * d]) / 3;
    const cy = (m.xy[2 * a + 1] + m.xy[2 * b + 1] + m.xy[2 * d + 1]) / 3;
    const cz = (m.z0[a] + m.z0[b] + m.z0[d]) / 3;
    const rF = Math.hypot(Lx - cx, Ly - cy, Lz - cz);
    let rB = Infinity;
    for (let k = 0; k < rim.length; k++) {
      rB = Math.min(rB, Math.hypot(cx - rim[k][0], cy - rim[k][1], cz - rim[k][2]) + rimToL[k]);
    }
    coefs[2 * e] = (rho * m.area[e]) / (2 * Math.PI * rF);
    coefs[2 * e + 1] = (g * rho * m.area[e]) / (2 * Math.PI * rB);
    const dF = (rF / c) * run.fs, dB = (rB / c) * run.fs;
    delays[2 * e] = Math.floor(dF); delays[2 * e + 1] = Math.floor(dB);
    taps.set(fractionalDelayTaps(dF - delays[2 * e]), 32 * e);
    taps.set(fractionalDelayTaps(dB - delays[2 * e + 1]), 32 * e + 16);
    maxLag = Math.max(maxLag, delays[2 * e], delays[2 * e + 1]);
  }
  return { E, coefs, delays, taps, maxLag: maxLag + DELAY_TAPS / 2 + 1 };
}

/**
 * Element means of the decoded record for global samples [j0, j1), sample-major (E per row);
 * rows for samples before 0 are zero (nothing has radiated yet).
 */
function elementMeans(run: CymsimRun, j0: number, j1: number): Float32Array<ArrayBuffer> {
  const acc = run.acceleration!, m = run.mesh, E = m.nElements, N = acc.nodes;
  const out = new Float32Array((j1 - j0) * E);
  const d0 = Math.max(0, j0);
  if (d0 >= j1) return out;
  const rows = acc.read(d0, j1);
  for (let j = d0; j < j1; j++) {
    const src = (j - d0) * N, dst = (j - j0) * E;
    for (let e = 0; e < E; e++) {
      const a = m.triangles[3 * e], b = m.triangles[3 * e + 1], c = m.triangles[3 * e + 2];
      out[dst + e] = (rows[src + a] + rows[src + b] + rows[src + c]) / 3;
    }
  }
  return out;
}

export interface RenderedListener { front: Float32Array; back: Float32Array; position: [number, number, number]; backend: "gpu" | "cpu" }
export interface RenderOptions { onProgress?: (fraction: number, backend: "gpu" | "cpu") => void; shouldStop?: () => boolean }

const GATHER_WGSL = /* wgsl */ `
struct U { E: u32, D: u32, chunkLen: u32, base: u32 }
@group(0) @binding(0) var<uniform> u: U;
@group(0) @binding(1) var<storage, read> coefs: array<f32>;
@group(0) @binding(2) var<storage, read> delays: array<u32>;
@group(0) @binding(3) var<storage, read> taps: array<f32>;
@group(0) @binding(4) var<storage, read> win: array<f32>;      // (D + chunkLen) rows of E means
@group(0) @binding(5) var<storage, read_write> outp: array<f32>; // 2 per sample: front, back

const WG = 256u;
var<workgroup> pf: array<f32, WG>;
var<workgroup> pb: array<f32, WG>;

@compute @workgroup_size(WG)
fn gather(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
  let s = wg.x; let tid = li.x;
  var f = 0.0; var b = 0.0;
  if (s < u.chunkLen) {
    let jg = i32(u.base + s);
    for (var e = tid; e < u.E; e += WG) {
      let i0f = i32(delays[2u * e]); let i0b = i32(delays[2u * e + 1u]);
      var af = 0.0; var ab = 0.0;
      for (var t = 0; t < 16; t++) {
        let lagF = i0f + t - 7;
        if (lagF >= 0 && lagF <= jg) {
          af += taps[e * 32u + u32(t)] * win[u32(i32(u.D + s) - lagF) * u.E + e];
        }
        let lagB = i0b + t - 7;
        if (lagB >= 0 && lagB <= jg) {
          ab += taps[e * 32u + 16u + u32(t)] * win[u32(i32(u.D + s) - lagB) * u.E + e];
        }
      }
      f += coefs[2u * e] * af; b += coefs[2u * e + 1u] * ab;
    }
  }
  pf[tid] = f; pb[tid] = b;
  workgroupBarrier();
  var step = WG / 2u;
  while (step > 0u) {
    if (tid < step) { pf[tid] += pf[tid + step]; pb[tid] += pb[tid + step]; }
    workgroupBarrier();
    step = step / 2u;
  }
  if (tid == 0u && s < u.chunkLen) { outp[2u * s] = pf[0]; outp[2u * s + 1u] = pb[0]; }
}
`;

let devicePromise: Promise<GPUDevice | null> | null = null;
function gpuDevice(): Promise<GPUDevice | null> {
  devicePromise ??= (async () => {
    try {
      const adapter = await navigator.gpu?.requestAdapter();
      return adapter ? await adapter.requestDevice() : null;
    } catch { return null; }
  })();
  return devicePromise;
}

async function renderGpu(device: GPUDevice, run: CymsimRun, geom: Geometry, opts: RenderOptions): Promise<RenderedListener | null> {
  const T = run.acceleration!.nSamples, E = geom.E, D = geom.maxLag;
  const chunk = Math.max(256, Math.min(16384, Math.floor(48e6 / (4 * E)) - D));
  const buf = (data: Float32Array<ArrayBuffer> | Uint32Array<ArrayBuffer>) => {
    const b = device.createBuffer({ size: data.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(b, 0, data);
    return b;
  };
  const coefs = buf(geom.coefs), delays = buf(geom.delays), taps = buf(geom.taps);
  const win = device.createBuffer({ size: (D + chunk) * E * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
  const outp = device.createBuffer({ size: 2 * chunk * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
  const stage = device.createBuffer({ size: 2 * chunk * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
  const uni = device.createBuffer({ size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
  const pipe = device.createComputePipeline({
    layout: "auto",
    compute: { module: device.createShaderModule({ code: GATHER_WGSL }), entryPoint: "gather" },
  });
  const bind = device.createBindGroup({
    layout: pipe.getBindGroupLayout(0),
    entries: [uni, coefs, delays, taps, win, outp].map((b, i) => ({ binding: i, resource: { buffer: b } })),
  });
  const front = new Float32Array(T), back = new Float32Array(T);
  try {
    for (let j0 = 0; j0 < T; j0 += chunk) {
      if (opts.shouldStop?.()) return null;
      const len = Math.min(chunk, T - j0);
      device.queue.writeBuffer(uni, 0, new Uint32Array([E, D, len, j0]));
      device.queue.writeBuffer(win, 0, elementMeans(run, j0 - D, j0 + len));
      const enc = device.createCommandEncoder();
      const pass = enc.beginComputePass();
      pass.setPipeline(pipe); pass.setBindGroup(0, bind); pass.dispatchWorkgroups(len);
      pass.end();
      enc.copyBufferToBuffer(outp, 0, stage, 0, 2 * len * 4);
      device.queue.submit([enc.finish()]);
      await stage.mapAsync(GPUMapMode.READ, 0, 2 * len * 4);
      const rows = new Float32Array(stage.getMappedRange(0, 2 * len * 4)).slice();
      stage.unmap();
      for (let s = 0; s < len; s++) { front[j0 + s] = rows[2 * s]; back[j0 + s] = rows[2 * s + 1]; }
      opts.onProgress?.((j0 + len) / T, "gpu");
    }
  } finally {
    for (const b of [coefs, delays, taps, win, outp, stage, uni]) b.destroy();
  }
  return { front, back, position: [0, 0, 0], backend: "gpu" };
}

async function renderCpu(run: CymsimRun, geom: Geometry, opts: RenderOptions): Promise<RenderedListener | null> {
  const T = run.acceleration!.nSamples, E = geom.E, D = geom.maxLag;
  const chunk = Math.max(256, Math.min(8192, Math.floor(16e6 / (4 * E)) - D));
  const front = new Float32Array(T), back = new Float32Array(T);
  const { coefs, delays, taps } = geom;
  for (let j0 = 0; j0 < T; j0 += chunk) {
    if (opts.shouldStop?.()) return null;
    const len = Math.min(chunk, T - j0);
    const win = elementMeans(run, j0 - D, j0 + len);
    for (let s = 0; s < len; s++) {
      const jg = j0 + s;
      let f = 0, b = 0;
      for (let e = 0; e < E; e++) {
        const i0f = delays[2 * e], i0b = delays[2 * e + 1];
        let af = 0, ab = 0;
        for (let t = 0; t < 16; t++) {
          const lagF = i0f + t - 7;
          if (lagF >= 0 && lagF <= jg) af += taps[32 * e + t] * win[(D + s - lagF) * E + e];
          const lagB = i0b + t - 7;
          if (lagB >= 0 && lagB <= jg) ab += taps[32 * e + 16 + t] * win[(D + s - lagB) * E + e];
        }
        f += coefs[2 * e] * af; b += coefs[2 * e + 1] * ab;
      }
      front[jg] = f; back[jg] = b;
    }
    opts.onProgress?.((j0 + len) / T, "cpu");
    await new Promise((r) => setTimeout(r, 0));   // let the page breathe between chunks
  }
  return { front, back, position: [0, 0, 0], backend: "cpu" };
}

/** Render the baffled (front) and back signals at ``position``; null when stopped or no record. */
export async function renderListener(run: CymsimRun, position: [number, number, number], opts: RenderOptions = {}): Promise<RenderedListener | null> {
  if (!run.acceleration || run.acceleration.nSamples === 0) return null;
  const geom = buildGeometry(run, position);
  const device = await gpuDevice();
  const out = device ? await renderGpu(device, run, geom, opts) : await renderCpu(run, geom, opts);
  if (out) out.position = [...position];
  return out;
}