// Driver: numbl setup stage -> WebGPU time loop -> numbl output stage.
//
//   node src/run.mjs                     full default render, writes cymbal2-gpu.wav
//   node src/run.mjs --pinned            the validate suite's reduced configuration
//   node src/run.mjs --compare           GPU vs CPU (numbl cymbal2_sim) on the pinned config
//   node src/run.mjs --mode mega|multi|persistent   backend (default multi)
//   node src/run.mjs --velocity 4 ...    parameter overrides

import { writeFileSync } from 'node:fs';
import { runSetup, runOutput, runCpuSim, loadMFiles } from './setup.mjs';
import { buildKernel, buildLayout } from './kernel.mjs';
import { buildStages } from './kernels2.mjs';
import { buildPersistent } from './kernels3.mjs';
import { installNodeWebGpu, requestDevice, runMegakernel, runMultiDispatch, runPersistent } from './gpu.mjs';
import { encodeWav } from './wav.mjs';

function parseArgs(argv) {
  const a = { params: {}, chunk: 2048, ws: 256, wav: 'cymbal2-gpu.wav', compare: false, bench: false, mode: 'multi' };
  for (let i = 2; i < argv.length; i++) {
    const t = argv[i];
    if (t === '--pinned') Object.assign(a.params, { sample_rate: 44100, oversample: 1, tail_s: 1.5 });
    else if (t === '--compare') { a.compare = true; Object.assign(a.params, { sample_rate: 44100, oversample: 1, tail_s: 1.5 }); }
    else if (t === '--bench') a.bench = true;
    else if (t === '--velocity') a.params.strike_velocity = Number(argv[++i]);
    else if (t === '--count') a.params.strike_count = Number(argv[++i]);
    else if (t === '--rate') a.params.strike_rate_hz = Number(argv[++i]);
    else if (t === '--oversample') a.params.oversample = Number(argv[++i]);
    else if (t === '--tail') a.params.tail_s = Number(argv[++i]);
    else if (t === '--chunk') a.chunk = Number(argv[++i]);
    else if (t === '--ws') a.ws = Number(argv[++i]);
    else if (t === '--wav') a.wav = argv[++i];
    else if (t === '--mode') a.mode = argv[++i];
    else if (t === '--set') { const [k, v] = argv[++i].split('='); a.params[k] = Number(v); }
    else throw new Error(`unknown argument ${t}`);
  }
  return a;
}

function hostDiagnostics(fhist, state, kernel, setup) {
  let peakF = 0, contactSteps = 0, events = 0, prev = 0;
  for (let i = 0; i < fhist.length; i++) {
    const f = fhist[i];
    if (f > peakF) peakF = f;
    if (f > 0) { contactSteps++; if (prev === 0) events++; }
    prev = f;
  }
  const L = kernel.layout;
  return {
    peakForceN: peakF,
    contactMs: (1000 * contactSteps) / setup.simRate,
    contactEvents: events,
    peakXiMm: state[L.SC + 12] * 1000,
    peakDispThickness: state[L.SC + 13] / setup.thickness,
    contactVel: Array.from(state.subarray(L.ST_CVEL, L.ST_CVEL + setup.K)),
    strokesPlayed: Math.min(state[L.SC + 4] + (state[L.SC + 3] > 0 ? 1 : 0), setup.K),
  };
}

const rms = (x) => {
  let s = 0;
  for (let i = 0; i < x.length; i++) s += x[i] * x[i];
  return Math.sqrt(s / x.length);
};

async function main() {
  const args = parseArgs(process.argv);

  console.log('setup (numbl): grid, per-wavenumber inverses, mallet, schedule ...');
  let t0 = performance.now();
  const setup = runSetup(args.params);
  console.log(`  ${(performance.now() - t0).toFixed(0)} ms; grid ${setup.nrt}x${setup.nth}, nfree ${setup.nfree}, nsteps ${setup.nsteps}, simRate ${setup.simRate}`);

  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}`);

  let kernel;
  let gpu;
  t0 = performance.now();
  const progress = (k, n, ms) => {
    if (k === n || k % (1024 * 64) === 0) {
      process.stdout.write(`\r  ${k}/${n} steps, ${(1000 * ms / k).toFixed(1)} us/step`);
    }
  };
  if (args.mode === 'persistent') {
    const L = buildLayout(setup);
    kernel = L;
    const pk = buildPersistent(L, { chunk: args.chunk >= 4096 ? 512 : Math.min(args.chunk, 2048) });
    console.log(`persistent kernel: ${pk.nwg} workgroups of ${pk.wg}, ${pk.chunk} steps/dispatch, ops ${(L.opsData.byteLength / 1024).toFixed(0)} KB`);
    gpu = await runPersistent(device, L, pk, progress);
  } else if (args.mode === 'mega') {
    kernel = buildKernel(setup, { ws: args.ws, chunk: args.chunk });
    console.log(`megakernel: ws ${args.ws}, ${args.chunk} steps/dispatch, ops ${(kernel.opsData.byteLength / 1024).toFixed(0)} KB`);
    gpu = await runMegakernel(device, kernel, setup.nsteps, progress);
  } else {
    const L = buildLayout(setup);
    kernel = L; // shares .layout for diagnostics
    const stages = buildStages(L, {});
    console.log(`multi-dispatch: ${stages.order.length} dispatches/step, wg ${stages.wg} x ${stages.nwg}, ops ${(L.opsData.byteLength / 1024).toFixed(0)} KB`);
    gpu = await runMultiDispatch(device, L, stages, progress);
  }
  console.log('');
  const usPerStep = (1000 * gpu.gpuMs) / setup.nsteps;
  console.log(`gpu integration: ${(gpu.gpuMs / 1000).toFixed(2)} s total, ${usPerStep.toFixed(1)} us/step`);

  let bad = 0;
  let firstBad = -1;
  for (let i = 0; i < gpu.acc.length; i++) if (!Number.isFinite(gpu.acc[i])) { bad++; if (firstBad < 0) firstBad = i; }
  if (process.env.DUMP_ACC) {
    const { writeFileSync: wf } = await import('node:fs');
    wf(process.env.DUMP_ACC, Buffer.from(gpu.acc.buffer, gpu.acc.byteOffset, gpu.acc.byteLength));
    console.log(`acc dumped; nonfinite ${bad}${bad ? ' first at ' + firstBad : ''}`);
  }
  if (bad > 0 && !process.env.DUMP_ACC) throw new Error(`simulation diverged: ${bad} non-finite acceleration samples (first at ${firstBad})`);
  if (bad > 0 && process.env.DUMP_ACC) { console.log('skipping output stage (diverged)'); return; }

  const d = hostDiagnostics(gpu.fhist, gpu.state, kernel, setup);
  console.log(`contact: peakF ${d.peakForceN.toFixed(2)} N, cvel ${d.contactVel[0].toFixed(3)} m/s, ` +
    `${d.contactMs.toFixed(2)} ms over ${d.contactEvents} events, peakXi ${d.peakXiMm.toFixed(3)} mm, ` +
    `peakDisp ${d.peakDispThickness.toFixed(1)} thicknesses`);

  if (args.compare) {
    console.log('cpu reference (numbl, unmodified cymbal2_sim) ...');
    const cpu = runCpuSim(setup.mfiles, args.params);
    console.log(`  ${(cpu.ms / 1000).toFixed(1)} s total, ${(1000 * cpu.ms / cpu.nsteps).toFixed(1)} us/step`);
    console.log(`  cpu: peakF ${cpu.peakForceN.toFixed(2)} N, cvel ${cpu.contactVel.toFixed(3)}, contactMs ${cpu.contactMs.toFixed(2)}, events ${cpu.contactEvents}, peakXi ${cpu.peakXiMm.toFixed(3)} mm`);

    // Contact window alignment and force agreement
    const firstC = (f) => f.findIndex((x) => x > 0);
    const kg = firstC(gpu.fhist), kc = firstC(cpu.fhist);
    console.log(`  first contact step: gpu ${kg}, cpu ${kc}${kg === kc ? '' : '  << MISALIGNED'}`);
    let femax = 0;
    for (let i = 0; i < cpu.fhist.length; i++) {
      const e = Math.abs(gpu.fhist[i] - cpu.fhist[i]);
      if (e > femax) femax = e;
    }
    console.log(`  max |force diff| ${femax.toExponential(2)} N (peak ${cpu.peakForceN.toFixed(1)} N)`);

    // Short-horizon acceleration agreement (before chaotic divergence)
    for (const w of [256, 1024, 4096]) {
      let num = 0, den = 0;
      for (let i = kc; i < Math.min(kc + w, cpu.acc.length); i++) {
        num += (gpu.acc[i] - cpu.acc[i]) ** 2;
        den += cpu.acc[i] ** 2;
      }
      console.log(`  acc rel-L2 over ${w} steps after contact: ${Math.sqrt(num / Math.max(den, 1e-30)).toExponential(2)}`);
    }
    console.log(`  acc rms: gpu ${rms(gpu.acc).toExponential(4)}, cpu ${rms(cpu.acc).toExponential(4)}`);
  }

  console.log('output stage (numbl cymbal2_output) ...');
  const out = runOutput(loadMFiles(), args.params, gpu.acc);
  console.log(`  y: ${out.y.length} samples, rms ${rms(out.y).toFixed(4)}, peakAccel ${out.peakAccel.toExponential(3)} m/s^2`);
  writeFileSync(args.wav, encodeWav(out.y, args.params.sample_rate ?? 48000));
  console.log(`wrote ${args.wav}`);
}

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