// Runs the WebGPU time loop (../instrument/gpuSim.ts) from Node against
// desktop WebGPU (the optional `webgpu` package, prebuilt Google Dawn) on the
// validation suite's reduced configuration, and checks the pinned contact
// diagnostics and output rms. This exercises exactly the code the browser
// worker runs; a machine without the package or without a GPU reports and
// exits 0 so CI can run it opportunistically.
//   npm i --no-save webgpu && npx tsx scripts/gpu-selftest.ts

import { readFileSync, readdirSync } from 'node:fs';
import { fileURLToPath } from 'node:url';
import path from 'node:path';
import { executeCode, RTV, type RuntimeValue, type RuntimeTensor } from 'numbl';

const MATLAB_DIR = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../../matlab');

async function installDawn(): Promise<boolean> {
  try {
    const mod = await import('webgpu');
    Object.assign(globalThis, mod.globals);
    Object.defineProperty(globalThis, 'navigator', {
      value: { gpu: mod.create([]) },
      configurable: true,
    });
    return true;
  } catch {
    return false;
  }
}

function mfiles() {
  return readdirSync(MATLAB_DIR)
    .filter((f) => f.endsWith('.m'))
    .sort()
    .map((name) => ({ name, source: readFileSync(path.join(MATLAB_DIR, name), 'utf8') }));
}

const scalar = (v: RuntimeValue): number => {
  if (typeof v === 'number') return v;
  const t = v as RuntimeTensor;
  if (t && t.kind === 'tensor' && t.data.length === 1) return t.data[0];
  throw new Error('expected scalar');
};
const f32 = (v: RuntimeValue): Float32Array => {
  const t = v as RuntimeTensor;
  return Float32Array.from(t.data);
};

async function main() {
  if (!(await installDawn())) {
    console.log('gpu-selftest: `webgpu` package not installed; skipping (npm i --no-save webgpu)');
    return;
  }
  const { requestGpuDevice, runGpuSim } = await import('../../instrument/gpuSim.ts');

  const files = mfiles();
  const P = `
P = cymbal2_defaults();
P.sample_rate = 44100;
P.oversample = 1;
P.tail_s = 1.5;
`;
  const r = executeCode(`${P}
S = cymbal2_system(P);
MB = cymbal2_mallet(P);
G = S.G;
o_D1 = G.D1; o_D2 = G.D2; o_F = G.F; o_Gi = G.Gi; o_Rd = G.Rd; o_R2d = G.R2d;
o_AiB = S.AiB; o_AiC = S.AiC; o_Ai = S.Ai; o_Fi = S.Fi;
o_ASiB = MB.ASiB; o_ASiC = MB.ASiC;
o_vM = MB.vM; o_vR = MB.vR; o_vK = MB.vK; o_tipcol = MB.tipcol;
K = max(1, round(P.strike_count));
tstr = P.pre_delay_ms * 1e-3 + (0:K - 1) / max(1e-3, P.strike_rate_hz);
vstr = zeros(1, K);
state = P.seed * 2654435761;
for i = 1:K
  state = mod(state * 1664525 + 1013904223, 4294967296);
  x = state / 4294967296;
  vstr(i) = max(0.05, P.strike_velocity * (1 + P.velocity_spread * (2 * x - 1)));
end
gap = P.strike_gap_m;
if gap <= 0
  gap = 1.05 * max(vstr)^2 / (2 * P.strike_accel) + 0.005;
end
travel = gap ./ vstr + vstr / (2 * P.strike_accel);
tlaunch = tstr - travel;
off = -min(0, min(tlaunch));
tlaunch = tlaunch + off;
duration = off + (K - 1) / max(1e-3, P.strike_rate_hz) + P.tail_s;
nsteps = floor(duration * P.sample_rate * P.oversample);
o_tlaunch = tlaunch; o_vstr = vstr;
c_nrt = G.nrt; c_Nr = G.Nr; c_nth = G.nth; c_hr = G.hr; c_hth = G.hth;
c_ht = S.ht; c_Eh = S.Eh; c_Rs = S.Rs; c_rhoh = S.rhoh;
c_fsc = S.forceScale; c_nlr0 = S.eomRow0;
c_ls = S.ls; c_js = S.js; c_lp = S.lp; c_jp = S.jp;
c_tipi = MB.tipi; c_nfree = MB.nfree; c_ps = MB.ps;
c_KM = P.head_stiffness; c_alM = P.head_alpha;
c_aH = P.strike_accel; c_aHr = P.retract_accel;
c_holdSteps = max(1, round(P.strike_hold_ms * 1e-3 / S.ht));
c_missDepth = gap + 0.06;
c_gap = gap; c_K = K; c_nsteps = nsteps;
`, { displayResults: false }, files, 'stage.m');
  const v = r.variableValues;
  const s = (n: string) => scalar(v[n]);
  const t = (n: string) => f32(v[n]);
  const setup = {
    ops: {
      D1: t('o_D1'), D2: t('o_D2'), F: t('o_F'), Gi: t('o_Gi'),
      Rd: t('o_Rd'), R2d: t('o_R2d'),
      AiB: t('o_AiB'), AiC: t('o_AiC'), Ai: t('o_Ai'), Fi: t('o_Fi'),
      ASiB: t('o_ASiB'), ASiC: t('o_ASiC'),
      vM: t('o_vM'), vR: t('o_vR'), vK: t('o_vK'), tipcol: t('o_tipcol'),
    },
    tlaunch: Array.from((v.o_tlaunch as RuntimeTensor).data),
    vstr: Array.from((v.o_vstr as RuntimeTensor).data),
    nrt: s('c_nrt'), Nr: s('c_Nr'), nth: s('c_nth'),
    hr: s('c_hr'), hth: s('c_hth'),
    ht: s('c_ht'), Eh: s('c_Eh'), Rs: s('c_Rs'), rhoh: s('c_rhoh'),
    fsc: s('c_fsc'), nlr0: s('c_nlr0'),
    ls: s('c_ls'), js: s('c_js'), lp: s('c_lp'), jp: s('c_jp'),
    tipi: s('c_tipi'), nfree: s('c_nfree'), ps: s('c_ps'),
    KM: s('c_KM'), alM: s('c_alM'), aH: s('c_aH'), aHr: s('c_aHr'),
    holdSteps: s('c_holdSteps'), missDepth: s('c_missDepth'),
    gap: s('c_gap'), K: s('c_K'), nsteps: s('c_nsteps'),
  };

  let device;
  try {
    device = await requestGpuDevice();
  } catch (e) {
    console.log(`gpu-selftest: no WebGPU device (${e instanceof Error ? e.message : e}); skipping`);
    return;
  }
  const t0 = performance.now();
  const sim = await runGpuSim(device, setup, () => {});
  device.destroy();
  const ms = performance.now() - t0;

  const ro = executeCode(`${P}
[y, peakAccel] = cymbal2_output(P, acc);
`, { displayResults: false, initialVariableValues: { acc: RTV.tensor(Float64Array.from(sim.acc), [1, sim.acc.length]) } }, files, 'stage.m');
  const y = f32(ro.variableValues.y);
  let ss = 0;
  for (let i = 0; i < y.length; i++) ss += y[i] * y[i];
  const rms = Math.sqrt(ss / y.length);
  let peakF = 0;
  for (let i = 0; i < sim.fhist.length; i++) peakF = Math.max(peakF, sim.fhist[i]);
  const firstC = sim.fhist.findIndex((x) => x > 0);

  console.log(
    `gpu-selftest: ${setup.nsteps} steps in ${(ms / 1000).toFixed(1)} s ` +
    `(${((1000 * ms) / setup.nsteps).toFixed(0)} us/step); ` +
    `peakF ${peakF.toFixed(2)} N (ref 311.14), cvel ${sim.contactVel[0].toFixed(3)} (ref -8.223), ` +
    `first contact step ${firstC + 1} (ref 1376), rms ${rms.toFixed(4)} (ref 0.0973)`,
  );
  const ok =
    Math.abs(peakF - 311.14) < 3 &&
    Math.abs(sim.contactVel[0] + 8.223) < 0.08 &&
    firstC + 1 === 1376 &&
    Math.abs(rms - 0.0973) < 0.012;
  if (!ok) {
    console.error('gpu-selftest: FAILED');
    process.exit(1);
  }
  console.log('gpu-selftest: OK');
}

main().then(() => process.exit(0)).catch((e) => {
  console.error(e);
  process.exit(1);
});