import { test } from "node:test";
import assert from "node:assert/strict";
import { accelerations, potentialEnergy, maxAngularFrequency, type Fields } from "../src/sim/cpuRef.ts";
import { type Grid, type Material, trapWeight } from "../src/params.ts";

const mat: Material = { D: 8.66, C: 1.665e8, G: 5.91e7, nu: 0.3, rhoh: 6.241 };

function fields(n: number): Fields {
  return { u: new Float64Array(n), v: new Float64Array(n), w: new Float64Array(n) };
}

/** Jacobi eigenvalues of a symmetric matrix (row-major, n x n). */
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 of one block ("w" or "uv"). */
function stiffness(g: Grid, block: "w" | "uv"): { K: Float64Array; n: number } {
  const N = g.n;
  const n = block === "w" ? N : 2 * N;
  const K = new Float64Array(n * n);
  const x = fields(N), y = fields(N);
  const eps = 1e-9;
  const mass = (k: number) => mat.rhoh * g.hx * g.hy * trapWeight(g, k % g.nx, Math.floor(k / g.nx));
  for (let col = 0; col < n; col++) {
    x.u.fill(0); x.v.fill(0); x.w.fill(0);
    if (block === "w") x.w[col] = eps;
    else if (col < N) x.u[col] = eps;
    else x.v[col - N] = eps;
    accelerations(g, mat, x, y);
    const mc = mass(col % N);
    for (let row = 0; row < n; row++) {
      const acc = block === "w" ? y.w[row] : row < N ? y.u[row] : y.v[row - N];
      const mr = mass(row % N);
      // K_sym = M^{1/2} (M^{-1} K) M^{-1/2}
      K[row * n + col] = (-acc / eps) * Math.sqrt(mr / mc);
    }
  }
  return { K, n };
}

test("bending stiffness is symmetric and has exactly three rigid modes", () => {
  const nx = 9, ny = 7;
  const g: Grid = { nx, ny, hx: 0.3 / (nx - 1), hy: 0.2 / (ny - 1), n: nx * ny };
  const { K, n } = stiffness(g, "w");
  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");
});

test("membrane stiffness is symmetric and has exactly three rigid modes", () => {
  const nx = 7, ny = 6;
  const g: Grid = { nx, ny, hx: 0.3 / (nx - 1), hy: 0.2 / (ny - 1), n: nx * ny };
  const { K, n } = stiffness(g, "uv");
  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}`);
  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");
});

test("free square plate frequencies match the Leissa values", () => {
  const nx = 21, ny = 21, a = 0.3;
  const g: Grid = { nx, ny, hx: a / (nx - 1), hy: a / (ny - 1), n: nx * ny };
  const { K, n } = stiffness(g, "w");
  const ev = symEig(K, n);
  const lam2 = ev.slice(3, 8).map((e) => Math.sqrt(e) * a * a * Math.sqrt(mat.rhoh / mat.D));
  // Completely free square plate, nu = 0.3 (Leissa / Gorman): first five nonzero modes.
  const ref = [13.468, 19.596, 24.27, 34.801, 34.801];
  for (let k = 0; k < ref.length; k++) {
    const relErr = Math.abs(lam2[k] - ref[k]) / ref[k];
    assert.ok(relErr < 0.03, `mode ${k}: got ${lam2[k].toFixed(3)}, expected ${ref[k]} (err ${(relErr * 100).toFixed(2)}%)`);
  }
});

test("forces are the negative gradient of the discrete energy (nonlinear state)", () => {
  const nx = 8, ny = 7;
  const g: Grid = { nx, ny, hx: 0.3 / (nx - 1), hy: 0.2 / (ny - 1), n: nx * ny };
  const f = fields(g.n), acc = fields(g.n);
  let seed = 7;
  const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 4294967296 - 0.5);
  for (let k = 0; k < g.n; k++) {
    f.w[k] = 0.02 * rnd(); // ~ 25 thicknesses, strongly nonlinear
    f.u[k] = 1e-4 * rnd();
    f.v[k] = 1e-4 * rnd();
  }
  accelerations(g, mat, 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 = mat.rhoh * g.hx * g.hy * trapWeight(g, k % nx, Math.floor(k / nx));
      const orig = arr[k];
      arr[k] = orig + delta;
      const Ep = potentialEnergy(g, mat, f);
      arr[k] = orig - delta;
      const Em = potentialEnergy(g, mat, 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));
    }
  }
  assert.ok(worst < 1e-4, `gradient mismatch ${worst}`);
});

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}`);
});

test("forces are the energy gradient with an initial-shape imperfection", () => {
  const nx = 8, ny = 7;
  const g: Grid = { nx, ny, hx: 0.3 / (nx - 1), hy: 0.2 / (ny - 1), n: nx * ny };
  const f = fields(g.n), acc = fields(g.n);
  const w0 = new Float64Array(g.n);
  let seed = 11;
  const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 4294967296 - 0.5);
  for (let k = 0; k < g.n; k++) {
    const x = (k % nx) / (nx - 1) - 0.5;
    w0[k] = 0.003 * (1 - 4 * x * x); // cylindrical bow, 3 mm
    f.w[k] = 0.01 * rnd();
    f.u[k] = 1e-4 * rnd();
    f.v[k] = 1e-4 * rnd();
  }
  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 = mat.rhoh * g.hx * g.hy * trapWeight(g, k % nx, Math.floor(k / nx));
      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));
    }
  }
  assert.ok(worst < 1e-4, `gradient mismatch ${worst}`);
  // The stress-free shape itself carries no membrane energy.
  const rest = fields(g.n);
  assert.ok(potentialEnergy(g, mat, rest, w0) < 1e-12, "w0 is stress free");
});