/**
 * Numerical proof of concept for the modal fast path (drum).
 *
 * The claim to verify, before building anything: the drum's mic trace can be
 * reproduced, to floating-point accuracy, without timestepping the membrane
 * or (after a one-time precompute) the air. It rests on three facts about
 * the discrete scheme as compiled today:
 *
 *   1. The membrane's 5-point Laplacian with pinned rim has exact product-
 *      sine eigenmodes, so each modal amplitude obeys a scalar two-term
 *      recursion — no PDE stepping needed.
 *   2. The coupling is one-way and everything is linear, so the source the
 *      air sees is a sum of fixed spatial patterns (prof .* rad(mode)) times
 *      those scalar series.
 *   3. The air update's transfer matrix is symmetric (the masked Laplacian
 *      in divergence form has w_i*w_j/h^2 off-diagonal entries; everything
 *      else is diagonal), so discrete reciprocity holds exactly: one run
 *      with an impulse at the mic yields every pattern-to-mic impulse
 *      response at once.
 *
 * This script builds the real drum session at a small grid, runs the real
 * GPU solver, and then checks each link in the chain against f64 CPU
 * references:
 *
 *   A. a CPU replica of the membrane leapfrog vs the GPU membrane state;
 *   B. the modal (closed-form) membrane vs the CPU replica — this is the
 *      "modes are exact" claim, and should agree to ~1e-12;
 *   C. a CPU replica of the whole coupled step vs the GPU mic trace;
 *   D. reciprocity on the air alone: direct pattern-source run vs the
 *      projected reciprocal run — should agree to f64 roundoff;
 *   E. end to end: mic trace synthesized from modal amplitudes convolved
 *      with reciprocal kernels, vs the direct CPU trace and the GPU trace.
 *
 * Usage: npx vite-node scripts/modal-poc.ts [--nx 32] [--steps 1200]
 */
import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
import { requestAcousticDevice } from '../src/device.ts';
import { StudioSession, specParams } from '../src/studio/session.ts';
import { presetByKey, cloneSpec } from '../src/studio/presets.ts';

const argOf = (name: string, dflt: number): number => {
  const i = process.argv.indexOf(name);
  return i >= 0 ? Number(process.argv[i + 1]) : dflt;
};
const NX = argOf('--nx', 32);
const STEPS = argOf('--steps', 1200);

const rms = (a: Float64Array | Float32Array): number => {
  let s = 0;
  for (let i = 0; i < a.length; i++) s += a[i] * a[i];
  return Math.sqrt(s / a.length);
};
const relErr = (a: ArrayLike<number>, b: ArrayLike<number>): number => {
  let d = 0;
  let r = 0;
  for (let i = 0; i < a.length; i++) {
    d += (a[i] - b[i]) * (a[i] - b[i]);
    r += b[i] * b[i];
  }
  return Math.sqrt(d / Math.max(r, 1e-300));
};
const report = (label: string, err: number, tol: number): boolean => {
  const ok = err < tol;
  console.log(`  ${ok ? 'ok  ' : 'FAIL'} ${label}: rel err ${err.toExponential(2)} (tol ${tol})`);
  return ok;
};

/* ------------------------------------------------------------- GPU run --- */

try {
  await installWebGpu();
} catch (e) {
  console.error(errMsg(e), '\n', NO_ADAPTER_HINT);
  process.exit(1);
}
const device = await requestAcousticDevice();
const preset = presetByKey('drum')!;
const spec = cloneSpec(preset.spec);
const { params, sceneParams } = specParams(spec);
const session = await StudioSession.create({ device, spec, params, sceneParams, nx: NX });

const air = session.air;
const g = session.grids['m1'];
const dt = session.dt;
const P = session.params;
console.log(
  `drum at nx = ${NX}: air ${air.nx}x${air.ny}x${air.nz}, membrane ${g.n1}x${g.n2}, ` +
    `dt = ${(1e6 * dt).toFixed(2)} µs, ${STEPS} steps`,
);

session.setMic(preset.mic.x, preset.mic.y, preset.mic.z);
session.excite();
session.step(STEPS);
await session.sync();
const gpuTrace = await session.recorder.read();
const gpuU = await session.read('u_m1');

/* ----------------------------------------------- CPU membrane leapfrog --- */

const { n1, n2, h1, h2, L1, L2 } = g;
const mN = g.npts;
const pin = g.pin;
const f0 = P.f0_m1;
const t60 = P.t60_m1;
const sig1 = P.sig1_m1;
const gair = P.gair_m1;
const cm = (2 * f0) / Math.sqrt(1 / (L1 * L1) + 1 / (L2 * L2));
const dmp = (6.91 / t60) * dt;

// The strike, exactly as the generated init writes it.
const um0 = new Float64Array(mN);
for (let k = 0; k < mN; k++) {
  const dx = g.xs64[k] - P.strikex_m1 * L1;
  const dy = g.ys64[k] - P.strikey_m1 * L2;
  const rr = (dx * dx + dy * dy) / (P.striker_m1 * P.striker_m1);
  um0[k] = -(dt * P.vel_m1) * pin[k] * Math.exp(-rr);
}

const lap2 = (u: Float64Array, out: Float64Array): void => {
  const ih1 = 1 / (h1 * h1);
  const ih2 = 1 / (h2 * h2);
  for (let j = 0; j < n2; j++) {
    for (let i = 0; i < n1; i++) {
      const k = i + n1 * j;
      const c = u[k];
      const xm = i > 0 ? u[k - 1] : 0;
      const xp = i < n1 - 1 ? u[k + 1] : 0;
      const ym = j > 0 ? u[k - n1] : 0;
      const yp = j < n2 - 1 ? u[k + n1] : 0;
      out[k] = ih1 * (xm - 2 * c + xp) + ih2 * (ym - 2 * c + yp);
    }
  }
};

// One membrane step: (u, um) -> un, matching the generated step line.
const memStep = (u: Float64Array, um: Float64Array, un: Float64Array, lu: Float64Array, lum: Float64Array): void => {
  lap2(u, lu);
  lap2(um, lum);
  for (let k = 0; k < mN; k++) {
    un[k] =
      pin[k] *
      ((2 * u[k] - (1 - dmp) * um[k] + dt * dt * cm * cm * lu[k] + 2 * sig1 * dt * (lu[k] - lum[k])) /
        (1 + dmp));
  }
};

// Run the membrane alone, keeping every step's acceleration.
console.log('\nA/B. membrane: CPU leapfrog and closed-form modes');
let u = new Float64Array(mN);
let um = Float64Array.from(um0);
const accCpu: Float64Array[] = [];
{
  const lu = new Float64Array(mN);
  const lum = new Float64Array(mN);
  for (let k = 0; k < STEPS; k++) {
    const un = new Float64Array(mN);
    memStep(u, um, un, lu, lum);
    const acc = new Float64Array(mN);
    for (let i = 0; i < mN; i++) acc[i] = (un[i] - 2 * u[i] + um[i]) / (dt * dt);
    accCpu.push(acc);
    um = u;
    u = un;
  }
}
let okAll = report('CPU leapfrog vs GPU u_m1 (f32)', relErr(u, gpuU), 5e-3);

/* -------------------------------------------------- modal membrane (B) --- */

// Interior product-sine modes of the pinned 5-point Laplacian.
const M1 = n1 - 2;
const M2 = n2 - 2;
// Strike projection b[m][n] via two small separable contractions.
const sinTab1 = new Float64Array(M1 * n1); // sin(m*pi*i/(n1-1))
const sinTab2 = new Float64Array(M2 * n2);
for (let m = 1; m <= M1; m++)
  for (let i = 0; i < n1; i++) sinTab1[(m - 1) * n1 + i] = Math.sin((m * Math.PI * i) / (n1 - 1));
for (let n = 1; n <= M2; n++)
  for (let j = 0; j < n2; j++) sinTab2[(n - 1) * n2 + j] = Math.sin((n * Math.PI * j) / (n2 - 1));

const projectModes = (v: Float64Array): Float64Array => {
  // v (n1 x n2) -> coefficients (M1 x M2), normalized.
  const tmp = new Float64Array(M1 * n2);
  for (let m = 0; m < M1; m++) {
    for (let j = 0; j < n2; j++) {
      let s = 0;
      for (let i = 0; i < n1; i++) s += sinTab1[m * n1 + i] * v[i + n1 * j];
      tmp[m + M1 * j] = s;
    }
  }
  const out = new Float64Array(M1 * M2);
  const norm = ((n1 - 1) / 2) * ((n2 - 1) / 2);
  for (let m = 0; m < M1; m++) {
    for (let n = 0; n < M2; n++) {
      let s = 0;
      for (let j = 0; j < n2; j++) s += sinTab2[n * n2 + j] * tmp[m + M1 * j];
      out[m + M1 * n] = s / norm;
    }
  }
  return out;
};

const b = projectModes(um0); // alpha[-1] per mode; alpha[0] = 0.

// Per-mode recursion, exactly the membrane step restricted to one mode.
const lamOf = (m: number, n: number): number => {
  const s1 = Math.sin((m * Math.PI) / (2 * (n1 - 1)));
  const s2 = Math.sin((n * Math.PI) / (2 * (n2 - 1)));
  return (4 / (h1 * h1)) * s1 * s1 + (4 / (h2 * h2)) * s2 * s2;
};

// Keep modes the strike actually excites.
let bMax = 0;
for (const v of b) bMax = Math.max(bMax, Math.abs(v));
const kept: { m: number; n: number; lam: number; b: number }[] = [];
for (let m = 1; m <= M1; m++)
  for (let n = 1; n <= M2; n++)
    if (Math.abs(b[m - 1 + M1 * (n - 1)]) > 1e-10 * bMax)
      kept.push({ m, n, lam: lamOf(m, n), b: b[m - 1 + M1 * (n - 1)] });
console.log(`  ${kept.length} of ${M1 * M2} modes carry the strike (rel cutoff 1e-10)`);

// Modal alpha series and accelerations.
const alphaAcc = (lam: number, bm: number): { acc: Float64Array; last: number } => {
  const acc = new Float64Array(STEPS);
  let prev = bm; // alpha[k-1]
  let cur = 0; // alpha[k]
  const cA = (2 - dt * dt * cm * cm * lam - 2 * sig1 * dt * lam) / (1 + dmp);
  const cB = (-(1 - dmp) + 2 * sig1 * dt * lam) / (1 + dmp);
  for (let k = 0; k < STEPS; k++) {
    const next = cA * cur + cB * prev;
    acc[k] = (next - 2 * cur + prev) / (dt * dt);
    prev = cur;
    cur = next;
  }
  return { acc, last: cur };
};

const modal = kept.map((md) => ({ ...md, ...alphaAcc(md.lam, md.b) }));

// Reconstruct u after STEPS from modal amplitudes and compare to the CPU run.
{
  const uModal = new Float64Array(mN);
  for (const md of modal) {
    for (let j = 0; j < n2; j++) {
      const sj = sinTab2[(md.n - 1) * n2 + j];
      for (let i = 0; i < n1; i++) {
        uModal[i + n1 * j] += md.last * sinTab1[(md.m - 1) * n1 + i] * sj;
      }
    }
  }
  okAll = report('modal membrane vs CPU leapfrog', relErr(uModal, u), 1e-9) && okAll;
}

/* ------------------------------------------------------- CPU air replica --- */

console.log('\nC. coupled CPU replica vs the GPU mic trace');
const { nx, ny, nz, npts, h } = air;
const F = session.scene.fields;
const wall = F.wall;
const sigF = F.sig;
const cF = F.c;
const prof = F.prof_m1;

// The mic index, as setMic computes it.
const micIx = Math.max(0, Math.min(nx - 1, Math.round((preset.mic.x + air.Lx / 2) / h - 0.5)));
const micIy = Math.max(0, Math.min(ny - 1, Math.round((preset.mic.y + air.Ly / 2) / h - 0.5)));
const micIz = Math.max(0, Math.min(nz - 1, Math.round((preset.mic.z + air.Lz / 2) / h - 0.5)));
const micIdx = micIx + nx * (micIy + ny * micIz);

// Bilinear sampling weights of the membrane at each air column (x, y) —
// exactly what rad_m1 bakes into its WGSL.
interface RadW { k00: number; k10: number; k01: number; k11: number; w00: number; w10: number; w01: number; w11: number }
const radW: (RadW | null)[] = [];
for (let iy = 0; iy < ny; iy++) {
  for (let ix = 0; ix < nx; ix++) {
    const xw = -air.Lx / 2 + (ix + 0.5) * h;
    const yw = -air.Ly / 2 + (iy + 0.5) * h;
    const uu = (xw + L1 / 2) / h1;
    const vv = (yw + L2 / 2) / h2;
    const ju = Math.floor(uu);
    const jv = Math.floor(vv);
    if (ju < 0 || ju >= n1 - 1 || jv < 0 || jv >= n2 - 1) {
      radW[ix + nx * iy] = null;
      continue;
    }
    const fu = uu - ju;
    const fv = vv - jv;
    radW[ix + nx * iy] = {
      k00: ju + n1 * jv, k10: ju + 1 + n1 * jv, k01: ju + n1 * (jv + 1), k11: ju + 1 + n1 * (jv + 1),
      w00: (1 - fu) * (1 - fv), w10: fu * (1 - fv), w01: (1 - fu) * fv, w11: fu * fv,
    };
  }
}
const radSample = (v: Float64Array, col: number): number => {
  const w = radW[col];
  if (!w) return 0;
  return w.w00 * v[w.k00] + w.w10 * v[w.k10] + w.w01 * v[w.k01] + w.w11 * v[w.k11];
};

// One air step; drv is the already-assembled source field (or null).
const airStep = (p: Float64Array, pm: Float64Array, pn: Float64Array, drv: ((i: number) => number) | null): void => {
  const ih2 = 1 / (h * h);
  for (let iz = 0; iz < nz; iz++) {
    for (let iy = 0; iy < ny; iy++) {
      for (let ix = 0; ix < nx; ix++) {
        const i = ix + nx * (iy + ny * iz);
        const pc = p[i];
        let s = 0;
        s += ix > 0 ? wall[i - 1] * (p[i - 1] - pc) : -pc;
        s += ix < nx - 1 ? wall[i + 1] * (p[i + 1] - pc) : -pc;
        s += iy > 0 ? wall[i - nx] * (p[i - nx] - pc) : -pc;
        s += iy < ny - 1 ? wall[i + nx] * (p[i + nx] - pc) : -pc;
        s += iz > 0 ? wall[i - nx * ny] * (p[i - nx * ny] - pc) : -pc;
        s += iz < nz - 1 ? wall[i + nx * ny] * (p[i + nx * ny] - pc) : -pc;
        const lp = ih2 * wall[i] * s;
        const sd = sigF[i] * dt;
        const src = drv ? drv(i) : 0;
        pn[i] =
          wall[i] *
          ((2 * pc - (1 - sd) * pm[i] + cF[i] * dt * cF[i] * dt * lp + dt * dt * src) / (1 + sd));
      }
    }
  }
};

// The coupled run: modal membrane is already verified ≡ leapfrog, so drive
// the air from the stored leapfrog accelerations.
const cpuTrace = new Float64Array(STEPS);
{
  let p = new Float64Array(npts);
  let pm = new Float64Array(npts);
  let pn = new Float64Array(npts);
  for (let k = 0; k < STEPS; k++) {
    const acc = accCpu[k];
    airStep(p, pm, pn, (i) => gair * prof[i] * radSample(acc, i % (nx * ny)));
    cpuTrace[k] = pn[micIdx];
    const t = pm;
    pm = p;
    p = pn;
    pn = t;
  }
}
okAll = report('CPU coupled trace vs GPU trace (f32)', relErr(cpuTrace, gpuTrace.subarray(0, STEPS)), 5e-3) && okAll;

/* ---------------------------------------------------- reciprocity (D/E) --- */

console.log('\nD. discrete reciprocity on the air');

// Reciprocal run: unit source at the mic at step 0, projected each step onto
// every kept mode's pattern prof .* rad(phi). The projection is separable:
// contract z first (plane = sum_z prof .* p), then the two sine factors.
// rad(phi_mn)(ix,iy) = sxA[m][ix] * syA[n][iy], because bilinear
// interpolation of a product of sines factorizes per axis.
const sxA = new Float64Array(M1 * nx);
const syA = new Float64Array(M2 * ny);
for (let ix = 0; ix < nx; ix++) {
  const xw = -air.Lx / 2 + (ix + 0.5) * h;
  const uu = (xw + L1 / 2) / h1;
  const ju = Math.floor(uu);
  if (ju < 0 || ju >= n1 - 1) continue;
  const fu = uu - ju;
  for (let m = 1; m <= M1; m++)
    sxA[(m - 1) * nx + ix] = (1 - fu) * sinTab1[(m - 1) * n1 + ju] + fu * sinTab1[(m - 1) * n1 + ju + 1];
}
for (let iy = 0; iy < ny; iy++) {
  const yw = -air.Ly / 2 + (iy + 0.5) * h;
  const vv = (yw + L2 / 2) / h2;
  const jv = Math.floor(vv);
  if (jv < 0 || jv >= n2 - 1) continue;
  const fv = vv - jv;
  for (let n = 1; n <= M2; n++)
    syA[(n - 1) * ny + iy] = (1 - fv) * sinTab2[(n - 1) * n2 + jv] + fv * sinTab2[(n - 1) * n2 + jv + 1];
}

const kernels = modal.map(() => new Float64Array(STEPS));
{
  let p = new Float64Array(npts);
  let pm = new Float64Array(npts);
  let pn = new Float64Array(npts);
  const plane = new Float64Array(nx * ny);
  const tmp = new Float64Array(M1 * ny);
  for (let k = 0; k < STEPS; k++) {
    airStep(p, pm, pn, k === 0 ? (i) => (i === micIdx ? 1 : 0) : null);
    // plane(ix,iy) = sum_z prof * pn
    plane.fill(0);
    for (let iz = 0; iz < nz; iz++) {
      const base = nx * ny * iz;
      for (let c2 = 0; c2 < nx * ny; c2++) plane[c2] += prof[base + c2] * pn[base + c2];
    }
    // tmp(m,iy) = sum_ix sxA[m][ix] * plane(ix,iy)
    tmp.fill(0);
    for (let m = 0; m < M1; m++) {
      for (let iy = 0; iy < ny; iy++) {
        let s = 0;
        for (let ix = 0; ix < nx; ix++) s += sxA[m * nx + ix] * plane[ix + nx * iy];
        tmp[m + M1 * iy] = s;
      }
    }
    for (let q = 0; q < modal.length; q++) {
      const { m, n } = modal[q];
      let s = 0;
      for (let iy = 0; iy < ny; iy++) s += syA[(n - 1) * ny + iy] * tmp[m - 1 + M1 * iy];
      kernels[q][k] = s;
    }
    const t = pm;
    pm = p;
    p = pn;
    pn = t;
  }
}

// Direct check for one mode: fire its pattern as an impulse, listen at the
// mic; must equal the reciprocal kernel to f64 roundoff.
{
  const q = modal.findIndex((md) => md.m === 2 && md.n === 3);
  const md = modal[q >= 0 ? q : 0];
  const pat = new Float64Array(npts);
  for (let i = 0; i < npts; i++) {
    const col = i % (nx * ny);
    pat[i] = prof[i] * sxA[(md.m - 1) * nx + (col % nx)] * syA[(md.n - 1) * ny + ((col / nx) | 0)];
  }
  const direct = new Float64Array(STEPS);
  let p = new Float64Array(npts);
  let pm = new Float64Array(npts);
  let pn = new Float64Array(npts);
  for (let k = 0; k < STEPS; k++) {
    airStep(p, pm, pn, k === 0 ? (i) => pat[i] : null);
    direct[k] = pn[micIdx];
    const t = pm;
    pm = p;
    p = pn;
    pn = t;
  }
  okAll = report(`reciprocity, mode (${md.m},${md.n})`, relErr(kernels[q >= 0 ? q : 0], direct), 1e-10) && okAll;
}

/* ------------------------------------------------- modal synthesis (E) --- */

console.log('\nE. modal synthesis vs direct traces');
const yModal = new Float64Array(STEPS);
for (let q = 0; q < modal.length; q++) {
  const hk = kernels[q];
  const a = modal[q].acc;
  // y[t] += gair * sum_k hk[t-k] * a[k]
  for (let t = 0; t < STEPS; t++) {
    let s = 0;
    for (let k = 0; k <= t; k++) s += hk[t - k] * a[k];
    yModal[t] += gair * s;
  }
}
okAll = report('modal synthesis vs CPU coupled trace', relErr(yModal, cpuTrace), 1e-6) && okAll;
okAll = report('modal synthesis vs GPU trace (f32)', relErr(yModal, gpuTrace.subarray(0, STEPS)), 5e-3) && okAll;

console.log(okAll ? '\nall checks passed' : '\nSOME CHECKS FAILED');
session.destroy();
process.exit(okAll ? 0 : 1);