// Runs the WGSL kernels under Dawn (the `webgpu` package) and compares them with the float64 CPU
// loop: same modal initial condition or drive, same projection, no hits. This is a plain script
// rather than a node:test file because Dawn's threads abort the process under the test runner.
// Run with `node tests/gpu.check.ts`; a failed check exits non-zero.
import assert from "node:assert/strict";
import { create, globals } from "webgpu";
import { cloneConfig } from "../src/config/yaml.ts";
import { CpuSimulator } from "../src/cpu/simulator.ts";
import { buildDrive, buildProjection, modalAmplitude, modalInitialState, modeShape, type DriveSpec } from "../src/drive.ts";
import { GpuSimulator } from "../src/gpu/simulator.ts";
import { buildSetup } from "../src/model/setup.ts";
import type { BackendOptions } from "../src/model/simulator.ts";
import { buildCondensationOperator } from "../src/precompute/assemble.ts";
import { computeModes } from "../src/precompute/pipeline.ts";
import { miniConfig, noProgress, OPTIONS } from "./helpers.ts";

Object.assign(globalThis, globals);

const gpuInstance = create([]);
const deviceP = gpuInstance.requestAdapter().then((a) => { if (!a) throw new Error("no WebGPU adapter"); return a.requestDevice(); });

let failures = 0;
async function test(name: string, fn: () => Promise<void>) {
  const t0 = performance.now();
  try { await fn(); console.log(`ok - ${name} (${(performance.now() - t0).toFixed(0)} ms)`); }
  catch (err) { failures++; console.log(`not ok - ${name}\n  ${(err as Error).message}`); }
}

const base = miniConfig();
const modes = computeModes(base, OPTIONS, noProgress);
const MODE = 3;

function mw(setup: Awaited<ReturnType<typeof buildSetup>>): Float64Array {
  const N = setup.mesh.n, Mw = new Float64Array(N);
  for (let n = 0; n < N; n++) Mw[n] = setup.shell.M[5 * n + 2];
  return Mw;
}

async function runBoth(nonlinear: boolean, mechanism: "decay" | "forced", nSamples: number, chunk: number) {
  const cfg = cloneConfig(base);
  cfg.numerics.nonlinear = nonlinear;
  const setup = await buildSetup(cfg);
  const shape = modeShape(modes, MODE, 1);
  const projection = buildProjection(modes, mw(setup), 1, 16);
  const opts: BackendOptions = { projection, traceMode: MODE };
  let initial: { q: Float64Array; v: Float64Array } | null = null;
  if (mechanism === "decay") {
    const cond = buildCondensationOperator(setup.mesh, setup.shell, noProgress);
    initial = modalInitialState(shape, cond.complement(shape), modalAmplitude(modes, shape, 2));
  } else {
    const spec: DriveSpec = { mode: MODE, mechanism: "forced", amplitude: 0.3, detuneCents: 0, rampTime: 0.003, driveTime: 1 };
    opts.drive = buildDrive(spec, modes, shape, mw(setup), cfg.material.sigma0, 2 * Math.PI * modes.freq[MODE], setup.dt);
  }
  const device = await deviceP;
  const cpu = new CpuSimulator(setup, opts), gpu = new GpuSimulator(device, setup, 4096, opts);
  if (initial) { cpu.setState(initial.q, initial.v); gpu.setState(initial.q, initial.v); }
  let projChecked = 0, traceErr = 0, traceRef = 0;
  for (let done = 0; done < nSamples; done += chunk) {
    const n = Math.min(chunk, nSamples - done);
    const rc = await cpu.runChunk(n), rg = await gpu.runChunk(n);
    assert.deepEqual(rg.projSamples, rc.projSamples);
    assert.equal(rg.pressureFront.length, rc.pressureFront.length);
    for (let s = 0; s < rc.proj.length; s++) {
      let num = 0, den = 0;
      for (let i = 0; i < rc.proj[s].length; i++) { num += (rc.proj[s][i] - rg.proj[s][i]) ** 2; den += rc.proj[s][i] ** 2; }
      if (den > 0) { assert.ok(Math.sqrt(num / den) < 1e-2, `snapshot ${projChecked} rel error ${Math.sqrt(num / den)}`); projChecked++; }
    }
    assert.equal(rg.traceBase, rc.traceBase);
    assert.equal(rg.trace.length, rc.trace.length);
    for (let i = 0; i < rc.trace.length; i++) { traceErr += (rc.trace[i] - rg.trace[i]) ** 2; traceRef += rc.trace[i] ** 2; }
  }
  if (traceRef > 0) assert.ok(Math.sqrt(traceErr / traceRef) < 1e-2, `trace rel error ${Math.sqrt(traceErr / traceRef)}`);
  const dc = await cpu.diagnostics(), dg = await gpu.diagnostics();
  gpu.dispose();
  assert.ok(projChecked > 5, `only ${projChecked} modal snapshots compared`);
  return { dc, dg };
}

const rel = (a: number, b: number) => Math.abs(a - b) / Math.max(Math.abs(b), 1e-30);

await test("GPU matches the CPU loop on a nonlinear free decay with projection", async () => {
  const { dc, dg } = await runBoth(true, "decay", 400, 150);
  assert.ok(rel(dg.kinetic, dc.kinetic) < 2e-3, `kinetic ${dg.kinetic} vs ${dc.kinetic}`);
  assert.ok(rel(dg.membrane, dc.membrane) < 2e-3, `membrane ${dg.membrane} vs ${dc.membrane}`);
  assert.ok(rel(dg.bending, dc.bending) < 2e-3, `bending ${dg.bending} vs ${dc.bending}`);
  assert.ok(rel(dg.maxWoverH, dc.maxWoverH) < 2e-3);
});

await test("GPU matches the CPU loop on a linear forced run", async () => {
  const { dc, dg } = await runBoth(false, "forced", 400, 130);
  assert.ok(rel(dg.kinetic, dc.kinetic) < 2e-3, `kinetic ${dg.kinetic} vs ${dc.kinetic}`);
  assert.ok(rel(dg.bending, dc.bending) < 2e-3, `bending ${dg.bending} vs ${dc.bending}`);
});

if (failures) { console.error(`${failures} check(s) failed`); process.exit(1); }
console.log("gpu checks passed");
process.exit(0);