/**
 * What the solver is held to.
 *
 * These run against the real pipeline — MATLAB source, numbl lowering,
 * generated WGSL, GPU — so they check the whole chain rather than any one
 * piece of it. The checks are physical wherever they can be: a wave should
 * travel at the speed the medium says, a scatterer that matches its
 * background should not scatter, an absorbing layer should absorb. What is
 * left over is the discretization's error, and that is what the numbers here
 * bound.
 *
 * Everything is in SI, at the app's own domain size (`DOMAIN`, 10 m) and
 * background speed (`C_AIR`, 343 m/s) — not a separate "toy" scale — so a
 * number that appears here means the same thing it would in the app, and the
 * synthetic scenes below (`uniform`, `closedBox`) share the app's own
 * absorbing-layer profile rather than inventing their own.
 *
 * Nothing here checks what the app looks like. That is for a browser.
 */
import { ModelSession } from '../src/mgpu/session.ts';
import type { MModel } from '../src/mgpu/registry.ts';
import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
import type { MScene } from '../src/scene/registry.ts';
import { mSceneByKey, defaultSceneParams } from '../src/scene/registry.ts';
import { C_AIR, DOMAIN } from '../src/units.ts';

export type Check = (name: string, ok: boolean, detail: string) => void;
export type Log = (s: string) => void;

/** A homogeneous medium with the usual absorbing edge — the case every
 *  scattering result is measured against. Same sponge profile as every real
 *  scene (src/tools/sponge.m via 0.2*L, 1700), so its absorption behaviour is
 *  exactly the app's, not a separately tuned stand-in. */
const uniform: MScene = {
  key: 'uniform',
  label: 'Uniform',
  blurb: 'No scatterer at all.',
  params: [],
  source: `
function [c, sig] = medium(x, y, L, c0)
  c = c0 + 0*x;
  sig = sponge(x, y, L, 0.2*L, 1700);
end
`,
};

/** The same, with nothing absorbing anywhere: a closed box. */
const closedBox: MScene = {
  ...uniform,
  key: 'closed',
  source: `
function [c, sig] = medium(x, y, c0)
  c = c0 + 0*x;
  sig = 0*x;
end
`,
};

/** A model whose "step" is one application of a Laplacian stencil, so the
 *  stencil can be measured on a field we chose. */
const stencilProbe = (op: 'lap2' | 'lap4'): MModel => ({
  key: `probe-${op}`,
  label: `probe ${op}`,
  blurb: '',
  state: ['p', 'pm', 't'],
  params: [],
  order: op === 'lap2' ? 2 : 4,
  source: `
function [p, pm, t] = init(npts)
  p = zeros(npts, 1);
  pm = zeros(npts, 1);
  t = zeros(npts, 1);
end

function [pn, pold, tn] = step(p, pm, t)
  pn = ${op}(p);
  pold = pm;
  tn = t;
end
`,
});

const maxAbs = (a: Float32Array): number => {
  let m = 0;
  for (const v of a) m = Math.max(m, Math.abs(v));
  return m;
};

/**
 * The Laplacian stencils, against a field whose Laplacian is known exactly.
 *
 * p = sin(kx*x) * sin(ky*y) has lap(p) = -(kx^2 + ky^2) * p. Both stencils
 * should reproduce that away from the boundary, the 5-point one to O((k*h)^2)
 * and the 9-point one to O((k*h)^4) — which at the resolution used here is
 * two orders of magnitude tighter. The margin matters more than either
 * number: it is what says the fourth-order stencil is actually fourth order
 * and not a mistyped second-order one. Purely a statement about the discrete
 * operator, so it does not depend on the medium or the domain size at all.
 */
export async function stencilChecks(device: GPUDevice, check: Check, log: Log): Promise<void> {
  const n = 64;
  const kx = 4 * Math.PI;
  const ky = 3 * Math.PI;
  const errs: Record<string, number> = {};

  for (const op of ['lap2', 'lap4'] as const) {
    const session = await ModelSession.create({
      device,
      model: stencilProbe(op),
      params: {},
      scene: uniform,
      sceneParams: {},
      n,
    });
    try {
      const { grid } = session;
      const p = new Float32Array(grid.npts);
      for (let i = 0; i < grid.npts; i++) {
        p[i] = Math.sin(kx * grid.x64[i]) * Math.sin(ky * grid.y64[i]);
      }
      session.reset();
      session.gpu.upload('p', p);
      session.step(1);
      const got = await session.read('p');

      // Interior only: the stencil takes the field outside the grid to be
      // zero, which is not what this analytic field does.
      const pad = 3;
      const want = -(kx * kx + ky * ky);
      let err = 0;
      for (let iy = pad; iy < n - pad; iy++) {
        for (let ix = pad; ix < n - pad; ix++) {
          const i = ix + n * iy;
          err = Math.max(err, Math.abs(got[i] - want * p[i]));
        }
      }
      errs[op] = err / Math.abs(want);
      log(`  ${op}: max relative error ${errs[op].toExponential(2)} at n = ${n}`);
    } finally {
      session.destroy();
    }
  }

  check(
    'lap2 matches the analytic Laplacian',
    errs.lap2 < 0.02,
    `relative error ${errs.lap2.toExponential(2)} (expect ~(k*h)^2/12)`,
  );
  check(
    'lap4 is far more accurate than lap2',
    errs.lap4 < errs.lap2 / 10,
    `${errs.lap4.toExponential(2)} vs ${errs.lap2.toExponential(2)}`,
  );
}

/**
 * A pulse from a point source should be a ring of radius c*(t - t0).
 *
 * This is the end-to-end statement that the thing solves the wave equation:
 * it exercises the source term, the model's own clock, the stencil, and the
 * timestep the host computed, and it fails if any of them is wrong by a
 * constant factor. 600 Hz on this grid (256 points over 10 m, h = 3.9 cm) is
 * about 15 cells per wavelength — comfortably resolved, and short enough that
 * a clean, narrow pulse is cheap to run.
 */
export async function propagationChecks(
  device: GPUDevice,
  check: Check,
  log: Log,
): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const t0 = 0.005;
  const session = await ModelSession.create({
    device,
    model,
    params: {
      ...defaultParams(model),
      f: 600,
      tw: 0.0012,
      t0,
      cw: 0,
      point: 1,
      x0: 0,
      y0: 0,
      w: 0.15,
    },
    scene: uniform,
    sceneParams: {},
    n: 512,
    L: DOMAIN,
  });
  try {
    session.reset();
    // Early enough that the front is still clear of the absorbing layer,
    // which starts at |x| = 0.3*DOMAIN and would pull the peak back towards
    // the interior.
    const until = 0.01;
    const steps = Math.round(until / session.dt);
    session.step(steps);
    const p = await session.read('p');
    const t = session.steps * session.dt;

    // Where the wavefront is, along +x from the source at the origin. An
    // energy centroid rather than the bare peak: at 600 Hz the wavelength is
    // 0.57 m, so the tallest individual fringe of an oscillating pulse can
    // sit anywhere within half a wavelength of the envelope's true centre
    // depending on carrier phase, which swamps the grid's own ~1% resolution.
    // Weighting position by p^2 averages over the fringes instead of picking
    // whichever one happens to be tallest.
    const { grid } = session;
    const cutoff = 0.275 * grid.L; // stays clear of the sponge, which starts at 0.3*L
    const iy = Math.floor(grid.ny / 2);
    let weighted = 0;
    let weight = 0;
    for (let ix = Math.floor(grid.nx / 2); ix < grid.nx; ix++) {
      const i = ix + grid.nx * iy;
      if (grid.x64[i] > cutoff) break;
      const w2 = p[i] * p[i];
      weighted += grid.x64[i] * w2;
      weight += w2;
    }
    const best = weighted / weight;
    const want = C_AIR * (t - t0);
    const err = Math.abs(best - want) / want;
    log(`  wavefront at r = ${best.toFixed(4)} m, expected ${want.toFixed(4)} m at t = ${(1000 * t).toFixed(3)} ms`);
    check(
      'a pulse travels at the medium speed',
      err < 0.03,
      `radius off by ${(100 * err).toFixed(1)}%`,
    );
  } finally {
    session.destroy();
  }
}

/**
 * The absorbing layer should leave next to nothing behind.
 *
 * A plane pulse is launched, crosses the grid, and is swallowed. What is
 * still in the interior long afterwards is what the sponge reflected, and it
 * is the honest measure of how open the open boundary is.
 */
export async function boundaryChecks(
  device: GPUDevice,
  check: Check,
  log: Log,
): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const session = await ModelSession.create({
    device,
    model,
    params: {
      ...defaultParams(model),
      f: 300, tw: 0.006, t0: 0.03, cw: 0, point: 0, x0: -0.25 * DOMAIN,
    },
    scene: uniform,
    sceneParams: {},
    n: 256,
    L: DOMAIN,
  });
  try {
    session.reset();
    const stepsTo = (t: number): number => Math.round(t / session.dt);
    session.step(stepsTo(0.07));
    const peak = maxAbs(await session.read('p'));

    session.step(stepsTo(0.3) - session.steps);
    const after = await session.read('p');

    // The interior only: the sponge itself is allowed to hold whatever it is
    // busy absorbing.
    const { grid } = session;
    const interior = 0.3 * grid.L; // exactly where the sponge starts
    let residual = 0;
    for (let i = 0; i < grid.npts; i++) {
      if (Math.abs(grid.x64[i]) < interior && Math.abs(grid.y64[i]) < interior) {
        residual = Math.max(residual, Math.abs(after[i]));
      }
    }
    const ratio = residual / peak;
    log(`  peak ${peak.toExponential(2)}, interior residual at t = 0.3 s is ${ratio.toExponential(2)} of it`);
    check(
      'the absorbing layer reflects little',
      ratio < 0.02,
      `residual ${(100 * ratio).toFixed(2)}% of the incident peak`,
    );
  } finally {
    session.destroy();
  }
}

/**
 * A scatterer whose speed matches the background is not a scatterer.
 *
 * Running the disk scene at cin = 1 must reproduce the uniform medium
 * exactly, which is a strong statement about the whole scene path: the
 * smoothed interface, the coordinates, the upload. And at cin = 3 there must
 * be a scattered field worth looking at, or the app would be drawing nothing.
 * Uses the real `disk` scene and its real domain, so what is checked is
 * exactly what the app runs.
 */
export async function scatteringChecks(
  device: GPUDevice,
  check: Check,
  log: Log,
): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const disk = mSceneByKey('disk')!;
  const params = {
    ...defaultParams(model),
    f: 400,
    tw: 0.002,
    t0: 0.008,
    cw: 0,
    point: 0,
    x0: -0.25 * DOMAIN,
  };
  const n = 256;
  const run = async (scene: MScene, sceneParams: Record<string, number>): Promise<Float32Array> => {
    const session = await ModelSession.create({
      device, model, params, scene, sceneParams, n, L: DOMAIN,
    });
    try {
      session.reset();
      session.step(Math.round(0.022 / session.dt));
      return await session.read('p');
    } finally {
      session.destroy();
    }
  };

  const plain = await run(uniform, {});
  const matched = await run(disk, { ...defaultSceneParams(disk), cin: 1, absorb: 0 });
  const hard = await run(disk, { ...defaultSceneParams(disk), cin: 3, absorb: 0 });

  const peak = maxAbs(plain);
  let dMatched = 0;
  let dHard = 0;
  for (let i = 0; i < plain.length; i++) {
    dMatched = Math.max(dMatched, Math.abs(matched[i] - plain[i]));
    dHard = Math.max(dHard, Math.abs(hard[i] - plain[i]));
  }
  log(`  matched disk differs by ${(dMatched / peak).toExponential(2)}, hard disk by ${(dHard / peak).toFixed(2)}`);
  check(
    'a speed-matched disk does not scatter',
    dMatched / peak < 1e-3,
    `scattered field ${(dMatched / peak).toExponential(2)} of the incident peak`,
  );
  check(
    'a hard disk scatters strongly',
    dHard / peak > 0.2,
    `scattered field ${(100 * dHard / peak).toFixed(0)}% of the incident peak`,
  );
}

/**
 * The timestep the host picks should be stable, and near the edge of being
 * unstable — a scheme that is merely stable because it is crawling is not
 * evidence of anything. Run a closed box (no absorption at all, so nothing
 * can hide a slow instability) at 95% of the computed limit and watch it
 * bounce around for a long time.
 */
export async function stabilityChecks(
  device: GPUDevice,
  check: Check,
  log: Log,
): Promise<void> {
  for (const key of ['leapfrog', 'leapfrog4']) {
    const model = mModelByKey(key)!;
    const session = await ModelSession.create({
      device,
      model,
      params: {
        ...defaultParams(model),
        f: 300, tw: 0.003, t0: 0.01, cw: 0, point: 1, x0: 0, y0: 0,
      },
      scene: closedBox,
      sceneParams: {},
      n: 128,
      L: DOMAIN,
      cfl: 0.95,
    });
    try {
      session.reset();
      session.step(Math.round(0.03 / session.dt));
      const early = maxAbs(await session.read('p'));
      session.step(Math.round(0.6 / session.dt));
      const late = maxAbs(await session.read('p'));
      log(`  ${key}: max|p| ${early.toExponential(2)} at t = 30 ms, ${late.toExponential(2)} at t = 630 ms`);
      check(
        `${key} is stable at 95% of the CFL limit`,
        Number.isFinite(late) && late < 5 * early,
        `max|p| went from ${early.toExponential(2)} to ${late.toExponential(2)} over 600 ms`,
      );
    } finally {
      session.destroy();
    }
  }
}

/**
 * The planner's kernel splitting must not change the answer.
 *
 * On any device worth running this on, the leapfrog update fits in one
 * kernel. Squeeze the budget down to two grid fields per kernel — what a
 * compatibility-mode device would allow — and the same line has to be
 * evaluated in half a dozen pieces through scratch buffers. The arithmetic is
 * the same; only the rounding of the intermediates differs, since each piece
 * is stored as f32 on the way out.
 */
export async function splitChecks(device: GPUDevice, check: Check, log: Log): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const disk = mSceneByKey('disk')!;
  const params = {
    ...defaultParams(model),
    f: 400, tw: 0.002, t0: 0.008, point: 0, x0: -0.25 * DOMAIN,
  };
  const run = async (operandBudget?: number): Promise<{ p: Float32Array; ops: string[] }> => {
    const session = await ModelSession.create({
      device,
      model,
      params,
      scene: disk,
      sceneParams: defaultSceneParams(disk),
      n: 128,
      L: DOMAIN,
      operandBudget,
    });
    try {
      session.reset();
      session.step(Math.round(0.02 / session.dt));
      return { p: await session.read('p'), ops: session.describe().step };
    } finally {
      session.destroy();
    }
  };

  const whole = await run();
  const split = await run(2);
  const peak = maxAbs(whole.p);
  let diff = 0;
  for (let i = 0; i < whole.p.length; i++) {
    diff = Math.max(diff, Math.abs(whole.p[i] - split.p[i]));
  }
  log(`  ${whole.ops.length} ops whole, ${split.ops.length} split; largest difference ${(diff / peak).toExponential(2)}`);
  check(
    'splitting a kernel does not change the answer',
    diff / peak < 1e-3,
    `fields differ by ${(diff / peak).toExponential(2)} of the peak`,
  );
  check(
    'a squeezed budget really does split the update',
    split.ops.length > whole.ops.length,
    `${split.ops.length} ops vs ${whole.ops.length}`,
  );
}

/**
 * The microphone records the field, at the point it is pointed at, once per
 * timestep.
 *
 * Checked against the field itself rather than against a description of it:
 * after N steps the trace must be N samples long, and its last sample must be
 * exactly — not approximately — the pressure sitting at the probe's grid point,
 * since both are the same f32 written by the same kernel. That pins the probe
 * index, the grid layout, and the fact that the recording dispatch really does
 * run once per step rather than once per submission.
 */
export async function microphoneChecks(
  device: GPUDevice,
  check: Check,
  log: Log,
): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const session = await ModelSession.create({
    device,
    model,
    params: {
      ...defaultParams(model),
      f: 600, tw: 0.0012, t0: 0.003, cw: 0, point: 1, x0: 0, y0: 0, w: 0.15,
    },
    scene: uniform,
    sceneParams: {},
    n: 128,
    L: DOMAIN,
  });
  try {
    const { grid } = session;
    const mx = 1.5;
    const my = 0.5;
    session.setMic(mx, my);
    session.reset();
    // Long enough for the wave to have reached the microphone and moved on.
    const steps = Math.round(0.01 / session.dt);
    session.step(steps);

    const trace = await session.recorder.read();
    const field = await session.read('p');
    const ix = Math.round((mx + grid.L / 2) / grid.h - 0.5);
    const iy = Math.round((my + grid.L / 2) / grid.h - 0.5);
    const at = field[ix + grid.nx * iy];

    let peak = 0;
    for (const v of trace) peak = Math.max(peak, Math.abs(v));
    log(`  ${trace.length} samples in ${steps} steps, peak ${peak.toExponential(2)}, last ${trace[trace.length - 1].toExponential(3)} vs field ${at.toExponential(3)}`);

    check(
      'the microphone records one sample per timestep',
      trace.length === steps,
      `${trace.length} samples for ${steps} steps`,
    );
    check(
      'the microphone records the field at its own grid point',
      trace.length > 0 && trace[trace.length - 1] === at,
      `last sample ${trace[trace.length - 1]} vs field ${at}`,
    );
    check(
      'the microphone hears the wave arrive',
      peak > 1e-3,
      `peak |p| at the microphone was ${peak.toExponential(2)}`,
    );

    // Moving it must move what it hears: the same run sampled at the origin,
    // where a point source is loudest, cannot match a point away from it.
    session.setMic(0, 0);
    session.reset();
    session.step(steps);
    const atSource = await session.recorder.read();
    let peak2 = 0;
    for (const v of atSource) peak2 = Math.max(peak2, Math.abs(v));
    log(`  peak at the source ${peak2.toExponential(2)}, at r = ${Math.hypot(mx, my).toFixed(2)} m ${peak.toExponential(2)}`);
    check(
      'moving the microphone changes what it hears',
      peak2 > peak,
      `${peak2.toExponential(2)} at the source vs ${peak.toExponential(2)} away from it`,
    );
  } finally {
    session.destroy();
  }
}

/**
 * A room rings, and sealing it makes it ring for longer.
 *
 * Measured as energy in the second half of the microphone's trace against the
 * first — a ratio rather than an envelope, because a small room beats between
 * its modes and any one window can land in a null. In the open field the same
 * pulse passes the microphone once and is gone, which is the contrast that
 * makes the number mean something. Uses the real `room` scene's own default
 * geometry, at the app's real domain size.
 *
 * The source is tuned for this test's own grid rather than reused from the
 * scene's `suggest`. A wall slower than the background (`cwall` = 0.15) has a
 * *shorter* wavelength inside itself than the background does at the same
 * frequency, by that same factor — the app's own docs on this scene call this
 * out — and an unresolved wall does not behave like a partial reflector, it
 * behaves like an absorber: at 256 grid points it swallowed the whole pulse
 * in a handful of bounces regardless of the `absorb` parameter, making every
 * room in an early version of this check look identically "sealed" no matter
 * what. So this uses the app's own 512-point grid, where the wall's own
 * wavelength is resolved to about ten cells at 220 Hz rather than four at the
 * scene's demo frequency (700 Hz) — and even then the ring is real but not
 * long: measured, late/early lands around 0.15, well above the open field's
 * 0.004 but short of the naive "rings for a while" threshold a lossless room
 * would give. That is the wall's transmission loss actually doing its job,
 * not a bug — 74% amplitude reflection per bounce (cwall = 0.15 gives
 * |c-1|/(c+1) = 0.74) empties a small room in a few tens of bounces.
 */
export async function roomChecks(device: GPUDevice, check: Check, log: Log): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const room = mSceneByKey('room')!;
  const params = {
    ...defaultParams(model),
    point: 1, x0: -0.7, y0: 0.5, w: 0.03, f: 220, tw: 0.002, t0: 0.008, cw: 0,
  };
  const roomDefaults = defaultSceneParams(room);
  const mic = { x: roomDefaults.side * 0.6, y: -roomDefaults.side * 0.5 };

  const listen = async (scene: MScene, sceneParams: Record<string, number>) => {
    const session = await ModelSession.create({
      device, model, params, scene, sceneParams, n: 512, L: DOMAIN,
    });
    try {
      session.setMic(mic.x, mic.y);
      session.reset();
      session.step(Math.round(0.15 / session.dt));
      const trace = await session.recorder.read();
      const half = Math.floor(trace.length / 2);
      let early = 0;
      let late = 0;
      for (let i = 0; i < half; i++) early += trace[i] * trace[i];
      for (let i = half; i < trace.length; i++) late += trace[i] * trace[i];
      return { early, late, ratio: late / Math.max(early, 1e-30) };
    } finally {
      session.destroy();
    }
  };

  const open = await listen(uniform, {});
  const sealed = await listen(room, { ...roomDefaults, gap: 0 });
  const wide = await listen(room, { ...roomDefaults, gap: 2 * roomDefaults.side });
  log(`  late/early energy — open field ${open.ratio.toExponential(3)}, sealed room ${sealed.ratio.toFixed(3)}, wide door ${wide.ratio.toFixed(3)}`);

  check(
    'a pulse in the open field does not come back',
    open.ratio < 0.02,
    `late/early energy ${open.ratio.toExponential(2)}`,
  );
  check(
    'a room rings after the pulse has passed',
    sealed.ratio > 0.05,
    `late/early energy ${sealed.ratio.toFixed(3)} inside the room, vs ${open.ratio.toExponential(2)} in the open`,
  );
  check(
    'sound leaves through the doorway',
    sealed.late > 1.2 * wide.late,
    `late energy ${sealed.late.toExponential(2)} sealed vs ${wide.late.toExponential(2)} with a wide door`,
  );
}

/**
 * What a step compiles to. A guard on the fusion passes: if one of them stops
 * firing, the model still gives the right answer, only several times slower,
 * and nothing else here would notice.
 */
export async function planChecks(device: GPUDevice, check: Check, log: Log): Promise<void> {
  const model = mModelByKey('leapfrog')!;
  const disk = mSceneByKey('disk')!;
  const session = await ModelSession.create({
    device,
    model,
    params: defaultParams(model),
    scene: disk,
    sceneParams: defaultSceneParams(disk),
    n: 128,
    L: DOMAIN,
  });
  try {
    const ops = session.describe().step;
    for (const line of ops) log(`  ${line}`);
    const kernels = ops.filter((o) => o.startsWith('kernel')).length;
    const stencils = ops.filter((o) => o.startsWith('stencil')).length;
    check(
      'the step uses exactly one stencil dispatch',
      stencils === 1,
      `${stencils} stencil ops`,
    );
    check(
      'the source term fuses into the update',
      kernels <= 6,
      `${kernels} element-wise kernels (5 expected: u, sd, pn, pold, tn)`,
    );
  } finally {
    session.destroy();
  }
}