// Runs the WGSL kernels under Dawn (the `webgpu` package) and compares them with the float64 CPU
// loop. This is a plain script rather than a node:test file: Dawn's threads abort the process
// under the node: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 { buildSetup } from "../src/model/setup.ts";
import { CpuSimulator } from "../src/cpu/simulator.ts";
import { GpuSimulator } from "../src/gpu/simulator.ts";
import { fixture } from "./helpers.ts";

Object.assign(globalThis, globals);

// one Dawn instance and device for the whole file (a second instance crashes the process)
const gpuInstance = create([]);
const adapterP = gpuInstance.requestAdapter().then((a) => { if (!a) throw new Error("no WebGPU adapter"); return a.requestDevice(); });
async function getDevice(): Promise<GPUDevice> { return adapterP; }

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}`); }
}

async function runBoth(cfg: any, nSamples: number, chunk: number) {
  const setup = await buildSetup(cfg);
  const device = await getDevice();
  const cpu = new CpuSimulator(setup), gpu = new GpuSimulator(device, setup, 4096);
  const pc: number[] = [], pg: number[] = [];
  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.equal(rg.firstSample, rc.firstSample); assert.equal(rg.pressureFront.length, rc.pressureFront.length);
    for (let i = 0; i < rc.pressureFront.length; i++) { pc.push(rc.pressureFront[i] - rc.pressureBack[i]); pg.push(rg.pressureFront[i] - rg.pressureBack[i]); }
  }
  const dc = await cpu.diagnostics(), dg = await gpu.diagnostics();
  gpu.dispose();
  return { setup, pc, pg, dc, dg };
}

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

test("GPU kernels match the CPU loop on a linear run", async () => {
  const cfg = structuredClone(fixture.config);
  cfg.numerics.nonlinear = false;
  const { pc, pg, dc, dg } = await runBoth(cfg, 420, 150);
  assert.ok(rel(dg.hits[0].peakForce, dc.hits[0].peakForce) < 1e-4, `peak ${dg.hits[0].peakForce} vs ${dc.hits[0].peakForce}`);
  assert.equal(Math.round(dg.hits[0].contactTime / 1e-6), Math.round(dc.hits[0].contactTime / 1e-6));
  assert.ok(rel(dg.hits[0].velocity, dc.hits[0].velocity) < 1e-4);
  assert.ok(rel(dg.kinetic, dc.kinetic) < 1e-3 && rel(dg.membrane, dc.membrane) < 1e-3 && rel(dg.bending, dc.bending) < 1e-3,
            `energies ${[dg.kinetic, dg.membrane, dg.bending]} vs ${[dc.kinetic, dc.membrane, dc.bending]}`);
  assert.ok(rel(dg.maxWoverH, dc.maxWoverH) < 1e-3);
  let num = 0, den = 0;
  for (let i = 0; i < pc.length; i++) { num += (pc[i] - pg[i]) ** 2; den += pc[i] ** 2; }
  assert.ok(pc.length > 200 && Math.sqrt(num / den) < 1e-3, `pressure rel error ${Math.sqrt(num / den)}`);
});

test("GPU kernels match the CPU loop on the nonlinear run (float32, short)", async () => {
  const { pc, pg, dc, dg } = await runBoth(fixture.config, 300, 100);
  assert.ok(rel(dg.hits[0].peakForce, dc.hits[0].peakForce) < 1e-3, `peak ${dg.hits[0].peakForce} vs ${dc.hits[0].peakForce}`);
  assert.ok(rel(dg.kinetic + dg.membrane + dg.bending, dc.kinetic + dc.membrane + dc.bending) < 1e-2);
  let num = 0, den = 0;
  for (let i = 0; i < pc.length; i++) { num += (pc[i] - pg[i]) ** 2; den += pc[i] ** 2; }
  assert.ok(Math.sqrt(num / den) < 1e-2, `pressure rel error ${Math.sqrt(num / den)}`);
});

console.log(failures ? `${failures} check(s) failed` : "all GPU checks passed");
// tear the device down and leave explicitly: Dawn aborts when destroyed by the natural exit
(await getDevice()).destroy();
process.exit(failures ? 1 : 0);