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

const fin = (a: Float64Array) => a.filter(Number.isFinite);

/** Relative L2 error over several rows at once, normalized by the norm of all reference rows together. */
function relL2Rows(a: Float64Array[], b: Float64Array[]): number {
  let num = 0;
  let den = 0;
  for (let k = 0; k < a.length; k++)
    for (let i = 0; i < a[k].length; i++) {
      num += (a[k][i] - b[k][i]) ** 2;
      den += a[k][i] ** 2;
    }
  return a.length ? Math.sqrt(num / Math.max(den, 1e-300)) : 0;
}

describe.skipIf(!handle)("explicit finite differences on WebGPU", () => {
  it.each([
    "sheet-on-thread",
    "gong-on-cords",
    "sheet-on-foam",
    "cymbal-on-stand",
    "cantilever-in-vise",
    "disk-in-hand",
  ])("reproduces the double-precision loop for %s", async (scenario) => {
    const spec = testSpec(scenario);
    spec.fd.duration = 0.01;
    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.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(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);
    // a coordinate or a support that stays at zero has no scale of its own: compare against the largest
    const eR = relL2Rows(rc.rigid.series, rg.rigid.series);
    const eS = relL2Rows(rc.supports.series, rg.supports.series);
    console.log(
      `gpu.test ${scenario}: rel L2 probe w ${eW.toExponential(2)}, frames ${eF.toExponential(2)}, probe vm ${eVm.toExponential(2)}, audio ${eA.toExponential(2)}, rigid ${eR.toExponential(2)}, supports ${eS.toExponential(2)}, ${rg.info.nsteps} steps in ${g.compute?.gpuMs?.toFixed(0)} ms on the GPU`,
    );
    expect(eW).toBeLessThan(2e-3);
    expect(eF).toBeLessThan(2e-3);
    expect(eVm).toBeLessThan(2e-2);
    expect(eA).toBeLessThan(2e-2);
    expect(eR).toBeLessThan(2e-3);
    expect(eS).toBeLessThan(2e-2);
    expect(maxAbs(rg.probe.w)).toBeGreaterThan(0);
  });

  it("keeps the swing of a released plate for many steps", async () => {
    const spec = testSpec();
    spec.excitation = { type: "release", x: 0.15, y: 0.02, displacement: 0.002, radius: 0.015 };
    spec.fd.duration = 0.05;
    const { rg, rc } = await both(spec);
    const eR = relL2(rc.rigid.series[1], rg.rigid.series[1]);
    const eW = relL2(rc.probe.w, rg.probe.w);
    console.log(
      `gpu.test release: rel L2 swing ${eR.toExponential(2)}, probe ${eW.toExponential(2)}`,
    );
    expect(eR).toBeLessThan(1e-3);
    expect(eW).toBeLessThan(2e-3);
  });

  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);
    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.t.length).toBe(rc.frames.t.length);
  });

  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(/Rec = plate_fd_loop\(S\);/);
    expect(staged).toMatch(/result = plate_fd_output\(S, Rec\);/);
    const stages = buildGpuStages(spec, FD_METHOD)!;
    expect(stages.loopFile).toBe("plate_fd_loop.m");
    expect(stages.setup).toMatch(/find\(S\.KB\)/);
    expect(stages.output).toMatch(/plate_fd_output\(S, Rec\)/);
  });
});