// The eigen solver: banded factorization and inertia, Lanczos slices against
// a dense Jacobi solve on small grids (with degenerate pairs), and the
// free-plate frequencies against the reference values.
import { test } from "node:test";
import assert from "node:assert/strict";
import { makeGrid, withMask, type Plate } from "../src/params.ts";
import { uniformPlate } from "../src/model/reference.ts";
import { assemble, stiffnessEntry, type LinearSystem } from "../src/model/assemble.ts";
import { factorize, solve, inertia } from "../src/modal/banded.ts";
import { tridiagonalEigen, sliceModes } from "../src/modal/lanczos.ts";
import { computeModes, localRunner } from "../src/modal/modes.ts";

const steel = { E: 200e9, nu: 0.3, rho: 7850, eta: 0, thermal: null };

/** All eigenvalues of M^{-1/2} K M^{-1/2} by cyclic Jacobi (small systems only). */
function denseEigenvalues(sys: LinearSystem): number[] {
  const n = sys.n;
  const A: number[][] = [];
  for (let i = 0; i < n; i++) {
    A.push([]);
    for (let j = 0; j < n; j++) A[i].push(stiffnessEntry(sys, i, j) / Math.sqrt(sys.mass[i] * sys.mass[j]));
  }
  for (let sweep = 0; sweep < 60; sweep++) {
    let off = 0;
    for (let i = 0; i < n; i++) for (let j = i + 1; j < n; j++) off += A[i][j] ** 2;
    if (off < 1e-24 * n * n) break;
    for (let p = 0; p < n; p++) {
      for (let q = p + 1; q < n; q++) {
        if (Math.abs(A[p][q]) < 1e-300) continue;
        const theta = (A[q][q] - A[p][p]) / (2 * A[p][q]);
        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][p], akq = A[k][q];
          A[k][p] = c * akp - s * akq; A[k][q] = s * akp + c * akq;
        }
        for (let k = 0; k < n; k++) {
          const apk = A[p][k], aqk = A[q][k];
          A[p][k] = c * apk - s * aqk; A[q][k] = s * apk + c * aqk;
        }
      }
    }
  }
  return Array.from({ length: n }, (_, i) => A[i][i]).sort((a, b) => a - b);
}

test("the banded LDL^T solves and counts eigenvalues", () => {
  const sys = assemble(uniformPlate(makeGrid(0.1, 0.08, 0.01), steel, 1e-3));
  const shift = -((2 * Math.PI * 50) ** 2);
  const f = factorize(sys.band, sys.mass, sys.n, sys.bw, shift);
  assert.equal(f.negative, 0, "K + 50 Hz shift is positive definite");
  const b = new Float64Array(sys.n).map((_, i) => Math.sin(i));
  const x = new Float64Array(sys.n);
  solve(f, b, x);
  // Check (K - shift M) x = b with the banded matvec.
  const { stiffnessMatvec } = require_assemble();
  const y = new Float64Array(sys.n);
  stiffnessMatvec(sys, x, y);
  let err = 0, nrm = 0;
  for (let i = 0; i < sys.n; i++) { const r = y[i] - shift * sys.mass[i] * x[i] - b[i]; err += r * r; nrm += b[i] * b[i]; }
  assert.ok(Math.sqrt(err / nrm) < 1e-9);
  const dense = denseEigenvalues(sys);
  for (const f of [30, 100, 300, 1000]) {
    const lam = (2 * Math.PI * f) ** 2;
    const expected = dense.filter((v) => v < lam).length;
    assert.equal(inertia(sys.band, sys.mass, sys.n, sys.bw, lam), expected, `count below ${f} Hz`);
  }
});

import * as assembleModule from "../src/model/assemble.ts";
function require_assemble(): typeof assembleModule { return assembleModule; }

test("the tridiagonal eigensolver agrees with a dense solve", () => {
  const d = [4, 3, 2, 5, 1, 3, 3], e = [1, 0.5, 2, 0.1, 1, 1];
  const { values, vectors } = tridiagonalEigen(d, e);
  const n = d.length;
  for (let k = 0; k < n; k++) {
    // T v = lambda v
    for (let i = 0; i < n; i++) {
      const tv = d[i] * vectors[k][i] + (i > 0 ? e[i - 1] * vectors[k][i - 1] : 0) + (i < n - 1 ? e[i] * vectors[k][i + 1] : 0);
      assert.ok(Math.abs(tv - values[k] * vectors[k][i]) < 1e-10);
    }
  }
  for (let k = 1; k < n; k++) assert.ok(values[k] >= values[k - 1]);
});

test("all modes of a small free square and a small disc match a dense solve, degenerate pairs included", async () => {
  const cases: { name: string; pl: Plate }[] = [
    { name: "square", pl: uniformPlate(makeGrid(0.1, 0.1, 0.1 / 9), steel, 1e-3) },
    { name: "disc", pl: uniformPlate(withMask(makeGrid(0.12, 0.12, 0.01), (x, y) => x * x + y * y <= 0.06 ** 2), steel, 1e-3) },
  ];
  for (const c of cases) {
    const sys = assemble(c.pl);
    const dense = denseEigenvalues(sys).filter((v) => v > (2 * Math.PI) ** 2);
    const fMax = Math.sqrt(dense[dense.length - 1]) / (2 * Math.PI) * 1.01;
    const ms = await computeModes(sys, fMax, localRunner(sys), { sliceSize: 12 });
    assert.equal(ms.missing, 0, `${c.name}: missing modes`);
    assert.equal(ms.count, dense.length, `${c.name}: ${ms.count} modes vs ${dense.length} dense`);
    for (let m = 0; m < ms.count; m++) {
      const lam = (2 * Math.PI * ms.freqs[m]) ** 2;
      assert.ok(Math.abs(lam / dense[m] - 1) < 1e-7, `${c.name}: mode ${m} ${lam} vs ${dense[m]}`);
    }
    // Mass-normalized shapes: phi' M phi = 1, and K phi = lambda M phi.
    const { stiffnessMatvec } = require_assemble();
    const y = new Float64Array(sys.n);
    for (const m of [0, 1, Math.floor(ms.count / 2), ms.count - 1]) {
      const phi = Float64Array.from(ms.shapes.subarray(m * sys.n, (m + 1) * sys.n));
      let mn = 0;
      for (let i = 0; i < sys.n; i++) mn += sys.mass[i] * phi[i] * phi[i];
      assert.ok(Math.abs(mn - 1) < 1e-5, `${c.name}: mass norm ${mn}`);
      stiffnessMatvec(sys, phi, y);
      const lam = (2 * Math.PI * ms.freqs[m]) ** 2;
      let r = 0, s = 0;
      for (let i = 0; i < sys.n; i++) { r += (y[i] - lam * sys.mass[i] * phi[i]) ** 2; s += (lam * sys.mass[i] * phi[i]) ** 2; }
      // The stored shapes are float32, which the stiffness amplifies by omega_max^2 / omega^2.
      assert.ok(Math.sqrt(r / s) < 1e-3, `${c.name}: residual ${Math.sqrt(r / s)}`);
    }
    // The float64 shapes of one slice satisfy the eigen equation to rounding.
    const lo = (2 * Math.PI) ** 2, hi = dense[Math.min(9, dense.length - 1)] * 1.0001;
    const slice = sliceModes(sys.band, sys.mass, sys.n, sys.bw, lo, hi, inertia(sys.band, sys.mass, sys.n, sys.bw, lo), inertia(sys.band, sys.mass, sys.n, sys.bw, hi));
    assert.equal(slice.lambda.length, slice.expected);
    for (let m = 0; m < slice.lambda.length; m++) {
      stiffnessMatvec(sys, slice.shapes[m], y);
      let r = 0, s = 0;
      for (let i = 0; i < sys.n; i++) { r += (y[i] - slice.lambda[m] * sys.mass[i] * slice.shapes[m][i]) ** 2; s += (slice.lambda[m] * sys.mass[i] * slice.shapes[m][i]) ** 2; }
      assert.ok(Math.sqrt(r / s) < 1e-6, `${c.name}: float64 residual ${Math.sqrt(r / s)}`);
    }
    // The disc has degenerate pairs: the first two frequencies agree closely.
    if (c.name === "disc") assert.ok(Math.abs(ms.freqs[1] / ms.freqs[0] - 1) < 0.02, "degenerate pair");
  }
});

test("free square plate frequencies converge to Gorman's values", async () => {
  // lambda = omega a^2 sqrt(rho h / D) for the completely free square (nu = 0.3).
  const gorman = [13.468, 19.596, 24.271, 34.801, 34.801, 61.093, 61.093];
  const a = 0.3, h = 1e-3;
  const D = (steel.E * h ** 3) / (12 * (1 - steel.nu ** 2));
  const scale = a * a * Math.sqrt((steel.rho * h) / D);
  const errors: number[] = [];
  for (const cells of [10, 20]) {
    const sys = assemble(uniformPlate(makeGrid(a, a, a / cells), steel, h));
    const ms = await computeModes(sys, (80 / scale) / (2 * Math.PI), localRunner(sys), { sliceSize: 20 });
    let worst = 0;
    for (let m = 0; m < 7; m++) worst = Math.max(worst, Math.abs((2 * Math.PI * ms.freqs[m] * scale) / gorman[m] - 1));
    errors.push(worst);
  }
  assert.ok(errors[0] < 0.1, `10 cells: ${errors[0]}`);
  assert.ok(errors[1] < 0.025, `20 cells: ${errors[1]}`);
  assert.ok(errors[1] < errors[0]);
});

test("free circular plate frequencies are close to Leissa's values", async () => {
  // lambda^2 = omega a^2 sqrt(rho h / D), completely free disc, nu = 0.3: (2,0) 5.253, (0,1) 9.084, (3,0) 12.23, (1,1) 20.52, (4,0) 21.6.
  const leissa = [5.253, 5.253, 9.084, 12.23, 12.23, 20.52, 20.52, 21.6, 21.6];
  const a = 0.15, h = 1e-3;
  const D = (steel.E * h ** 3) / (12 * (1 - steel.nu ** 2));
  const scale = a * a * Math.sqrt((steel.rho * h) / D);
  const g = withMask(makeGrid(2 * a, 2 * a, (2 * a) / 40), (x, y) => x * x + y * y <= a * a);
  const sys = assemble(uniformPlate(g, steel, h));
  const ms = await computeModes(sys, (25 / scale) / (2 * Math.PI), localRunner(sys), { sliceSize: 20 });
  for (let m = 0; m < leissa.length; m++) {
    const lam = 2 * Math.PI * ms.freqs[m] * scale;
    assert.ok(Math.abs(lam / leissa[m] - 1) < 0.03, `mode ${m}: ${lam} vs ${leissa[m]}`);
  }
});