import { describe, expect, it } from "vitest";
import { NodeRunner } from "../src/node";
import {
  FD_METHOD,
  LIBRARY_FILES,
  buildGpuStages,
  buildRunBundle,
  parsePlateResult,
  withEdits,
  type ProblemSpec,
  type RunResult,
} from "../src/index";
import { requestGpuDevice, type GpuHandle } from "../src/gpu/fdLoop";
import { maxAbs, relL2, testSpec } from "./helpers";

/**
 * Desktop WebGPU for Node through the optional `webgpu` package (Google Dawn):
 * `npm i --no-save webgpu`. Without the package, or on a machine without a GPU
 * adapter, the suite is skipped rather than failed, so CI can run it
 * opportunistically.
 */
async function dawn(): Promise<GpuHandle | null> {
  try {
    const specifier = "webgpu";
    const mod = (await import(specifier)) as {
      create: (flags: string[]) => GPU;
      globals: Record<string, unknown>;
    };
    Object.assign(globalThis, mod.globals);
    Object.defineProperty(globalThis, "navigator", {
      value: { gpu: mod.create([]) },
      configurable: true,
      writable: true,
    });
    return await requestGpuDevice();
  } catch (e) {
    console.log(
      `gpu.test: no WebGPU device under Node (${e instanceof Error ? e.message : e}); skipping`,
    );
    return null;
  }
}

const handle = await dawn();
const cpu = new NodeRunner();
const gpu = new NodeRunner({ gpu: () => Promise.resolve(handle!) });

async function both(spec: ProblemSpec, edits: Record<string, string> = {}) {
  const stages = buildGpuStages(spec, FD_METHOD)!;
  const g = await gpu.run({
    bundle: withEdits(buildRunBundle(spec, FD_METHOD, { gpu: true }), edits),
    gpu: stages,
  });
  const c = await cpu.run({ bundle: withEdits(buildRunBundle(spec, FD_METHOD), edits) });
  if (!g.ok) throw new Error(`gpu run failed: ${g.error}\n${g.output.join("")}`);
  if (!c.ok) throw new Error(`cpu run failed: ${c.error}\n${c.output.join("")}`);
  return { g, c, rg: parsePlateResult(g.value), rc: parsePlateResult(c.value) };
}

describe.skipIf(!handle)("explicit finite differences on WebGPU", () => {
  it.each(["simply-supported", "clamped", "free", "pinned-center"] as const)(
    "reproduces the double-precision loop for %s edges",
    async (bc) => {
      const spec = testSpec();
      spec.bc = bc;
      const { g, rg, rc } = await both(spec);
      expect(g.compute?.loop).toBe("webgpu");
      expect(rg.info.extra.loop).toBe("webgpu");
      expect(rg.info.stable).toBe(true);
      expect(rg.info.notes.some((n) => /GPU/.test(n))).toBe(true);
      expect(rg.t.length).toBe(rc.t.length);
      expect(rg.frames.t.length).toBe(rc.frames.t.length);
      const eW = relL2(rc.probe.w, rg.probe.w);
      const eF = relL2(rc.frames.w.data, rg.frames.w.data);
      const eVm = relL2(rc.probe.vm, rg.probe.vm);
      const eA = relL2(rc.audio.p, rg.audio.p);
      const eEnv = relL2(rc.envelope.vm.data, rg.envelope.vm.data);
      console.log(
        `gpu.test ${bc}: rel L2 probe w ${eW.toExponential(2)}, frames ${eF.toExponential(2)}, probe vm ${eVm.toExponential(2)}, audio ${eA.toExponential(2)}, envelope ${eEnv.toExponential(2)}, ${rg.info.nsteps} steps in ${g.compute?.gpuMs?.toFixed(0)} ms on the GPU`,
      );
      expect(eW).toBeLessThan(1e-3);
      expect(eF).toBeLessThan(1e-3);
      expect(eVm).toBeLessThan(1e-2);
      expect(eA).toBeLessThan(1e-2);
      expect(eEnv).toBeLessThan(1e-2);
      expect(maxAbs(rg.probe.w)).toBeGreaterThan(0);
    },
  );

  it.each(["simply-supported", "clamped", "free", "pinned-center"] as const)(
    "reproduces the double-precision loop on an ellipse with %s edges",
    async (bc) => {
      const spec = testSpec();
      spec.plate.shape = "ellipse";
      spec.bc = bc;
      const { g, rg, rc } = await both(spec);
      expect(g.compute?.loop).toBe("webgpu");
      expect(rg.info.stable).toBe(true);
      expect(rg.frames.w.data.filter((v) => Number.isNaN(v)).length).toBe(
        rc.frames.w.data.filter((v) => Number.isNaN(v)).length,
      );
      const fin = (a: Float64Array) => a.filter(Number.isFinite);
      const eW = relL2(rc.probe.w, rg.probe.w);
      const eF = relL2(fin(rc.frames.w.data), fin(rg.frames.w.data));
      const eVm = relL2(rc.probe.vm, rg.probe.vm);
      const eA = relL2(rc.audio.p, rg.audio.p);
      console.log(
        `gpu.test ellipse ${bc}: rel L2 probe w ${eW.toExponential(2)}, frames ${eF.toExponential(2)}, probe vm ${eVm.toExponential(2)}, audio ${eA.toExponential(2)}`,
      );
      expect(eW).toBeLessThan(1e-3);
      expect(eF).toBeLessThan(1e-3);
      expect(eVm).toBeLessThan(1e-2);
      expect(eA).toBeLessThan(1e-2);
      expect(maxAbs(rg.probe.w)).toBeGreaterThan(0);
    },
  );

  it("stops at the blow-up where the double-precision loop does", async () => {
    const spec = testSpec();
    spec.fd.cfl = 1.03;
    const { rg, rc } = await both(spec);
    expect(rg.info.stable).toBe(false);
    expect(rc.info.stable).toBe(false);
    expect(rg.info.diverged_step).toBeGreaterThan(0);
    // the unstable mode grows from rounding noise once the strike is over, so the two runs need not agree to the step
    expect(Math.abs(rg.info.diverged_step - rc.info.diverged_step)).toBeLessThanOrEqual(
      Math.max(10, 0.1 * rc.info.nsteps),
    );
    expect(rg.probe.w.every(Number.isFinite)).toBe(true);
    expect(rg.frames.w.data.every(Number.isFinite)).toBe(true);
    expect(rg.frames.t.length).toBe(rc.frames.t.length);
    expect(rg.info.notes.some((n) => /diverged/.test(n))).toBe(true);
  });

  it("falls back to numbl when the loop file is edited, and says so", async () => {
    const spec = testSpec();
    const loop = LIBRARY_FILES.find((f) => f.name === "plate_fd_loop.m")!.source;
    const { g, rg, rc } = await both(spec, {
      "plate_fd_loop.m": loop + "\n% edited in the code view\n",
    });
    expect(g.compute?.loop).toBe("numbl");
    expect(g.compute?.note).toMatch(/edited/);
    expect(rg.info.extra.loop).toBe("numbl");
    expect(relL2(rc.probe.w, rg.probe.w)).toBeLessThan(1e-12);
  });
});

describe("explicit finite differences without a GPU", () => {
  it("runs the whole bundle in numbl when the GPU is asked for but absent", async () => {
    const spec = testSpec();
    const r: RunResult = await cpu.run({
      bundle: buildRunBundle(spec, FD_METHOD, { gpu: true }),
      gpu: buildGpuStages(spec, FD_METHOD)!,
    });
    expect(r.ok).toBe(true);
    if (!r.ok) return;
    expect(r.compute?.loop).toBe("numbl");
    expect(r.compute?.note).toMatch(/not available/);
    const res = parsePlateResult(r.value);
    expect(res.info.extra.loop).toBe("numbl");
    expect(res.info.stable).toBe(true);
  });

  it("writes the staged script and the plain script for the same problem", () => {
    const spec = testSpec();
    const plain = buildRunBundle(spec, FD_METHOD).files[0].source;
    const staged = buildRunBundle(spec, FD_METHOD, { gpu: true }).files[0].source;
    expect(plain).toMatch(/result = plate_fd\(params\);/);
    expect(staged).toMatch(/S = plate_fd_setup\(params\);/);
    expect(staged).toMatch(/R = plate_fd_loop\(S\);/);
    expect(staged).toMatch(/result = plate_fd_output\(S, R\);/);
    const stages = buildGpuStages(spec, FD_METHOD)!;
    expect(stages.loopFile).toBe("plate_fd_loop.m");
    expect(stages.setup).toMatch(/find\(S\.Km\)/);
    expect(stages.output).toMatch(/plate_fd_output\(S, R\)/);
  });
});