/**
 * The test suite. Two layers:
 *
 *  - unit checks on the host: the expression evaluator, the bake container
 *    round-trip, the bake key's tuning-blindness.
 *  - physics validation: each instrument, rendered from its own bake, must
 *    match the timestepped soundbox solver for the equivalent scene, run
 *    through numbl's MATLAB-to-WebGPU pipeline on the same device. That is
 *    the sense in which numbl stays the source of truth for the physics:
 *    the instrument files' hand-parameterized WGSL is held, numerically,
 *    to the compiled solver's output. The gate is soundbox's own 5e-3.
 *
 * The soundbox import assumes the vault layout (a sibling checkout at
 * ../soundbox with its node_modules installed).
 */
import { parseBake, serializeBake } from '../src/host/bake.ts';
import { runPhase } from '../src/host/engine.ts';
import { raymarchWgsl, slotPartials } from '../src/host/view.ts';
import { spectrumOf } from '../src/fft.ts';
import { evalDerived, evalExpr } from '../src/expr.ts';
import { requestDevice } from '../src/host/device.ts';
import { InstrumentSession } from '../src/host/session.ts';
import { parseInstrument, paramValues, type InstrumentFile } from '../src/format.ts';
import { buildInstrument } from '../author/build.ts';
import { errMsg, installWebGpu } from './nodeWebGpu.ts';
import { readFileSync } from 'node:fs';

// soundbox: the same scenes through the numbl pipeline's timestepped solver.
import { StudioSession, specParams } from '../../soundbox/src/studio/session.ts';
import { presetByKey, cloneSpec } from '../../soundbox/src/studio/presets.ts';
import type { Params, SceneSpec } from '../../soundbox/src/studio/types.ts';

let passed = 0;
let failed = 0;
const check = (name: string, ok: boolean, detail?: string): void => {
  console.log(`  ${ok ? 'ok ' : 'FAIL'} ${name}${!ok && detail ? ` — ${detail}` : ''}`);
  if (ok) passed++;
  else failed++;
};

const relErr = (a: Float32Array, b: Float32Array, n: number): number => {
  let d = 0;
  let r = 0;
  for (let i = 0; i < n; i++) {
    d += (a[i] - b[i]) * (a[i] - b[i]);
    r += b[i] * b[i];
  }
  return Math.sqrt(d / Math.max(r, 1e-300));
};

/* ------------------------------------------------------- unit checks --- */

console.log('host checks');
{
  const s = { nx: 96, Lx: 1 };
  check('expr arithmetic', evalExpr('nx / 2 + 3 * 4', s) === 60);
  check('expr functions', evalExpr('max(5, min(256, floor(0.32 / 0.002) + 1))', {}) === 161);
  check('expr clamp/round', evalExpr('clamp(round(2.6), 0, 2)', {}) === 2);
  let threw = false;
  try {
    evalExpr('nx + undefined_name', s);
  } catch {
    threw = true;
  }
  check('expr rejects unknown names', threw);
}
{
  const bake = {
    key: 'k',
    instrument: 'x',
    scope: { a: 1.5 },
    arrays: { kernels: Float32Array.from([1, 2, 3.5]), other: Float32Array.from([-1]) },
  };
  const round = parseBake(serializeBake(bake));
  check(
    'bake container round-trips',
    round.key === 'k' &&
      round.scope.a === 1.5 &&
      round.arrays.kernels.length === 3 &&
      round.arrays.kernels[2] === 3.5 &&
      round.arrays.other[0] === -1,
  );
}

const instruments: Record<string, InstrumentFile> = {};
for (const name of ['drum', 'dulcimer', 'chime']) {
  const text = readFileSync(new URL(`../instruments/${name}.instrument.json`, import.meta.url), 'utf8');
  instruments[name] = parseInstrument(text);
  const rebuilt = buildInstrument(
    name === 'drum'
      ? { name: 'drum', description: instruments[name].description, object: 'membrane', body: 'shell' }
      : name === 'dulcimer'
        ? { name, description: instruments[name].description, object: 'string', body: 'box', defaults: { mic_x: 0.12, mic_y: 0.08, mic_z: 0.1 } }
        : { name, description: instruments[name].description, object: 'bar', body: 'none', defaults: { mic_x: 0.12, mic_y: 0.08, mic_z: 0.1 } },
  );
  check(
    `${name}.instrument.json is the builder's output`,
    `${JSON.stringify(rebuilt, null, 2)}\n` === text,
    'run npm run author',
  );
}

console.log(await installWebGpu());
const device = await requestDevice();

{
  const sess = (v: Record<string, number>): InstrumentSession =>
    new InstrumentSession(device, instruments.drum, paramValues(instruments.drum, v));
  const kA = await sess({}).bakeKey();
  const kB = await sess({ f0: 200, strikex: 0.7, t60: 0.4 }).bakeKey();
  const kC = await sess({ shellr: 0.2 }).bakeKey();
  const kD = await sess({ mic_x: 0.1 }).bakeKey();
  const kE = await sess({ nx: 64 }).bakeKey();
  check('tuning does not touch the bake key', kA === kB);
  check('geometry touches the bake key', kA !== kC && kA !== kD);
  check('accuracy touches the bake key', kA !== kE);
}

/* -------------------------------------------- solver-matched renders --- */

const NX = 32;
const STEPS = 6000;

/** The timestepped reference: the equivalent soundbox scene, compiled by
 *  numbl, stepped on the same device, recorded at the microphone. */
async function reference(
  presetKey: string,
  edits: Params,
  mic: { x: number; y: number; z: number },
): Promise<Float32Array> {
  const spec: SceneSpec = cloneSpec(presetByKey(presetKey)!.spec);
  const s = await StudioSession.create({ device, spec, nx: NX, params: edits });
  s.setMic(mic.x, mic.y, mic.z);
  s.excite();
  s.step(STEPS);
  await s.sync();
  const trace = await s.recorder.read();
  s.destroy();
  return trace;
}

interface Case {
  name: string;
  preset: string;
  /** instrawave parameter overrides beyond nx/ir_seconds. */
  set: Record<string, number>;
  /** the same edits in soundbox's suffixed names. */
  edits: Params;
  mic: { x: number; y: number; z: number };
}

const mic1 = { x: 0.12, y: 0.08, z: 0.1 };
const micDrum = { x: 0.15, y: 0.1, z: 0.12 };
const cases: Case[] = [
  { name: 'chime', preset: 'chime', set: {}, edits: {}, mic: mic1 },
  { name: 'dulcimer', preset: 'dulcimer', set: {}, edits: {}, mic: mic1 },
  { name: 'drum', preset: 'drum', set: {}, edits: {}, mic: micDrum },
];

console.log(`\nagainst the timestepped solver (nx = ${NX}, ${STEPS} steps)`);
const bakes: Record<string, Awaited<ReturnType<InstrumentSession['precompute']>>> = {};
for (const c of cases) {
  const inst = instruments[c.name];
  const session = new InstrumentSession(device, inst, paramValues(inst, c.set));
  const dt = session.dt;
  // A kernel at least as long as the trace, so truncation plays no part.
  session.values = paramValues(inst, {
    ...c.set,
    nx: NX,
    ir_seconds: (STEPS + 100) * dt,
    ...(c.name === 'drum' ? { mode_margin: 4 } : {}),
  });
  const bake = await session.precompute();
  bakes[c.name] = bake;
  const { trace } = await session.render(bake, STEPS * session.dt);
  const ref = await reference(c.preset, c.edits, c.mic);
  const n = Math.min(trace.length, ref.length, STEPS);
  const err = relErr(trace, ref, n);
  console.log(`  ${c.name}: ${bake.scope.n_rows} kernel rows; rel err ${err.toExponential(2)}`);
  check(`${c.name} matches the solver`, err < 5e-3, `rel err ${err.toExponential(2)}`);
}

// The default (geometric) drum mode set, as a number: what the band limit
// at the default mode_margin costs against the full solver. Measured on the
// development machine: 5.3e-2 at margin 1.25, 4.0e-3 at 2 (the default),
// and a 2.4e-3 f32 floor from 2.5 up.
{
  const inst = instruments.drum;
  const margin = inst.params.find((p) => p.key === 'mode_margin')!.value;
  const session = new InstrumentSession(device, inst, paramValues(inst, {}));
  const dt = session.dt;
  session.values = paramValues(inst, { nx: NX, ir_seconds: (STEPS + 100) * dt });
  const bake = await session.precompute();
  const { trace } = await session.render(bake, STEPS * session.dt);
  const ref = await reference('drum', {}, micDrum);
  const err = relErr(trace, ref, Math.min(trace.length, ref.length, STEPS));
  console.log(`  drum at default mode_margin ${margin}: ${bake.scope.n_rows} rows; rel err ${err.toExponential(2)}`);
  check('the default drum mode set stays close', err < 8e-3, `rel err ${err.toExponential(2)}`);
}

// Retune and restrike from the same bake: tuning moves, the bake does not.
{
  const inst = instruments.drum;
  const set = { f0: 150, t60: 0.7, sig1: 0.002, strikex: 0.62, striker: 0.03 };
  const session = new InstrumentSession(
    device,
    inst,
    paramValues(inst, {
      ...set,
      nx: NX,
      ir_seconds: (STEPS + 100) * new InstrumentSession(device, inst, paramValues(inst, {})).dt,
      mode_margin: 4,
    }),
  );
  const { trace } = await session.render(bakes.drum, STEPS * session.dt);
  const ref = await reference(
    'drum',
    { f0_m1: 150, t60_m1: 0.7, sig1_m1: 0.002, strikex_m1: 0.62, striker_m1: 0.03 },
    micDrum,
  );
  const err = relErr(trace, ref, Math.min(trace.length, ref.length, STEPS));
  console.log(`  retuned, restruck drum from the same bake: rel err ${err.toExponential(2)}`);
  check('the same bake serves a retuned drum', err < 5e-3, `rel err ${err.toExponential(2)}`);
}

// The visual section: the medium fields run headless through the same
// engine, the raymarch shader compiles, the markers evaluate.
console.log('\nvisual');
for (const name of ['drum', 'dulcimer', 'chime'] as const) {
  const inst = instruments[name];
  const visual = inst.visual!;
  // nx = 64: at 32 the default walls are under half a cell and the mask
  // never saturates (a real property, which the viewer draws as a haze).
  const scope = evalDerived(inst.derived, paramValues(inst, { nx: 64 }));
  const fields = await runPhase({
    device,
    shaders: inst.shaders,
    phase: { buffers: visual.buffers, setup: visual.setup },
    scope,
    outputs: visual.volumes.map((v) => v.buffer),
    label: `${name}-visual`,
  });
  const wall = fields.wall;
  let solid = 0;
  let air = 0;
  for (const v of wall) {
    if (v < 0.5) solid++;
    if (v > 0.95) air++;
  }
  const hasBody = name !== 'chime';
  check(
    `${name}: the medium fields draw ${hasBody ? 'a body in air' : 'open air'}`,
    air > 0.5 * wall.length && (hasBody ? solid > 100 : solid === 0),
    `${solid} solid cells, ${air} air cells of ${wall.length}`,
  );
  let profPeak = 0;
  for (const v of fields[visual.volumes[1].buffer]) profPeak = Math.max(profPeak, Math.abs(v));
  check(`${name}: the coupling profile is present`, profPeak > 0, `peak ${profPeak}`);

  device.pushErrorScope('validation');
  const module = device.createShaderModule({ code: raymarchWgsl(visual) });
  await device.createRenderPipelineAsync({
    layout: 'auto',
    vertex: { module, entryPoint: 'vmain' },
    fragment: {
      module,
      entryPoint: 'fmain',
      targets: [{ format: 'bgra8unorm' }],
      constants: {
        NX: evalExpr(visual.dims[0], scope),
        NY: evalExpr(visual.dims[1], scope),
        NZ: evalExpr(visual.dims[2], scope),
      },
    },
  });
  const err = await device.popErrorScope();
  check(`${name}: the raymarch shader compiles`, err === null, err?.message);
  const half = visual.extents.map((e) => evalExpr(e, scope) / 2);
  let markersOk = true;
  for (const m of visual.markers) {
    const pts = m.kind === 'point' ? [m.at] : m.kind === 'line' ? [m.from, m.to] : [m.corner];
    for (const p of pts) {
      const v = p.map((x) => evalExpr(x, scope));
      if (v.some((x, i) => !Number.isFinite(x) || Math.abs(x) > half[i] * 1.001)) markersOk = false;
    }
  }
  check(`${name}: the markers land inside the domain`, markersOk);
}

// The spectrum tools: a pure sine peaks at its own bin, and a rendered
// chime's strongest measured peak lands on the strongest predicted partial.
{
  const dt = 1 / 48000;
  const sine = new Float32Array(1 << 15);
  for (let i = 0; i < sine.length; i++) sine[i] = Math.sin(2 * Math.PI * 440 * i * dt);
  const { db, binHz } = spectrumOf(sine, dt);
  let best = 0;
  for (let i = 1; i < db.length; i++) if (db[i] > db[best]) best = i;
  check(
    'the spectrum of a 440 Hz sine peaks at 440 Hz',
    Math.abs(best * binHz - 440) < 2 * binHz,
    `peak at ${(best * binHz).toFixed(1)} Hz`,
  );
}

// Stimuli are LTI: two events must equal the sum of their single renders.
{
  const inst = instruments.chime;
  const session = new InstrumentSession(device, inst, paramValues(inst, { nx: NX }));
  const dt = session.dt;
  const bake = await session.precompute();
  const seconds = 0.4;
  const t1 = 0.15;
  const evB = { t: t1, set: { f1: 620, strikepos: 0.44 } };
  const both = await session.render(bake, seconds, [{ t: 0 }, evB]);
  const a = await session.render(bake, seconds, [{ t: 0 }]);
  const b = await session.render(bake, seconds, [evB]);
  const sum = new Float32Array(both.trace.length);
  for (let i = 0; i < sum.length; i++) sum[i] = a.trace[i] + b.trace[i];
  const err = relErr(both.trace, sum, sum.length);
  console.log(`  two stimuli vs summed singles: rel err ${err.toExponential(2)}`);
  check('stimuli superpose', err < 1e-5, `rel err ${err.toExponential(2)}`);
  let head = 0;
  const startB = Math.round(t1 / dt);
  for (let i = 0; i < Math.min(startB, b.trace.length); i++) head = Math.max(head, Math.abs(b.trace[i]));
  check('a later event is silent before its time', head === 0, `peak ${head}`);

  // The predicted partials, probed from the slots buffer, against the
  // measured spectrum of the same render.
  const vp = inst.visual!.partials!;
  const r = await session.render(bake, 1.2, [{ t: 0 }], [vp.buffer]);
  const partials = slotPartials(r.probes[vp.buffer], vp, r.dt);
  const { db, binHz } = spectrumOf(r.trace, r.dt);
  let best = 0;
  for (let i = 1; i < db.length; i++) if (db[i] > db[best]) best = i;
  const peakHz = best * binHz;
  const nearest = partials.reduce((a, p) => (Math.abs(p.freq - peakHz) < Math.abs(a.freq - peakHz) ? p : a), partials[0]);
  console.log(
    `  chime partials: ${partials.map((p) => p.freq.toFixed(0)).join(', ')} Hz; ` +
      `measured peak ${peakHz.toFixed(0)} Hz`,
  );
  check(
    'the measured peak lands on a predicted partial',
    partials.length >= 1 && Math.abs(peakHz - nearest.freq) < Math.max(3 * binHz, 0.01 * nearest.freq),
    `peak ${peakHz.toFixed(1)} Hz vs nearest predicted ${nearest?.freq.toFixed(1)} Hz`,
  );
}

console.log(`\n${passed} passed, ${failed} failed`);
process.exit(failed === 0 ? 0 : 1);