// 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, type Fields } from "../src/sim/cpuRef.ts";
import { type Grid, type Material, countInside, nodeWeight, withEllipseMask } from "../src/params.ts";

// 0.79 mm steel with nu = 0.3, the value the reference tables use.
const mat: Material = { D: 8.66, C: 1.665e8, G: 5.91e7, nu: 0.3, rhoh: 6.241 };

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(g: Grid, k: number): number {
  return mat.rhoh * g.hx * g.hy * nodeWeight(g, k % g.nx, Math.floor(k / g.nx));
}

/** 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 sheet. */
function stiffness(g: Grid, block: "w" | "uv"): { K: Float64Array; n: number } {
  const nodes: number[] = [];
  for (let k = 0; k < g.n; k++) if (massOf(g, 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(g, mat, x, y);
    const mc = massOf(g, 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(g, nodes[row % N]) / mc);
    }
  }
  return { K, n };
}

function checkSymmetricWithThreeRigidModes(K: Float64Array, n: number): number[] {
  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}`);
  const ev = symEig(K, n);
  assert.ok(Math.abs(ev[2]) < 1e-6 * ev[n - 1], "third rigid mode");
  assert.ok(ev[3] > 1e-4 * ev[n - 1], `fourth mode is not rigid (${ev[3]} vs ${ev[n - 1]})`);
  return ev;
}

/** Largest relative mismatch between the forces and the centerd-difference gradient of the energy. */
function gradientMismatch(g: Grid, f: Fields, w0?: Float64Array): number {
  const acc = fields(g.n);
  accelerations(g, mat, f, acc, w0);
  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(g, k);
      if (mass === 0) continue;
      const orig = arr[k];
      arr[k] = orig + delta;
      const Ep = potentialEnergy(g, mat, f, w0);
      arr[k] = orig - delta;
      const Em = potentialEnergy(g, mat, f, w0);
      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(rectangle(9, 7), "w");
  checkSymmetricWithThreeRigidModes(K, n);
});

test("membrane stiffness is symmetric and has exactly three rigid modes", () => {
  const { K, n } = stiffness(rectangle(7, 6), "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 errors: number[][] = [];
  for (const N of [16, 21]) {
    const { K, n } = stiffness(rectangle(N, N, a, a), "w");
    const ev = symEig(K, n);
    const lam2 = ev.slice(3, 3 + ref.length).map((e) => Math.sqrt(e) * a * a * Math.sqrt(mat.rhoh / mat.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(g, f) < 1e-4, "flat");
  assert.ok(gradientMismatch(g, f, w0) < 1e-4, "bowed");
  // The stress-free shape itself carries no energy.
  assert.ok(potentialEnergy(g, mat, fields(g.n), w0) < 1e-12, "w0 is stress free");
});

test("loss terms only remove energy", () => {
  const g = rectangle(9, 8);
  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(g, 1.0, 1e-3, 5.0, v, acc);
  let power = 0;
  for (let k = 0; k < g.n; k++) power += massOf(g, 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);
  const cp = Math.sqrt(mat.C / mat.rhoh);
  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}`);
});

// ---------- Masked (elliptical) grids ----------

test("elliptical sheet: both stiffness blocks are symmetric with exactly three rigid modes", () => {
  const g = withEllipseMask(rectangle(13, 10));
  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 sheet carries mass (no whiskers left after the mask is trimmed).
    if (g.mask![k]) assert.ok(massOf(g, k) > 0, `node ${k} has no mass`);
  }
  const b = stiffness(g, "w");
  checkSymmetricWithThreeRigidModes(b.K, b.n);
  const m = stiffness(g, "uv");
  checkSymmetricWithThreeRigidModes(m.K, m.n);
});

test("elliptical sheet: forces are the energy gradient, flat and bowed", () => {
  const g = withEllipseMask(rectangle(12, 9));
  const { f, w0 } = randomState(g, 5);
  assert.ok(gradientMismatch(g, f) < 1e-4, "flat");
  assert.ok(gradientMismatch(g, f, w0) < 1e-4, "bowed");
});

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(g, "w");
  const ev = symEig(K, n);
  const lam2 = ev.slice(3, 12).map((e) => Math.sqrt(e) * a * a * Math.sqrt(mat.rhoh / mat.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)}%)`);
  }
});