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

describe("implicit finite differences", () => {
  it("assembles the same operator as the explicit ghost-node stencils, for both edge conditions", async () => {
    const check = `
      for bc = {'simply-supported', 'clamped'}
        nx = 19; ny = 13; dx = 0.3/(nx-1); dy = 0.2/(ny-1);
        dom = plate_domain(struct('shape', 'rectangle', 'Lx', 0.3, 'Ly', 0.2), nx, ny);
        [L, B] = plate_matrices(dom, bc{1});
        w = zeros(ny, nx); w(2:end-1, 2:end-1) = reshape(sin(1:(ny-2)*(nx-2)), ny-2, nx-2);
        Lw = plate_laplacian(plate_pad(w, bc{1}), dx, dy);
        Bw = zeros(ny, nx);
        Bw(2:end-1,2:end-1) = (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;
        wi = w(2:end-1,2:end-1); Bi = Bw(2:end-1,2:end-1);
        err.(strrep(bc{1}, '-', '_')) = norm(B*wi(:) - Bi(:)) / norm(Bi(:));
        sp.(strrep(bc{1}, '-', '_')) = issparse(B) && issparse(L);
      end
      out = struct('err', err, 'sp', sp);
    `;
    const run = await runner.run({
      bundle: {
        files: [{ name: "c.m", source: check }, ...LIBRARY_FILES],
        entry: "c.m",
        resultVariable: "out",
      },
    });
    expect(run.ok).toBe(true);
    if (!run.ok || !isPlainStruct(run.value)) return;
    const err = run.value.fields.err;
    const sp = run.value.fields.sp;
    if (!isPlainStruct(err) || !isPlainStruct(sp)) throw new Error("unexpected shape");
    expect(err.fields.simply_supported as number).toBeLessThan(1e-12);
    expect(err.fields.clamped as number).toBeLessThan(1e-12);
    expect(sp.fields.simply_supported).toBe(true);
    expect(sp.fields.clamped).toBe(true);
  });

  it("agrees with the explicit scheme on the same grid and with the modal solution", async () => {
    const spec = testSpec();
    spec.implicit = {
      nx: spec.fd.nx,
      ny: spec.fd.ny,
      duration: spec.fd.duration,
      steps_per_sample: 2,
      theta: 0.25,
    };
    spec.output.duration = spec.fd.duration;
    spec.output.grid = [spec.fd.nx, spec.fd.ny];
    const imp = await solve(spec, "implicit");
    const exp = await solve(spec, "fd");
    const modal = await solve(spec, "modal");
    expect(imp.result.method).toBe("implicit");
    expect(imp.result.info.stable).toBe(true);
    expect(imp.result.info.nsteps).toBeLessThan(exp.result.info.nsteps);
    const cmpE = await runner.run({
      bundle: buildCompareBundle(),
      inputs: { A: exp.run.ok ? exp.run.value : 0, B: imp.run.ok ? imp.run.value : 0 },
    });
    const cmpM = await runner.run({
      bundle: buildCompareBundle(),
      inputs: { A: modal.run.ok ? modal.run.value : 0, B: imp.run.ok ? imp.run.value : 0 },
    });
    expect(cmpE.ok && cmpM.ok).toBe(true);
    if (!cmpE.ok || !cmpM.ok) return;
    const e = parseCompareResult(cmpE.value);
    const m = parseCompareResult(cmpM.value);
    // same spatial operator, different time discretization: the two FD schemes agree closely
    expect(e.w_rel_l2).toBeLessThan(0.02);
    expect(m.w_rel_l2).toBeLessThan(0.05);
    expect(m.vm_rel_l2).toBeLessThan(0.15);
    expect(m.frames_rel_l2.length).toBeGreaterThan(0);
  });

  it("is unconditionally stable for theta = 1/4 and blows up for theta = 0 at the same step", async () => {
    const spec = testSpec();
    spec.implicit = { nx: 31, ny: 21, duration: 0.01, steps_per_sample: 0.25, theta: 0.25 };
    const coarse = await solve(spec, "implicit");
    expect(coarse.result.info.stable).toBe(true);
    expect(maxAbs(coarse.result.probe.w)).toBeLessThan(1e-3);
    spec.implicit = { nx: 31, ny: 21, duration: 0.002, steps_per_sample: 1, theta: 0 };
    const explicitLike = await solve(spec, "implicit");
    expect(explicitLike.result.info.stable).toBe(false);
    expect(explicitLike.result.info.diverged_step).toBeGreaterThan(0);
    expect(explicitLike.result.probe.w.every(Number.isFinite)).toBe(true);
  });

  it("supports clamped edges", async () => {
    const spec = testSpec();
    spec.bc = "clamped";
    spec.implicit = { nx: 31, ny: 21, duration: 0.003, steps_per_sample: 1, theta: 0.25 };
    const r = await solve(spec, "implicit");
    expect(r.result.info.stable).toBe(true);
    expect(maxAbs(r.result.probe.w)).toBeGreaterThan(0);
  });
});