// 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 opts = { movieEvery: 7, dispEvery: 50, recordAcceleration: true };
  const cpu = new CpuSimulator(setup, opts), gpu = new GpuSimulator(device, setup, 4096, opts);
  const pc: number[] = [], pg: number[] = [];
  let framesChecked = 0, accChecked = 0, dispChecked = 0;
  const relErr = (a: ArrayLike<number>, b: ArrayLike<number>) => { let num = 0, den = 0; for (let i = 0; i < a.length; i++) { num += (a[i] - b[i]) ** 2; den += a[i] ** 2; } return den > 0 ? Math.sqrt(num / den) : 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.equal(rg.firstSample, rc.firstSample); assert.equal(rg.pressureFront.length, rc.pressureFront.length);
    assert.deepEqual(rg.frameSamples, rc.frameSamples);
    for (let f = 0; f < rc.frames.length; f++) {
      let num = 0, den = 0;
      for (let i = 0; i < rc.frames[f].length; i++) { num += (rc.frames[f][i] - rg.frames[f][i]) ** 2; den += rc.frames[f][i] ** 2; }
      if (den > 0) { assert.ok(Math.sqrt(num / den) < 1e-2, `frame ${f} rel error ${Math.sqrt(num / den)}`); framesChecked++; }
    }
    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]); }
    // the acceleration record and the file displacement frames, sample by sample
    assert.equal(rg.acceleration!.length, rc.acceleration!.length);
    if (rc.acceleration!.length) { const err = relErr(rc.acceleration!, rg.acceleration!); assert.ok(err < 1e-2, `acceleration record rel error ${err}`); accChecked += rc.pressureFront.length; }
    assert.deepEqual(rg.dispSamples, rc.dispSamples);
    for (let f = 0; f < rc.dispFrames.length; f++) { const err = relErr(rc.dispFrames[f], rg.dispFrames[f]); assert.ok(err < 1e-2, `disp frame ${f} rel error ${err}`); dispChecked++; }
  }
  const dc = await cpu.diagnostics(), dg = await gpu.diagnostics();
  assert.ok(framesChecked > 10, `only ${framesChecked} movie frames compared`);
  assert.ok(accChecked > 100 && dispChecked >= 5, `only ${accChecked} acceleration samples and ${dispChecked} displacement frames compared`);
  // the exported state restarts the GPU loop where it was: the next chunk matches a fresh run of the same length
  const st = await gpu.exportState();
  const gpu2 = new GpuSimulator(device, setup, 4096, opts);
  await gpu2.importState(st);
  const [a, b] = [await gpu.runChunk(50), await gpu2.runChunk(50)];
  assert.equal(a.firstSample, b.firstSample);
  assert.ok(relErr(a.pressureFront, b.pressureFront) < 1e-5, "resumed GPU run diverges from the continued one");
  gpu2.dispose(); 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);