// Dataset generator: renders many drumstick strikes headlessly and writes one
// WAV clip plus one line of metadata per clip. Serves dist/ (run `npm run
// build` first) to a headless Chromium with WebGPU through Playwright, and
// drives several pages in parallel, since one simulation is latency-bound and
// leaves a large GPU mostly idle.
//
//   node --experimental-strip-types scripts/generate.ts --out DIR --count N [--start I] [--workers W] [--seed S] [--fps F]
//
// Clip i (global index) is drawn from a generator seeded by (seed, i), so the
// dataset is reproducible and can be extended or regenerated in pieces.
// Existing clips in DIR are skipped. Output: DIR/clips/NNNNNN.wav (16-bit,
// peak-normalized; the peak pressure in Pa is in the metadata) and DIR/index.jsonl.

import { createServer } from "node:http";
import { readFileSync, existsSync, mkdirSync, appendFileSync, writeFileSync, readdirSync } from "node:fs";
import { fileURLToPath } from "node:url";
import { dirname, join, extname } from "node:path";
import { chromium, type Page, type Browser } from "playwright";
import { defaultParams, makeGrid, makeMaterial, type PhysParams } from "../src/params.ts";
import { encodeFeatures, initialShape, randomFeatures, shapeStats, type ShapeFeature } from "../src/shape.ts";

const root = join(dirname(fileURLToPath(import.meta.url)), "..");

// ---------- Arguments ----------

const argv = process.argv.slice(2);
function arg(name: string, def: string): string {
  const k = argv.indexOf(`--${name}`);
  return k >= 0 && k + 1 < argv.length ? argv[k + 1] : def;
}
const OUT = arg("out", "dataset");
const COUNT = parseInt(arg("count", "10"));
const START = parseInt(arg("start", "0"));
const WORKERS = parseInt(arg("workers", "4"));
const SEED = parseInt(arg("seed", "1"));
const FPS = parseFloat(arg("fps", "120"));
const DURATION = parseFloat(arg("duration", "3"));
const JOB_TIMEOUT_MS = parseFloat(arg("timeout", "1800")) * 1000;

// ---------- Random numbers ----------

/** mulberry32, seeded by a hash of (seed, index). */
function rng(seed: number, index: number): () => number {
  let a = (Math.imul(seed, 0x9e3779b1) ^ Math.imul(index + 1, 0x85ebca6b)) >>> 0;
  for (let k = 0; k < 4; k++) { a = Math.imul(a ^ (a >>> 15), 0x2c1b3c6d); a = Math.imul(a ^ (a >>> 12), 0x297a2d39); }
  return () => {
    a = (a + 0x6d2b79f5) >>> 0;
    let t = a;
    t = Math.imul(t ^ (t >>> 15), t | 1);
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

// ---------- Parameter sampling ----------

export interface Sample {
  id: string;
  index: number;
  params: PhysParams;
  /** Shape statistics on the simulation grid. */
  shapeScale: number;
  peakToPeak: number;
  maxSlope: number;
  nx: number;
  ny: number;
}

const NU = 0.29;
const SLOPE_LIMIT = 0.1;
const SPACING = parseFloat(arg("spacing", "0.008"));

/**
 * The targets (E, rho, h, Lx, Ly) are drawn log-uniformly and independently
 * over ranges of about +-30 percent around the defaults. The nuisance variables
 * (initial shape, strike point and speed, losses, listener direction) are
 * drawn independently of them. The grid spacing is fixed (8 mm by default),
 * so the grid dispersion at a given frequency varies with the bending wave
 * constant of the sheet; see the note in the dataset README.
 */
export function drawSample(seed: number, index: number): Sample {
  const rnd = rng(seed, index);
  const U = (a: number, b: number) => a + (b - a) * rnd();
  const logU = (a: number, b: number) => Math.exp(U(Math.log(a), Math.log(b)));
  const IN = 0.0254;
  const E = logU(120e9, 220e9);
  const rho = logU(6000, 9500);
  const h = logU(0.6e-3, 1.1e-3);
  const Lx = logU(9 * IN, 16 * IN);
  const Ly = logU(13 * IN, 24 * IN);
  const p: PhysParams = {
    ...defaultParams,
    shape: "rectangle",
    Lx, Ly, h, E, nu: NU, rho,
    features: [],
    sigma0: defaultParams.sigma0 * logU(0.7, 1.4),
    sigma1: defaultParams.sigma1 * logU(0.7, 1.4),
    sigmaIn: 20,
    excitation: "strike",
    strikeFracX: U(0.6, 0.8),
    strikeFracY: U(0.25, 0.45),
    strikeVelocity: logU(2, 4),
    hitSequence: "single",
    listenerDist: 1,
    listenerAzimuth: U(-15, 15),
    listenerElevation: U(-15, 15),
    spacing: SPACING,
    fps: FPS,
    duration: DURATION,
  };
  // Initial shape: flat with a small probability, otherwise a random feature
  // list with heights of the order of 0.5 to 5 mm (the default is a 3 mm bow);
  // redrawn at half the scale while the largest slope leaves the shallow-shell range.
  const grid = makeGrid(p);
  let shapeScale = rnd() < 0.08 ? 0 : logU(0.5e-3, 5e-3);
  let features: ShapeFeature[] = [];
  let stats = { peakToPeak: 0, maxSlope: 0 };
  for (let attempt = 0; shapeScale > 0 && attempt < 12; attempt++) {
    features = randomFeatures(p, shapeScale, rnd);
    stats = shapeStats(initialShape({ ...p, features }, grid), grid);
    if (stats.maxSlope <= SLOPE_LIMIT) break;
    shapeScale *= 0.5;
  }
  p.features = features;
  return { id: String(index).padStart(6, "0"), index, params: p, shapeScale, peakToPeak: stats.peakToPeak, maxSlope: stats.maxSlope, nx: grid.nx, ny: grid.ny };
}

/** Query string for the app, in the form form.ts applyUrl reads. */
export function toQuery(p: PhysParams): string {
  const parts: string[] = ["autostart=1"];
  for (const k of Object.keys(defaultParams) as (keyof PhysParams)[]) {
    if (k === "features") {
      parts.push(`features=${encodeURIComponent(p.features.length ? encodeFeatures(p.features) : "none")}`);
    } else if (p[k] !== defaultParams[k]) {
      const v = p[k];
      parts.push(`${k}=${encodeURIComponent(typeof v === "number" ? String(v) : typeof v === "boolean" ? (v ? "1" : "0") : String(v))}`);
    }
  }
  return parts.join("&");
}

/** Checks that the page parsed the targets as intended (the link is the app's own encoding of its parameters). */
function checkLink(link: string, p: PhysParams): void {
  const q = new URLSearchParams(link.includes("?") ? link.slice(link.indexOf("?") + 1) : link);
  for (const k of ["E", "rho", "h", "Lx", "Ly", "strikeVelocity", "sigma0", "sigma1"] as const) {
    const got = parseFloat(q.get(k) ?? String(defaultParams[k]));
    if (!(Math.abs(got - p[k]) <= 1e-5 * Math.abs(p[k]))) throw new Error(`parameter ${k} mismatch: page has ${got}, wanted ${p[k]} (link ${link})`);
  }
}

// ---------- WAV ----------

export function wavBytes(samples: Float32Array, fs: number, scale: number): Uint8Array {
  const buf = new ArrayBuffer(44 + samples.length * 2);
  const v = new DataView(buf);
  const str = (o: number, s: string) => { for (let k = 0; k < s.length; k++) v.setUint8(o + k, s.charCodeAt(k)); };
  str(0, "RIFF"); v.setUint32(4, 36 + samples.length * 2, true); str(8, "WAVE");
  str(12, "fmt "); v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
  v.setUint32(24, fs, true); v.setUint32(28, fs * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
  str(36, "data"); v.setUint32(40, samples.length * 2, true);
  for (let k = 0; k < samples.length; k++) {
    const s = Math.max(-1, Math.min(1, samples[k] * scale));
    v.setInt16(44 + 2 * k, Math.round(s * 32767), true);
  }
  return new Uint8Array(buf);
}

// ---------- Static server for dist/ ----------

const MIME: Record<string, string> = { ".html": "text/html", ".js": "text/javascript", ".css": "text/css", ".wgsl": "text/plain" };
export function serveDist(): Promise<string> {
  const dist = join(root, "dist");
  if (!existsSync(join(dist, "index.html"))) throw new Error("dist/index.html not found; run `npm run build` first");
  const server = createServer((req, res) => {
    const path = decodeURIComponent((req.url ?? "/").split("?")[0]);
    const file = join(dist, path === "/" ? "index.html" : path);
    try {
      const data = readFileSync(file);
      res.writeHead(200, { "content-type": MIME[extname(file)] ?? "application/octet-stream" });
      res.end(data);
    } catch {
      res.writeHead(404); res.end();
    }
  });
  return new Promise((resolve) => server.listen(0, "127.0.0.1", () => {
    const addr = server.address();
    resolve(`http://127.0.0.1:${typeof addr === "object" && addr ? addr.port : 0}`);
  }));
}

// ---------- Rendering one clip ----------

export const ARGS = ["--enable-unsafe-webgpu", "--enable-features=Vulkan", "--use-angle=vulkan", "--ignore-gpu-blocklist", "--no-sandbox"];

interface Status { phase: string; t: number; totalSteps: number; wmax: number; peakPressure: number; samples: number; stepsPerSec: number; dt: number; nx: number; ny: number; stick: { maxForce: number } | null }

export async function renderClip(page: Page, url: string, s: Sample): Promise<Record<string, unknown>> {
  const errors: string[] = [];
  const onErr = (e: Error) => errors.push(e.message);
  page.on("pageerror", onErr);
  const t0 = Date.now();
  await page.goto(`${url}/?${toQuery(s.params)}`);
  let st: Status | null = null;
  for (;;) {
    await page.waitForTimeout(250);
    st = await page.evaluate<Status | null>("window.steelsim ? window.steelsim.status() : null");
    const noGpu = await page.evaluate<string | null>("window.steelsim ? window.steelsim.noGpu() : null");
    if (noGpu) throw new Error(`no WebGPU: ${noGpu}`);
    if (errors.length) throw new Error(`page errors: ${errors.join("; ")}`);
    if (st && st.phase === "done") break;
    if (Date.now() - t0 > JOB_TIMEOUT_MS) throw new Error("timeout");
  }
  page.off("pageerror", onErr);
  const audio = new Float32Array(await page.evaluate<number[]>("window.steelsim.audio()"));
  const link = await page.evaluate<string>("window.steelsim.link()");
  checkLink(link, s.params);
  const wall = (Date.now() - t0) / 1000;
  if (!st) throw new Error("no status");
  if (!isFinite(st.wmax) || st.wmax > 0.1 || !isFinite(st.peakPressure)) throw new Error(`unstable: wmax ${st.wmax}, peak ${st.peakPressure}`);
  let peak = 0, sumsq = 0;
  for (let k = 0; k < audio.length; k++) { peak = Math.max(peak, Math.abs(audio[k])); sumsq += audio[k] * audio[k]; }
  const scale = peak > 0 ? 0.99 / peak : 0;
  const n = Math.min(audio.length, Math.round(DURATION * s.params.fs));
  writeFileSync(join(OUT, "clips", `${s.id}.wav`), wavBytes(audio.subarray(0, n), s.params.fs, scale));
  const p = s.params;
  const mat = makeMaterial(p);
  return {
    id: s.id,
    index: s.index,
    seed: SEED,
    // Targets
    E: p.E, rho: p.rho, h: p.h, Lx: p.Lx, Ly: p.Ly, nu: p.nu,
    // Derived
    kappa: Math.sqrt(mat.D / mat.rhoh),
    mass: mat.rhoh * p.Lx * p.Ly,
    // Nuisance
    features: p.features.length ? encodeFeatures(p.features) : "none",
    shapeScale: s.shapeScale, peakToPeak: s.peakToPeak, maxSlope: s.maxSlope,
    sigma0: p.sigma0, sigma1: p.sigma1, sigmaIn: p.sigmaIn,
    strikeFracX: p.strikeFracX, strikeFracY: p.strikeFracY, strikeVelocity: p.strikeVelocity,
    listenerDist: p.listenerDist, listenerAzimuth: p.listenerAzimuth, listenerElevation: p.listenerElevation,
    // Numerics
    spacing: p.spacing, nx: st.nx, ny: st.ny, dt: st.dt, fps: p.fps, fs: p.fs, duration: DURATION, totalSteps: st.totalSteps,
    // Audio
    samples: n, peakPa: peak, rmsPa: Math.sqrt(sumsq / Math.max(1, audio.length)), wavScalePa: peak > 0 ? 1 / scale : 0,
    maxForce: st.stick?.maxForce ?? null, wmax: st.wmax,
    // Bookkeeping
    wallSeconds: wall, stepsPerSec: st.stepsPerSec, link, finished: new Date().toISOString(),
  };
}

// ---------- Main ----------

async function main(): Promise<void> {
  mkdirSync(join(OUT, "clips"), { recursive: true });
  const indexFile = join(OUT, "index.jsonl");
  const done = new Set<string>();
  if (existsSync(indexFile)) for (const line of readFileSync(indexFile, "utf8").split("\n")) if (line.trim()) done.add(JSON.parse(line).id);
  const existing = new Set(readdirSync(join(OUT, "clips")).map((f) => f.replace(/\.wav$/, "")));
  const jobs: Sample[] = [];
  for (let i = START; i < START + COUNT; i++) {
    const s = drawSample(SEED, i);
    if (done.has(s.id) && existing.has(s.id)) continue;
    jobs.push(s);
  }
  console.log(`${jobs.length} clips to render (${COUNT - jobs.length} already done), ${WORKERS} workers, out ${OUT}`);
  if (!jobs.length) return;
  const url = await serveDist();
  let next = 0, finished = 0, failed = 0;
  const t0 = Date.now();
  const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
  // Each worker owns a browser, so that a lost GPU process costs one worker a
  // relaunch rather than every worker its WebGPU; the starts are staggered.
  const worker = async (w: number) => {
    await sleep(1500 * w);
    let browser: Browser = await chromium.launch({ headless: true, args: ARGS });
    let page: Page = await (await browser.newContext()).newPage();
    const relaunch = async () => {
      try { await browser.close(); } catch { /* already gone */ }
      await sleep(3000);
      browser = await chromium.launch({ headless: true, args: ARGS });
      page = await (await browser.newContext()).newPage();
    };
    while (next < jobs.length) {
      const s = jobs[next++];
      for (let attempt = 0; ; attempt++) {
        try {
          const rec = await renderClip(page, url, s);
          appendFileSync(indexFile, JSON.stringify(rec) + "\n");
          finished++;
          const rate = (finished * 3600) / ((Date.now() - t0) / 1000);
          console.log(`[w${w}] ${s.id} done in ${(rec.wallSeconds as number).toFixed(0)} s: ${s.nx}x${s.ny}, E ${(s.params.E / 1e9).toFixed(0)} GPa, rho ${s.params.rho.toFixed(0)}, h ${(s.params.h * 1e3).toFixed(2)} mm, ${(s.params.Lx * 100).toFixed(0)}x${(s.params.Ly * 100).toFixed(0)} cm, peak ${(rec.peakPa as number).toFixed(2)} Pa | ${finished} done, ${failed} failed, ${rate.toFixed(0)} clips/h`);
          break;
        } catch (e) {
          const msg = String(e).replace(/\n/g, " ");
          if (/no WebGPU|closed|crashed|Protocol error|timeout/i.test(msg) && attempt < 3) {
            console.log(`[w${w}] ${s.id} retrying after: ${msg}`);
            await relaunch();
            continue;
          }
          failed++;
          console.log(`[w${w}] ${s.id} FAILED: ${msg}`);
          appendFileSync(join(OUT, "failures.log"), `${s.id}\t${new Date().toISOString()}\t${msg}\n`);
          break;
        }
      }
    }
    await browser.close();
  };
  await Promise.all(Array.from({ length: WORKERS }, (_, w) => worker(w)));
  console.log(`finished: ${finished} clips, ${failed} failures, ${((Date.now() - t0) / 60000).toFixed(1)} min`);
  process.exit(0);
}

if (process.argv[1] && fileURLToPath(import.meta.url) === process.argv[1]) void main().catch((e) => { console.error(e); process.exit(1); });