// GPU-versus-CPU self-test: one time step of the WGSL kernel on a smooth,
// strongly nonlinear random state is compared with the float64 reference.
// Runs in the browser (?selftest=1) and is driven by tests/gpu.test.ts.

import { defaultParams, makeMaterial, nodeWeight, type PhysParams } from "./params.ts";
import { initialShape, parseFeatures } from "./shape.ts";
import { accelerations, addDamping, fields } from "./sim/cpuRef.ts";
import { Simulation } from "./sim/simulation.ts";

export interface SelfTestCase {
  name: string;
  /** Maximum relative error of the acceleration of each component, over the sheet. */
  accError: { u: number; v: number; w: number };
  /** Maximum error of the compensated w update relative to the largest increment dt * w_t. */
  positionError: number;
  nodes: number;
  dt: number;
}

export const ACC_TOLERANCE = 1e-3;
export const POSITION_TOLERANCE = 1e-2;

const cases: { name: string; overrides: Partial<PhysParams> }[] = [
  { name: "rectangle, flat", overrides: { shape: "rectangle", features: [] } },
  { name: "rectangle, bowed across the width", overrides: { shape: "rectangle" } },
  { name: "ellipse, dish", overrides: { shape: "ellipse", features: parseFeatures("bow:x,0.0015,0;bow:y,0.0015,0")! } },
  { name: "rectangle, twist, dent, crease, and random field", overrides: { shape: "rectangle", features: parseFeatures("bow:x,0.002,0.5;twist:0.001;dent:0.3,0.6,0.03,0.0015;crease:0.6,0.4,25,-1.5,0.015;noise:11,0.0005,0.05")! } },
  { name: "rectangle, bowed, in-plane mass x4", overrides: { shape: "rectangle", inPlaneMass: 4 } },
];

/** A smooth random field: a few sinusoids across the sheet plus a little white noise. */
function smoothField(nx: number, ny: number, amplitude: number, noise: number, rnd: () => number): Float64Array {
  const f = new Float64Array(nx * ny);
  const terms = Array.from({ length: 6 }, () => ({ kx: 1 + Math.floor(4 * rnd()), ky: 1 + Math.floor(4 * rnd()), px: 2 * Math.PI * rnd(), py: 2 * Math.PI * rnd(), a: rnd() - 0.5 }));
  for (let j = 0; j < ny; j++) {
    for (let i = 0; i < nx; i++) {
      let s = 0;
      for (const t of terms) s += t.a * Math.cos((t.kx * Math.PI * i) / (nx - 1) + t.px) * Math.cos((t.ky * Math.PI * j) / (ny - 1) + t.py);
      f[j * nx + i] = amplitude * s + noise * (rnd() - 0.5);
    }
  }
  return f;
}

export async function runSelfTest(device: GPUDevice): Promise<SelfTestCase[]> {
  const results: SelfTestCase[] = [];
  let seed = 2024;
  const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 4294967296);
  for (const c of cases) {
    const params: PhysParams = { ...defaultParams, ...c.overrides, spacing: 0.012 };
    const sim = new Simulation(device, params);
    const g = sim.grid;
    const n = g.n;
    const m = makeMaterial(params);
    const w0 = Float64Array.from(initialShape(params, g));

    // State in float32 (what the GPU sees), then the same values in float64 for the reference.
    const pos = new Float32Array(4 * n), vel = new Float32Array(3 * n);
    const fw = smoothField(g.nx, g.ny, 2e-3, 5e-5, rnd); // about 2.5 thicknesses: strongly nonlinear
    const fu = smoothField(g.nx, g.ny, 2e-5, 1e-6, rnd);
    const fv = smoothField(g.nx, g.ny, 2e-5, 1e-6, rnd);
    const fdw = smoothField(g.nx, g.ny, 0.5, 1e-2, rnd);
    const fdu = smoothField(g.nx, g.ny, 1e-2, 1e-3, rnd);
    const fdv = smoothField(g.nx, g.ny, 1e-2, 1e-3, rnd);
    for (let k = 0; k < n; k++) {
      if (nodeWeight(g, k % g.nx, Math.floor(k / g.nx)) === 0) continue;
      pos[k] = fu[k]; pos[n + k] = fv[k]; pos[2 * n + k] = fw[k]; pos[3 * n + k] = 0;
      vel[k] = fdu[k]; vel[n + k] = fdv[k]; vel[2 * n + k] = fdw[k];
    }
    const x = fields(n), v = fields(n);
    for (let k = 0; k < n; k++) {
      x.u[k] = pos[k]; x.v[k] = pos[n + k]; x.w[k] = pos[2 * n + k];
      v.u[k] = vel[k]; v.v[k] = vel[n + k]; v.w[k] = vel[2 * n + k];
    }

    sim.writeState(pos, vel);
    await sim.stepOnce();
    const out = await sim.readState();
    sim.destroy();

    // Reference: a = elastic (in-plane scaled) + losses, v' = v + dt a, w' = w + dt v'.
    const acc = fields(n);
    accelerations(g, m, x, acc, w0);
    for (let k = 0; k < n; k++) { acc.u[k] /= params.inPlaneMass; acc.v[k] /= params.inPlaneMass; }
    addDamping(g, params.sigma0, params.sigma1, params.sigmaIn, v, acc);
    const dt = sim.dt;
    const err = { u: 0, v: 0, w: 0 };
    let scale = { u: 0, v: 0, w: 0 };
    let posErr = 0, posScale = 0;
    for (let k = 0; k < n; k++) {
      if (nodeWeight(g, k % g.nx, Math.floor(k / g.nx)) === 0) continue;
      const gu = (out.vel[k] - vel[k]) / dt, gv = (out.vel[n + k] - vel[n + k]) / dt, gw = (out.vel[2 * n + k] - vel[2 * n + k]) / dt;
      err.u = Math.max(err.u, Math.abs(gu - acc.u[k]));
      err.v = Math.max(err.v, Math.abs(gv - acc.v[k]));
      err.w = Math.max(err.w, Math.abs(gw - acc.w[k]));
      scale = { u: Math.max(scale.u, Math.abs(acc.u[k])), v: Math.max(scale.v, Math.abs(acc.v[k])), w: Math.max(scale.w, Math.abs(acc.w[k])) };
      const vNew = v.w[k] + dt * acc.w[k];
      const wNew = x.w[k] + dt * vNew;
      posErr = Math.max(posErr, Math.abs(out.pos[2 * n + k] + out.pos[3 * n + k] - wNew));
      posScale = Math.max(posScale, Math.abs(dt * vNew));
    }
    results.push({
      name: c.name,
      accError: { u: err.u / scale.u, v: err.v / scale.v, w: err.w / scale.w },
      positionError: posErr / posScale,
      nodes: n,
      dt,
    });
  }
  return results;
}

export function selfTestPassed(r: SelfTestCase): boolean {
  return r.accError.u < ACC_TOLERANCE && r.accError.v < ACC_TOLERANCE && r.accError.w < ACC_TOLERANCE && r.positionError < POSITION_TOLERANCE;
}