import { describe, expect, it } from "vitest";
import { LIBRARY_FILES, buildCompareBundle, isPlainStruct, parseCompareResult } from "../src/index";
import { maxAbs, runner, solve, testSpec } from "./helpers";

/** Run a MATLAB snippet against the library and return the struct it leaves in `out`. */
async function script(source: string) {
  const run = await runner.run({
    bundle: {
      files: [{ name: "c.m", source }, ...LIBRARY_FILES],
      entry: "c.m",
      resultVariable: "out",
    },
  });
  if (!run.ok) throw new Error(`${run.error}\n${run.output.join("")}`);
  if (!isPlainStruct(run.value)) throw new Error("unexpected shape");
  return run.value.fields as Record<string, number>;
}

describe("free edges", () => {
  it("assembles a symmetric stiffness whose null space is exactly the rigid-body motions and whose interior is the 13-point stencil", async () => {
    const out = await script(`
      nx = 19; ny = 13; dx = 0.3/(nx-1); dy = 0.2/(ny-1); nu = 0.3;
      dom = plate_domain(struct('shape', 'rectangle', 'Lx', 0.3, 'Ly', 0.2), nx, ny);
      [L, K, M, idx] = plate_matrices(dom, 'free', nu);
      x = linspace(0, 0.3, nx); y = linspace(0, 0.2, ny); [X, Y] = meshgrid(x, y);
      out.symm = norm(full(K - K'), 1) / norm(full(K), 1);
      out.null1 = max(abs(K * ones(nx*ny, 1)));
      out.nullx = max(abs(K * X(:))) * dx^4;
      out.nully = max(abs(K * Y(:))) * dx^4;
      out.lnull = max(abs(L * ones(nx*ny, 1)));
      % a curved field: the interior rows must be the biharmonic stencil, per unit mass
      W = sin(3*X) .* cos(5*Y) + X.^2 .* Y;
      Kw = reshape(K * W(:), ny, nx); Mv = reshape(full(diag(M)), ny, nx);
      Lw = plate_laplacian(plate_pad(W, 'simply-supported'), dx, dy);   % pad only feeds the outer ring, dropped below
      Bw = (Lw(3:end,2:end-1) - 2*Lw(2:end-1,2:end-1) + Lw(1:end-2,2:end-1))/dy^2 + (Lw(2:end-1,3:end) - 2*Lw(2:end-1,2:end-1) + Lw(2:end-1,1:end-2))/dx^2;
      inner = Kw(3:end-2, 3:end-2) ./ Mv(3:end-2, 3:end-2);
      ref = Bw(2:end-1, 2:end-1);
      out.interior = norm(inner(:) - ref(:)) / norm(ref(:));
      out.m_inner = Mv(5, 7); out.m_edge = Mv(1, 7); out.m_corner = Mv(1, 1);
      out.nunk = numel(idx);
    `);
    expect(out.symm).toBeLessThan(1e-12);
    expect(out.null1).toBeLessThan(1e-6);
    expect(out.nullx).toBeLessThan(1e-9);
    expect(out.nully).toBeLessThan(1e-9);
    expect(out.lnull).toBeLessThan(1e-9);
    expect(out.interior).toBeLessThan(1e-10);
    expect(out.m_inner).toBe(1);
    expect(out.m_edge).toBe(0.5);
    expect(out.m_corner).toBe(0.25);
    expect(out.nunk).toBe(19 * 13);
  });

  it("reproduces the frequencies of a completely free square plate (Leissa, nu = 0.3)", async () => {
    // omega a^2 sqrt(rho h / D) for the first three elastic modes of the free square plate,
    // Ritz values tabulated by Leissa (1969): 13.49, 19.79, 24.43. The energy method converges
    // to them from below at about first order in dx (the edge treatment), so a 17-node grid
    // should land within a couple of percent, and the three rigid modes must sit at zero.
    const out = await script(`
      n = 17; a = 1; h = a/(n-1); nu = 0.3;
      dom = plate_domain(struct('shape', 'rectangle', 'Lx', a, 'Ly', a), n, n);
      [~, K, M] = plate_matrices(dom, 'free', nu);
      Dh = diag(1 ./ sqrt(full(diag(M)))); S = Dh * full(K) * Dh;
      S = (S + S') / 2;
      ev = sort(eig(S));
      out.zero3 = max(abs(ev(1:3))) / ev(end);
      out.f1 = sqrt(ev(4)) * a^2; out.f2 = sqrt(ev(5)) * a^2; out.f3 = sqrt(ev(6)) * a^2;
      out.lam_ratio = ev(end) / (4/h^2 + 4/h^2)^2;
    `);
    expect(out.zero3).toBeLessThan(1e-9);
    expect(Math.abs(out.f1 - 13.49) / 13.49).toBeLessThan(0.02);
    // and it lies below the simply supported square's 2 pi^2 = 19.74: free edges are the softest support
    expect(out.f1).toBeLessThan(2 * Math.PI ** 2);
    expect(Math.abs(out.f2 - 19.79) / 19.79).toBeLessThan(0.03);
    expect(Math.abs(out.f3 - 24.43) / 24.43).toBeLessThan(0.03);
    // the explicit scheme's CFL bound, derived for the interior stencil, remains an upper bound
    expect(out.lam_ratio).toBeLessThan(1);
    expect(out.lam_ratio).toBeGreaterThan(0.9);
  });

  it("runs the explicit scheme on a free plate: stable below the CFL limit, no rigid-body drift, corners stress-free", async () => {
    const spec = testSpec();
    spec.bc = "free";
    spec.fd.cfl = 0.98;
    const { result } = await solve(spec, "fd");
    expect(result.info.stable).toBe(true);
    expect(result.probe.w.every(Number.isFinite)).toBe(true);
    expect(maxAbs(result.probe.w)).toBeGreaterThan(0);
    expect(result.info.notes.some((n) => /rigid-body/.test(n))).toBe(true);
    // the last frame carries no mean displacement and no mean tilt in the mass inner product
    // (trapezoidal node weights: 1 inside, 1/2 on an edge, 1/4 at a corner), which is the
    // inner product the projector uses
    const { nx, ny, nf, data } = result.frames.w;
    const k = nf - 1;
    let sum = 0,
      sx = 0,
      sy = 0,
      peak = 0,
      mass = 0;
    for (let j = 0; j < nx; j++)
      for (let i = 0; i < ny; i++) {
        const v = data[i + j * ny + k * nx * ny];
        const m = (j === 0 || j === nx - 1 ? 0.5 : 1) * (i === 0 || i === ny - 1 ? 0.5 : 1);
        mass += m;
        sum += m * v;
        sx += m * v * (j - (nx - 1) / 2);
        sy += m * v * (i - (ny - 1) / 2);
        peak = Math.max(peak, Math.abs(v));
      }
    expect(peak).toBeGreaterThan(0);
    expect(Math.abs(sum) / mass).toBeLessThan(1e-9 * peak);
    expect(Math.abs(sx) / (mass * nx)).toBeLessThan(1e-9 * peak);
    expect(Math.abs(sy) / (mass * ny)).toBeLessThan(1e-9 * peak);
    // at a free corner both curvatures and the twist vanish, so every stress there is zero
    const vm = result.frames.vm;
    for (const [i, j] of [
      [0, 0],
      [0, nx - 1],
      [ny - 1, 0],
      [ny - 1, nx - 1],
    ])
      expect(vm.data[i + j * ny + k * nx * ny]).toBe(0);
    // and the edges move, unlike a supported plate
    let edgePeak = 0;
    for (let j = 0; j < nx; j++)
      edgePeak = Math.max(edgePeak, Math.abs(data[0 + j * ny + k * nx * ny]));
    expect(edgePeak).toBeGreaterThan(0.05 * peak);
  });

  it("diverges just above the CFL limit, as the supported plate does", async () => {
    const spec = testSpec();
    spec.bc = "free";
    spec.fd.cfl = 1.04;
    const { result } = await solve(spec, "fd");
    expect(result.info.stable).toBe(false);
    expect(result.probe.w.every(Number.isFinite)).toBe(true);
  });

  it("agrees between the implicit and explicit schemes", async () => {
    const spec = testSpec();
    spec.bc = "free";
    spec.implicit = {
      nx: spec.fd.nx,
      ny: spec.fd.ny,
      duration: spec.fd.duration,
      steps_per_sample: 2,
      theta: 0.25,
    };
    const imp = await solve(spec, "implicit");
    const exp = await solve(spec, "fd");
    expect(imp.result.info.stable).toBe(true);
    expect(imp.result.info.extra.nunknowns).toBe(spec.fd.nx * spec.fd.ny);
    const cmp = await runner.run({
      bundle: buildCompareBundle(),
      inputs: { A: exp.run.ok ? exp.run.value : 0, B: imp.run.ok ? imp.run.value : 0 },
    });
    expect(cmp.ok).toBe(true);
    if (!cmp.ok) return;
    const c = parseCompareResult(cmp.value);
    expect(c.w_rel_l2).toBeLessThan(0.03);
  });

  it("is refused by the modal method with a legible message", async () => {
    const spec = testSpec();
    spec.bc = "free";
    await expect(solve(spec, "modal")).rejects.toThrow(/simply supported/);
  });
});

describe("pinned centre", () => {
  it("removes the centre node from the system and leaves exactly the two rotations about it at zero frequency", async () => {
    const out = await script(`
      n = 15; a = 1; h = a/(n-1); nu = 0.3;
      dom = plate_domain(struct('shape', 'rectangle', 'Lx', a, 'Ly', a), n, n);
      [L, K, M, idx, fixed] = plate_matrices(dom, 'pinned-center', nu);
      ic = (floor(n/2) + 1) + floor(n/2) * n;
      out.fixed = fixed; out.ic = ic;
      out.row = max(abs(K(ic, :))); out.col = max(abs(K(:, ic))); out.lrow = max(abs(L(ic, :)));
      out.mdiag = full(M(ic, ic));
      x = linspace(0, a, n); [X, Y] = meshgrid(x, x);
      out.rotx = max(abs(K * (X(:) - X(ic)))) * h^4; out.roty = max(abs(K * (Y(:) - Y(ic)))) * h^4;
      out.const = max(abs(K * ones(n*n, 1))) * h^4;
      % the pinned node's own row and column are zero, so drop it before looking at the spectrum
      keep = true(n*n, 1); keep(ic) = false;
      Dh = diag(1 ./ sqrt(full(diag(M)))); S = Dh * full(K) * Dh; S = (S + S') / 2; S = S(keep, keep);
      ev = sort(eig(S));
      out.zero2 = max(abs(ev(1:2))) / ev(end);
      out.third = ev(3) / ev(end);
      out.f1 = sqrt(ev(3)) * a^2; out.f2 = sqrt(ev(4)) * a^2;
      [~, Kf, Mf] = plate_matrices(dom, 'free', nu);
      Dhf = diag(1 ./ sqrt(full(diag(Mf)))); Sf = Dhf * full(Kf) * Dhf; Sf = (Sf + Sf') / 2;
      evf = sort(eig(Sf));
      out.free_f1 = sqrt(evf(4)) * a^2;
    `);
    expect(out.fixed).toBe(out.ic);
    expect(out.row).toBe(0);
    expect(out.col).toBe(0);
    expect(out.lrow).toBe(0);
    expect(out.mdiag).toBe(1);
    expect(out.rotx).toBeLessThan(1e-9);
    expect(out.roty).toBeLessThan(1e-9);
    expect(out.const).toBeGreaterThan(1e-3);
    expect(out.zero2).toBeLessThan(1e-9);
    // the pin removes the translation, so the third eigenvalue is an elastic mode, not a zero
    expect(out.third).toBeGreaterThan(1e-6);
    // Rayleigh's interlacing theorem: one constraint lifts each eigenvalue to at most the next one,
    // so the first elastic mode of the pinned plate (the rim moving in phase, centre held) lies
    // between the free plate's third zero and its fundamental
    expect(out.f1).toBeGreaterThan(0);
    expect(out.f1).toBeLessThan(out.free_f1);
    // the free plate's first elastic mode twists about both diagonals and is zero at the centre
    // node, so it is an eigenvector of the pinned plate as well, at the same frequency
    expect(Math.abs(out.f2 - out.free_f1) / out.free_f1).toBeLessThan(1e-8);
  });

  it("holds the centre at zero through both schemes while the edges move, and reports the support", async () => {
    const spec = testSpec();
    spec.bc = "pinned-center";
    spec.implicit = {
      nx: spec.fd.nx,
      ny: spec.fd.ny,
      duration: spec.fd.duration,
      steps_per_sample: 2,
      theta: 0.25,
    };
    for (const id of ["fd", "implicit"]) {
      const { result, run } = await solve(spec, id);
      expect(result.info.stable).toBe(true);
      expect(result.info.notes.some((n) => /pinned/.test(n))).toBe(true);
      const { nx, ny, nf, data } = result.frames.w;
      const jc = Math.floor(nx / 2);
      const ic = Math.floor(ny / 2);
      expect(result.info.extra.support_x).toBeCloseTo(result.grid.x[jc], 12);
      expect(result.info.extra.support_y).toBeCloseTo(result.grid.y[ic], 12);
      let peak = 0;
      let edgePeak = 0;
      for (let k = 0; k < nf; k++) {
        expect(data[ic + jc * ny + k * nx * ny]).toBe(0);
        for (let j = 0; j < nx; j++) {
          edgePeak = Math.max(edgePeak, Math.abs(data[0 + j * ny + k * nx * ny]));
          for (let i = 0; i < ny; i++)
            peak = Math.max(peak, Math.abs(data[i + j * ny + k * nx * ny]));
        }
      }
      expect(peak).toBeGreaterThan(0);
      expect(edgePeak).toBeGreaterThan(0.05 * peak);
      if (id === "fd") (spec as unknown as { _fd: typeof run })._fd = run;
    }
  });

  it("agrees between the implicit and explicit schemes", async () => {
    const spec = testSpec();
    spec.bc = "pinned-center";
    spec.implicit = {
      nx: spec.fd.nx,
      ny: spec.fd.ny,
      duration: spec.fd.duration,
      steps_per_sample: 2,
      theta: 0.25,
    };
    const imp = await solve(spec, "implicit");
    const exp = await solve(spec, "fd");
    const cmp = await runner.run({
      bundle: buildCompareBundle(),
      inputs: { A: exp.run.ok ? exp.run.value : 0, B: imp.run.ok ? imp.run.value : 0 },
    });
    expect(cmp.ok).toBe(true);
    if (!cmp.ok) return;
    expect(parseCompareResult(cmp.value).w_rel_l2).toBeLessThan(0.03);
  });

  it("is refused by the modal method", async () => {
    const spec = testSpec();
    spec.bc = "pinned-center";
    await expect(solve(spec, "modal")).rejects.toThrow(/simply supported/);
  });
});