// Tests of the spatial discretization on the float64 reference (no GPU needed).

import { test } from "node:test";
import assert from "node:assert/strict";
import { accelerations, addDamping, fields, maxAngularFrequency, potentialEnergy, uniformPlate, type Fields } from "../src/sim/cpuRef.ts";
import { type Grid, type Material, type Plate, countInside, insideNode, membraneStiffness, mountPatch, nodeMass, nodeX, nodeY, withEllipseMask, withMask } from "../src/params.ts";

// Steel with nu = 0.3, the value the reference tables use; 0.79 mm thick.
const mat: Material = { E: 170e9, nu: 0.3, rho: 7900 };
const H = 0.79e-3;

function rectangle(nx: number, ny: number, Lx = 0.3, Ly = 0.2): Grid {
  return { nx, ny, hx: Lx / (nx - 1), hy: Ly / (ny - 1), n: nx * ny };
}

function lcg(seed: number): () => number {
  return () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 4294967296 - 0.5);
}

function massOf(pl: Plate, k: number): number {
  return nodeMass(pl, k % pl.grid.nx, Math.floor(k / pl.grid.nx));
}

/** A plate whose thickness rises linearly from 0.5 to 1.5 mm across the width. */
function taperedPlate(g: Grid, w0?: Float64Array, mount: Plate["mount"] = null): Plate {
  const pl = uniformPlate(g, mat, H, w0);
  for (let j = 0; j < g.ny; j++) {
    for (let i = 0; i < g.nx; i++) {
      if (insideNode(g, i, j)) pl.h[j * g.nx + i] = 1e-3 * (1 + nodeX(g, i) / ((g.nx - 1) * g.hx));
    }
  }
  return { ...pl, mount };
}

/** Eigenvalues of a symmetric matrix (row-major, n x n) by cyclic Jacobi. */
function symEig(A: Float64Array, n: number): number[] {
  const a = Float64Array.from(A);
  for (let sweep = 0; sweep < 60; sweep++) {
    let off = 0;
    for (let p = 0; p < n; p++) for (let q = p + 1; q < n; q++) off += a[p * n + q] ** 2;
    if (off < 1e-22) break;
    for (let p = 0; p < n; p++) {
      for (let q = p + 1; q < n; q++) {
        const apq = a[p * n + q];
        if (Math.abs(apq) < 1e-300) continue;
        const theta = (a[q * n + q] - a[p * n + p]) / (2 * apq);
        const t = Math.sign(theta || 1) / (Math.abs(theta) + Math.sqrt(theta * theta + 1));
        const c = 1 / Math.sqrt(t * t + 1), s = t * c;
        for (let k = 0; k < n; k++) {
          const akp = a[k * n + p], akq = a[k * n + q];
          a[k * n + p] = c * akp - s * akq;
          a[k * n + q] = s * akp + c * akq;
        }
        for (let k = 0; k < n; k++) {
          const apk = a[p * n + k], aqk = a[q * n + k];
          a[p * n + k] = c * apk - s * aqk;
          a[q * n + k] = s * apk + c * aqk;
        }
      }
    }
  }
  const ev: number[] = [];
  for (let k = 0; k < n; k++) ev.push(a[k * n + k]);
  return ev.sort((x, y) => x - y);
}

/** Mass-symmetrized linear stiffness M^{-1/2} K M^{-1/2} of one block ("w" or "uv") over the nodes of the plate. */
function stiffness(pl: Plate, block: "w" | "uv"): { K: Float64Array; n: number } {
  const g = pl.grid;
  const nodes: number[] = [];
  for (let k = 0; k < g.n; k++) if (massOf(pl, k) > 0) nodes.push(k);
  const N = nodes.length;
  const n = block === "w" ? N : 2 * N;
  const K = new Float64Array(n * n);
  const x = fields(g.n), y = fields(g.n);
  const eps = 1e-9;
  for (let col = 0; col < n; col++) {
    x.u.fill(0); x.v.fill(0); x.w.fill(0);
    if (block === "w") x.w[nodes[col]] = eps;
    else if (col < N) x.u[nodes[col]] = eps;
    else x.v[nodes[col - N]] = eps;
    accelerations(pl, x, y);
    const mc = massOf(pl, nodes[col % N]);
    for (let row = 0; row < n; row++) {
      const acc = block === "w" ? y.w[nodes[row]] : row < N ? y.u[nodes[row]] : y.v[nodes[row - N]];
      K[row * n + col] = (-acc / eps) * Math.sqrt(massOf(pl, nodes[row % N]) / mc);
    }
  }
  return { K, n };
}

function checkSymmetric(K: Float64Array, n: number): void {
  let asym = 0, scale = 0;
  for (let r = 0; r < n; r++) {
    for (let c = 0; c < n; c++) {
      asym = Math.max(asym, Math.abs(K[r * n + c] - K[c * n + r]));
      scale = Math.max(scale, Math.abs(K[r * n + c]));
    }
  }
  assert.ok(asym < 1e-6 * scale, `asymmetry ${asym} vs scale ${scale}`);
}

/** `gap` is the smallest ratio of the fourth eigenvalue to the largest that counts as non-rigid (a tapered plate on a fine grid has softer low modes). */
function checkSymmetricWithThreeRigidModes(K: Float64Array, n: number, gap = 1e-4): number[] {
  checkSymmetric(K, n);
  const ev = symEig(K, n);
  assert.ok(Math.abs(ev[2]) < 1e-6 * ev[n - 1], "third rigid mode");
  assert.ok(ev[3] > gap * ev[n - 1], `fourth mode is not rigid (${ev[3]} vs ${ev[n - 1]})`);
  return ev;
}

/** Largest relative mismatch between the forces and the centered-difference gradient of the energy. */
function gradientMismatch(pl: Plate, f: Fields): number {
  const g = pl.grid;
  const acc = fields(g.n);
  accelerations(pl, f, acc);
  const delta = 1e-7;
  let worst = 0;
  for (const [arr, accArr] of [[f.w, acc.w], [f.u, acc.u], [f.v, acc.v]] as const) {
    for (let k = 0; k < g.n; k += 5) {
      const mass = massOf(pl, k);
      if (mass === 0) continue;
      const orig = arr[k];
      arr[k] = orig + delta;
      const Ep = potentialEnergy(pl, f);
      arr[k] = orig - delta;
      const Em = potentialEnergy(pl, f);
      arr[k] = orig;
      const Ffd = -(Ep - Em) / (2 * delta);
      const Fan = accArr[k] * mass;
      worst = Math.max(worst, Math.abs(Ffd - Fan) / (Math.abs(Fan) + 1e-3));
    }
  }
  return worst;
}

/** A random state with w of about 25 thicknesses (strongly nonlinear) and a cylindrical 3 mm bow. */
function randomState(g: Grid, seed: number): { f: Fields; w0: Float64Array } {
  const rnd = lcg(seed);
  const f = fields(g.n);
  const w0 = new Float64Array(g.n);
  for (let k = 0; k < g.n; k++) {
    const x = (k % g.nx) / (g.nx - 1) - 0.5;
    w0[k] = 0.003 * (1 - 4 * x * x);
    f.w[k] = 0.02 * rnd();
    f.u[k] = 1e-4 * rnd();
    f.v[k] = 1e-4 * rnd();
  }
  return { f, w0 };
}

test("bending stiffness is symmetric and has exactly three rigid modes", () => {
  const { K, n } = stiffness(uniformPlate(rectangle(9, 7), mat, H), "w");
  checkSymmetricWithThreeRigidModes(K, n);
});

test("membrane stiffness is symmetric and has exactly three rigid modes", () => {
  const { K, n } = stiffness(uniformPlate(rectangle(7, 6), mat, H), "uv");
  checkSymmetricWithThreeRigidModes(K, n);
});

test("free square plate frequencies converge to the reference values", () => {
  // Completely free square plate, nu = 0.3, lambda^2 = omega a^2 sqrt(rho h / D)
  // (Gorman 1978; Leissa's Ritz values are slightly higher). The scheme is
  // second order: the errors are about 0.8 percent at 20 cells per side and
  // 0.3 percent at 30 (the pair at 34.80 is degenerate).
  const ref = [13.468, 19.596, 24.271, 34.801, 34.801, 61.093, 61.093];
  const a = 0.3;
  const D = (mat.E * H ** 3) / (12 * (1 - mat.nu ** 2));
  const errors: number[][] = [];
  for (const N of [16, 21]) {
    const { K, n } = stiffness(uniformPlate(rectangle(N, N, a, a), mat, H), "w");
    const ev = symEig(K, n);
    const lam2 = ev.slice(3, 3 + ref.length).map((e) => Math.sqrt(e) * a * a * Math.sqrt((mat.rho * H) / D));
    errors.push(lam2.map((v, k) => Math.abs(v - ref[k]) / ref[k]));
  }
  for (let k = 0; k < ref.length; k++) {
    assert.ok(errors[1][k] < 0.03, `mode ${k}: error ${(errors[1][k] * 100).toFixed(2)}% at 20 cells`);
    assert.ok(errors[1][k] < 0.7 * errors[0][k], `mode ${k}: error does not shrink from 15 to 20 cells`);
  }
});

test("forces are the negative gradient of the discrete energy, flat and bowed", () => {
  const g = rectangle(8, 7);
  const { f, w0 } = randomState(g, 7);
  assert.ok(gradientMismatch(uniformPlate(g, mat, H), f) < 1e-4, "flat");
  const bowed = uniformPlate(g, mat, H, w0);
  assert.ok(gradientMismatch(bowed, f) < 1e-4, "bowed");
  // The stress-free shape itself carries no energy.
  assert.ok(potentialEnergy(bowed, fields(g.n)) < 1e-12, "w0 is stress free");
});

test("loss terms only remove energy", () => {
  const g = rectangle(9, 8);
  const pl = { ...uniformPlate(g, mat, H), loss: { sigma0: 1.0, sigma1: 1e-3, sigmaIn: 5.0 } };
  const rnd = lcg(3);
  const v = fields(g.n), acc = fields(g.n);
  for (let k = 0; k < g.n; k++) { v.u[k] = rnd(); v.v[k] = rnd(); v.w[k] = rnd(); }
  addDamping(pl, v, acc);
  let power = 0;
  for (let k = 0; k < g.n; k++) power += massOf(pl, k) * (v.u[k] * acc.u[k] + v.v[k] * acc.v[k] + v.w[k] * acc.w[k]);
  assert.ok(power < 0, `power ${power} should be negative`);
});

test("maximum frequency estimate is close to the in-plane CFL estimate", () => {
  const g: Grid = { nx: 50, ny: 50, hx: 0.004, hy: 0.004, n: 2500 };
  const wmax = maxAngularFrequency(g, mat, H);
  const cp = Math.sqrt(membraneStiffness(mat, H) / (mat.rho * H));
  const cfl = (2 * cp * Math.SQRT2) / 0.004; // 2 c sqrt(1/hx^2 + 1/hy^2)
  assert.ok(wmax > 0.5 * cfl && wmax < 1.5 * cfl, `wmax ${wmax} vs cfl ${cfl}`);
  // The in-plane wave speed does not depend on the thickness, so neither does the estimate.
  const wmax2 = maxAngularFrequency(g, mat, 2 * H);
  assert.ok(Math.abs(wmax2 - wmax) < 0.05 * wmax, `wmax ${wmax} vs ${wmax2} at twice the thickness`);
});

// ---------- Variable thickness ----------

test("tapered plate: both stiffness blocks are symmetric with exactly three rigid modes", () => {
  const pl = taperedPlate(rectangle(9, 7));
  const b = stiffness(pl, "w");
  checkSymmetricWithThreeRigidModes(b.K, b.n);
  const m = stiffness(taperedPlate(rectangle(7, 6)), "uv");
  checkSymmetricWithThreeRigidModes(m.K, m.n);
});

test("tapered plate: forces are the energy gradient, flat and bowed", () => {
  const g = rectangle(8, 7);
  const { f, w0 } = randomState(g, 11);
  assert.ok(gradientMismatch(taperedPlate(g), f) < 1e-4, "flat");
  assert.ok(gradientMismatch(taperedPlate(g, w0), f) < 1e-4, "bowed");
});

test("a uniform plate of twice the thickness has twice the bending frequencies", () => {
  // Bending frequencies scale with h, membrane frequencies do not.
  const thin = stiffness(uniformPlate(rectangle(9, 8), mat, H), "w");
  const thick = stiffness(uniformPlate(rectangle(9, 8), mat, 2 * H), "w");
  const e1 = symEig(thin.K, thin.n), e2 = symEig(thick.K, thick.n);
  for (let k = 3; k < 10; k++) assert.ok(Math.abs(e2[k] / e1[k] - 4) < 1e-6, `mode ${k}: ratio ${e2[k] / e1[k]}`);
  const thinM = stiffness(uniformPlate(rectangle(6, 5), mat, H), "uv");
  const thickM = stiffness(uniformPlate(rectangle(6, 5), mat, 2 * H), "uv");
  const m1 = symEig(thinM.K, thinM.n), m2 = symEig(thickM.K, thickM.n);
  for (let k = 3; k < 10; k++) assert.ok(Math.abs(m2[k] / m1[k] - 1) < 1e-6, `membrane mode ${k}: ratio ${m2[k] / m1[k]}`);
});

// ---------- Mount ----------

test("mounted plate: the stiffness is symmetric and the rigid modes feel the mount", () => {
  const g = rectangle(9, 7);
  const pl = taperedPlate(g, undefined, { x: 0, y: 0, radius: 0.04, stiffness: 1e4, damping: 0 });
  assert.ok(mountPatch(pl)!.nodes.length > 1, "the mount covers several nodes");
  for (const block of ["w", "uv"] as const) {
    const { K, n } = stiffness(pl, block);
    checkSymmetric(K, n);
  }
  // A uniform translation of the plate in any direction feels the whole mount stiffness.
  const x = fields(g.n), y = fields(g.n);
  for (const comp of ["u", "v", "w"] as const) {
    x.u.fill(0); x.v.fill(0); x.w.fill(0);
    x[comp].fill(1e-6);
    accelerations(pl, x, y);
    let F = 0;
    for (let k = 0; k < g.n; k++) F += massOf(pl, k) * y[comp][k];
    assert.ok(Math.abs(-F / 1e-6 - 1e4) < 1e-6 * 1e4, `${comp}: mount force ${-F / 1e-6} vs 1e4`);
  }
  // A tilt about the x axis feels a restoring moment from the finite mount patch.
  x.u.fill(0); x.v.fill(0);
  for (let k = 0; k < g.n; k++) x.w[k] = 1e-6 * nodeY(g, Math.floor(k / g.nx));
  accelerations(pl, x, y);
  let M = 0;
  for (let k = 0; k < g.n; k++) M += massOf(pl, k) * y.w[k] * nodeY(g, Math.floor(k / g.nx));
  assert.ok(M < 0, `tilt moment ${M} should restore`);
});

test("mounted plate: forces are the energy gradient and the dashpot removes energy", () => {
  const g = rectangle(8, 7);
  const { f, w0 } = randomState(g, 13);
  const pl = { ...taperedPlate(g, w0, { x: 0.02, y: -0.01, radius: 0.03, stiffness: 5e3, damping: 2 }), loss: { sigma0: 0, sigma1: 0, sigmaIn: 0 } };
  assert.ok(gradientMismatch(pl, f) < 1e-4, "gradient");
  const rnd = lcg(5);
  const v = fields(g.n), acc = fields(g.n);
  for (let k = 0; k < g.n; k++) { v.u[k] = rnd(); v.v[k] = rnd(); v.w[k] = rnd(); }
  addDamping(pl, v, acc);
  let power = 0, touched = 0;
  for (let k = 0; k < g.n; k++) {
    power += massOf(pl, k) * (v.u[k] * acc.u[k] + v.v[k] * acc.v[k] + v.w[k] * acc.w[k]);
    if (acc.w[k] !== 0) touched++;
  }
  assert.ok(power < 0, `power ${power} should be negative`);
  assert.ok(touched > 0 && touched < countInside(g), "the dashpot acts on the mount nodes only");
});

// ---------- Masked grids ----------

test("elliptical plate: both stiffness blocks are symmetric with exactly three rigid modes", () => {
  const g = withEllipseMask(rectangle(13, 10));
  const pl = uniformPlate(g, mat, H);
  assert.ok(countInside(g) < g.n && countInside(g) > g.n / 2, "mask removes some nodes");
  for (let k = 0; k < g.n; k++) {
    // Every node of the plate carries mass (no whiskers left after the mask is trimmed).
    if (g.mask![k]) assert.ok(massOf(pl, k) > 0, `node ${k} has no mass`);
  }
  const b = stiffness(pl, "w");
  checkSymmetricWithThreeRigidModes(b.K, b.n);
  const m = stiffness(pl, "uv");
  checkSymmetricWithThreeRigidModes(m.K, m.n);
});

test("disc with a center hole: symmetric stiffness with three rigid modes, forces are the energy gradient", () => {
  const nx = 15, a = 0.1;
  const g0: Grid = { nx, ny: nx, hx: (2 * a) / (nx - 1), hy: (2 * a) / (nx - 1), n: nx * nx };
  const g = withMask(g0, (x, y) => Math.hypot(x, y) <= a && Math.hypot(x, y) >= 0.02);
  assert.ok(!insideNode(g, 7, 7), "the center node lies in the hole");
  const pl = taperedPlate(g);
  const b = stiffness(pl, "w");
  checkSymmetricWithThreeRigidModes(b.K, b.n, 1e-5);
  const { f } = randomState(g, 9);
  assert.ok(gradientMismatch(pl, f) < 1e-4, "gradient");
});

test("free circular plate frequencies approach the reference values", () => {
  const nx = 25, a = 0.15; // radius; 24 cells across the diameter
  const g = withEllipseMask({ nx, ny: nx, hx: (2 * a) / (nx - 1), hy: (2 * a) / (nx - 1), n: nx * nx });
  const { K, n } = stiffness(uniformPlate(g, mat, H), "w");
  const ev = symEig(K, n);
  const D = (mat.E * H ** 3) / (12 * (1 - mat.nu ** 2));
  const lam2 = ev.slice(3, 12).map((e) => Math.sqrt(e) * a * a * Math.sqrt((mat.rho * H) / D));
  // Completely free circular plate, nu = 0.3 (Leissa, Table 2.4): modes (2,0) x2, (0,1),
  // (3,0) x2, (1,1) x2, (4,0) x2 by nodal diameters and circles. The staircase edge is
  // first-order accurate: about 2 percent here, under 1 percent at 60 cells across.
  const ref = [5.358, 5.358, 9.003, 12.439, 12.439, 20.475, 20.475, 21.835, 21.835];
  for (let k = 0; k < ref.length; k++) {
    const relErr = Math.abs(lam2[k] - ref[k]) / ref[k];
    assert.ok(relErr < 0.04, `mode ${k}: got ${lam2[k].toFixed(3)}, expected ${ref[k]} (err ${(relErr * 100).toFixed(2)}%)`);
  }
});

test("node coordinates are centered on the bounding box", () => {
  const g = rectangle(9, 7, 0.4, 0.3);
  assert.ok(Math.abs(nodeX(g, 0) + 0.2) < 1e-12 && Math.abs(nodeX(g, 8) - 0.2) < 1e-12);
  assert.ok(Math.abs(nodeY(g, 0) + 0.15) < 1e-12 && Math.abs(nodeY(g, 6) - 0.15) < 1e-12);
});