// Tests of the design script: the functions a script defines for itself,
// compilation and the checks, the evaluation on the grid, and the presets.

import { test } from "node:test";
import assert from "node:assert/strict";
import { DesignError, compileDesign, evaluatePlate, profileStats, removePlane, runDesignScript } from "../src/design.ts";
import { PRESETS } from "../src/presets.ts";
import { countInside, insideNode, mountPatch, nodeWeight, nodeX, nodeY, thicknessRange } from "../src/params.ts";

const minimal = `design.material.E = 170e9; design.material.nu = 0.29; design.material.rho = 7900;
design.bounds.width = 0.3; design.bounds.height = 0.2;
design.thickness = 1e-3;`;

/** A scalar from the workspace. */
function num(vars: Record<string, unknown>, name: string): number {
  const v = vars[name] as number | { data: Float64Array };
  return typeof v === "number" ? v : v.data[0];
}
function data(vars: Record<string, unknown>, name: string): Float64Array {
  return (vars[name] as { data: Float64Array }).data;
}

test("a minimal script compiles with defaults filled in", () => {
  const d = compileDesign(minimal);
  assert.equal(d.name, "Untitled design");
  assert.equal(d.hasInside, false);
  assert.equal(d.mount, null);
  assert.equal(d.thicknessConst, 1e-3);
  assert.equal(d.profileConst, 0);
  assert.deepEqual(d.loss, { sigma0: 1.15, sigma1: 5e-4, sigmaIn: 20 });
  assert.ok(Math.abs(d.Lx - 0.3) < 1e-12 && Math.abs(d.Ly - 0.2) < 1e-12);
});

test("a script may define its own functions, and they are vectorized", () => {
  const v = runDesignScript(`
mm = 1e-3;
s = [smoothstep(0, 1, -1), smoothstep(0, 1, 2), smoothstep(0, 1, 0.5)];
p0 = dome(0, 0.05, 0.01);
p1 = dome(0.06, 0.05, 0.01);
[X, Y] = meshgrid(-0.1:0.05:0.1, -0.1:0.05:0.1);
dd = dents(X, Y, [0 0.05], [0 0], [0.02 0.02], [1e-3 -1e-3]);
u = seededrand(5, 3); u2 = seededrand(5, 3);

function v = ramp(a, b, x)
v = min(max((x - a) ./ (b - a), 0), 1);
end
function v = smoothstep(e0, e1, x)
t = ramp(e0, e1, x);
v = t.^2 .* (3 - 2*t);
end
function v = dome(d, radius, height)
v = 0.5 * height * (1 + cos(pi * d / radius)) .* (d < radius);
end
function v = dents(x, y, cx, cy, radius, height)
v = zeros(size(x));
for k = 1:numel(cx)
  v = v + height(k) * exp(-((x - cx(k)).^2 + (y - cy(k)).^2) / radius(k)^2);
end
end
function u = seededrand(seed, n)
s = floor(seed);
u = zeros(1, n);
for k = 1:n
  s = mod(s * 1664525 + 1013904223, 4294967296);
  u(k) = s / 4294967296;
end
end
`);
  assert.deepEqual(Array.from(data(v, "s")), [0, 1, 0.5], "one local function may call another");
  assert.ok(Math.abs(num(v, "p0") - 0.01) < 1e-12);
  assert.equal(num(v, "p1"), 0, "the dome vanishes beyond its rim");
  const dd = data(v, "dd"); // 5 x 5, column-major; the center sits at index 12
  assert.ok(Math.abs(dd[12] - (1e-3 - 1e-3 * Math.exp(-(0.05 ** 2) / 0.02 ** 2))) < 1e-15, "two dents of opposite sign");
  assert.deepEqual(Array.from(data(v, "u")), Array.from(data(v, "u2")), "the seeded generator is deterministic");
  assert.ok(data(v, "u").every((x) => x >= 0 && x < 1));
});

test("a handle may call a function the script defines below it", () => {
  // The trailer that samples the fields is appended after the script's own
  // functions, and it calls the handles from inside a function of its own.
  const d = compileDesign(`design.material.E = 170e9; design.material.nu = 0.29; design.material.rho = 7900;
design.bounds.width = 0.2; design.bounds.height = 0.2;
t0 = 1e-3;
design.thickness = @(x, y, r, theta) t0 * (1 + taper(r));
design.inside = @(x, y, r, theta) r <= 0.1;

function v = taper(r)
v = min(max(r / 0.1, 0), 1);
end
`);
  const pl = evaluatePlate(d, 0.01);
  const hr = thicknessRange(pl);
  assert.ok(Math.abs(hr.min - 1e-3) < 2e-4 && Math.abs(hr.max - 2e-3) < 1e-4, `thickness ${hr.min}..${hr.max}`);
});

test("scripts that do not describe a plate are rejected with a message", () => {
  const bad: [string, RegExp][] = [
    [`x = 5;`, /struct named design/],
    [`design = 5;`, /design must be a struct/],
    [`x = (1 + ;`, /line 1/],
    [`error('boom');`, /boom/],
    [`design.bounds.width = 0.3; design.bounds.height = 0.2; design.thickness = 1e-3;`, /material/],
    [`design.material.E = 170e9; design.material.nu = 0.6; design.material.rho = 7900; design.bounds.width = 0.3; design.bounds.height = 0.2; design.thickness = 1e-3;`, /nu/],
    [`design.material.E = 170e9; design.material.nu = 0.3; design.material.rho = 7900; design.bounds.width = 0.3; design.thickness = 1e-3;`, /bounds.height/],
    [`design.material.E = 170e9; design.material.nu = 0.3; design.material.rho = 7900; design.bounds.width = 0.3; design.bounds.height = 0.2;`, /thickness/],
    [`${minimal} design.thickness = 'thin';`, /thickness must be a number or a function handle/],
    [`${minimal} design.thickness = -1;`, /thickness must be positive/],
    [`${minimal} design.inside = 5;`, /inside must be a function handle/],
    [`${minimal} design.mount.stiffness = 1e4;`, /mount.radius/],
    [`${minimal} design.loss = 3;`, /loss must be a struct/],
  ];
  for (const [src, re] of bad) assert.throws(() => compileDesign(src), (e: unknown) => e instanceof DesignError && re.test(e.message), src);
});

test("evaluation rejects fields that are not positive, not finite, not vectorized, or that throw", () => {
  const zero = compileDesign(`${minimal} design.thickness = @(x, y, r, theta) x;`);
  assert.throws(() => evaluatePlate(zero, 0.02), /thickness must be positive/);
  const nan = compileDesign(`${minimal} design.profile = @(x, y, r, theta) 0*x./0;`);
  assert.throws(() => evaluatePlate(nan, 0.02), /profile is NaN/);
  const shape = compileDesign(`${minimal} design.profile = @(x, y, r, theta) [1 2 3];`);
  assert.throws(() => evaluatePlate(shape, 0.02), /profile returned a 1 x 3 array/);
  const thrower = compileDesign(`${minimal} design.inside = @(x, y, r, theta) no_such_function(x);`);
  assert.throws(() => evaluatePlate(thrower, 0.02), /inside: .*no_such_function/);
  const empty = compileDesign(`${minimal} design.inside = @(x, y, r, theta) false & x;`);
  assert.throws(() => evaluatePlate(empty, 0.02), /no grid cells/);
  const scalarFn = evaluatePlate(compileDesign(`${minimal} design.thickness = @(x, y, r, theta) 2e-3;`), 0.02);
  assert.equal(thicknessRange(scalarFn).min, 2e-3, "a handle returning a scalar is spread over the grid");
});

test("a circle with a hole excludes the center and the corners", () => {
  const d = compileDesign(`design.material.E = 170e9; design.material.nu = 0.3; design.material.rho = 7900;
design.bounds.width = 0.2; design.bounds.height = 0.2; design.thickness = 1e-3;
design.inside = @(x, y, r, theta) (r <= 0.1) & (r >= 0.015);`);
  const pl = evaluatePlate(d, 0.01);
  const g = pl.grid;
  assert.equal(g.nx, 21);
  const c = (g.nx - 1) / 2;
  assert.ok(!insideNode(g, c, c), "center node lies in the hole");
  assert.ok(!insideNode(g, 0, 0), "corner node lies outside the disc");
  assert.ok(insideNode(g, c, 3), "a node on the radius belongs to the disc");
  let weight = 0;
  for (let j = 0; j < g.ny; j++) for (let i = 0; i < g.nx; i++) weight += nodeWeight(g, i, j);
  const area = weight * g.hx * g.hy;
  const exact = Math.PI * (0.1 ** 2 - 0.015 ** 2);
  assert.ok(Math.abs(area / exact - 1) < 0.05, `the ${countInside(g)} nodes cover ${area} vs ${exact}`);
  assert.equal(thicknessRange(pl).max, 1e-3);
});

test("fields are sampled at the right nodes (meshgrid orientation)", () => {
  const d = compileDesign(`${minimal} design.thickness = @(x, y, r, theta) 10 + x + 10*y;`);
  const pl = evaluatePlate(d, 0.05);
  const g = pl.grid;
  assert.ok(g.nx !== g.ny, "the grid is not square, so a transposed sampling would show");
  for (let j = 0; j < g.ny; j++) for (let i = 0; i < g.nx; i++) {
    const want = 10 + nodeX(g, i) + 10 * nodeY(g, j);
    assert.ok(Math.abs(pl.h[j * g.nx + i] - want) < 1e-12, `node ${i}, ${j}: ${pl.h[j * g.nx + i]} vs ${want}`);
  }
});

test("the profile is made plane-free and the mount patch is found", () => {
  const d = compileDesign(`${minimal}
design.profile = @(x, y, r, theta) 0.01 + 0.05*x - 0.02*y + 0.001*x.^2;
design.mount.radius = 0.02; design.mount.stiffness = 1e4; design.mount.damping = 1;`);
  const pl = evaluatePlate(d, 0.01);
  const g = pl.grid;
  let s = 0, sx = 0, sy = 0;
  for (let j = 0; j < g.ny; j++) {
    for (let i = 0; i < g.nx; i++) {
      const wgt = nodeWeight(g, i, j), v = pl.w0[j * g.nx + i];
      s += wgt * v; sx += wgt * nodeX(g, i) * v; sy += wgt * nodeY(g, j) * v;
    }
  }
  assert.ok(Math.abs(s) < 1e-12 && Math.abs(sx) < 1e-12 && Math.abs(sy) < 1e-12, "mean plane removed");
  const stats = profileStats(pl);
  assert.ok(stats.peakToPeak > 0 && stats.peakToPeak < 1e-4, `peak to peak ${stats.peakToPeak}`);
  const mp = mountPatch(pl)!;
  assert.ok(mp.nodes.length >= 9 && mp.nodes.length <= 21, `mount patch has ${mp.nodes.length} nodes`);
  assert.ok(mp.weight > 0);
});

test("removePlane leaves a plane-free field unchanged", () => {
  const g = { nx: 7, ny: 5, hx: 0.05, hy: 0.05, n: 35 };
  const f = new Float64Array(g.n);
  for (let j = 0; j < g.ny; j++) for (let i = 0; i < g.nx; i++) f[j * g.nx + i] = nodeX(g, i) ** 2 - 0.0125;
  const before = Float64Array.from(f);
  removePlane(f, g);
  // The parabola x^2 has a nonzero mean but no tilt; only the mean should move.
  const shift = f[0] - before[0];
  for (let k = 0; k < g.n; k++) assert.ok(Math.abs(f[k] - before[k] - shift) < 1e-12);
});

test("the presets' noise field is deterministic with unit RMS", () => {
  // The same generator the dented crash uses, as a script-local function
  // returning a handle that calls another script-local function.
  const v = runDesignScript(`
a = noise2d(7, 0.05); b = noise2d(7, 0.05); c = noise2d(8, 0.05);
[X, Y] = meshgrid(-0.5:0.01:0.49, -0.5:0.01:0.49);
A = a(X, Y); B = b(X, Y); C = c(X, Y);
same = isequal(A, B); different = ~isequal(A, C);
rms = sqrt(mean(A(:).^2));

function u = seededrand(seed, n)
s = floor(seed);
u = zeros(1, n);
for k = 1:n
  s = mod(s * 1664525 + 1013904223, 4294967296);
  u(k) = s / 4294967296;
end
end
function f = noise2d(seed, len)
n = 64;
u = reshape(seededrand(seed, 5*n), 5, n);
kx = sqrt(-2*log(1 - u(1,:))) .* cos(2*pi*u(2,:)) * sqrt(2) / len;
ky = sqrt(-2*log(1 - u(3,:))) .* cos(2*pi*u(4,:)) * sqrt(2) / len;
ph = 2*pi*u(5,:);
f = @(x, y) noisesum(kx, ky, ph, x, y);
end
function v = noisesum(kx, ky, ph, x, y)
v = zeros(size(x));
for k = 1:numel(kx)
  v = v + cos(kx(k) * x + ky(k) * y + ph(k));
end
v = v / sqrt(numel(kx) / 2);
end
`);
  assert.equal(v.same, true);
  assert.equal(v.different, true);
  const rms = num(v, "rms");
  assert.ok(Math.abs(rms - 1) < 0.2, `rms ${rms}`);
});

test("every preset compiles and evaluates on the default grid", () => {
  for (const p of PRESETS) {
    const d = compileDesign(p.source);
    const pl = evaluatePlate(d, 0.008);
    const hr = thicknessRange(pl);
    assert.ok(hr.min > 0 && hr.max < 5e-3, `${p.id}: thickness ${hr.min}..${hr.max}`);
    assert.ok(countInside(pl.grid) > 100, `${p.id}: ${countInside(pl.grid)} nodes`);
    const stats = profileStats(pl);
    assert.ok(stats.peakToPeak < 0.05, `${p.id}: peak to peak ${stats.peakToPeak}`);
    if (pl.mount) {
      const mp = mountPatch(pl)!;
      assert.ok(mp.nodes.length > 1, `${p.id}: mount patch has ${mp.nodes.length} nodes`);
    }
  }
});