/**
 * The correctness suite, shared by both GPU stacks: `npm run test:node`
 * drives it through desktop Dawn, `npm run test:gpu` through headless
 * Chrome's own WebGPU (SwiftShader when there is no hardware).
 *
 * Each case compiles and runs a MATLAB script and checks the values it
 * prints (via fprintf) against references computed here in f64. Tolerances
 * are f32-scale: the GPU computes in single precision.
 */
import { runScript } from '../src/mgpu/session.ts';
import { formatFailure } from '../src/mgpu/errors.ts';

export interface Case {
  name: string;
  source: string;
  /** Expected values for each %-printed number, with relative tolerance. */
  expect?: { value: number; rel?: number; abs?: number }[];
  /** Substrings that must appear in the output. */
  contains?: string[];
  /** Expected compile failure instead of a run. */
  fails?: RegExp;
}

const N = (v: number, rel = 2e-5): { value: number; rel: number } => ({ value: v, rel });

/** f64 reference for the fused-chain case. */
function refFused(): number {
  const n = 10000;
  let acc = 0;
  for (let i = 0; i < n; i++) {
    const x = i / (n - 1);
    acc += 3 * x * x - 2 * x + Math.sin(2 * Math.PI * x) / (1 + x * x);
  }
  return acc / n;
}

export const cases: Case[] = [
  {
    name: 'fused elementwise chain',
    source: `
n = 10000;
x = linspace(0, 1, n);
y = 3*x.^2 - 2*x + sin(2*pi*x)./(1 + x.^2);
fprintf('%.6f\\n', sum(y(:)) / n);
`,
    // 3e-4 relative: SwiftShader's sin/exp are a touch less accurate than
    // hardware drivers'.
    expect: [N(refFused(), 3e-4)],
  },
  {
    name: 'gemm small vs f64 reference',
    source: `
n = 64;
A = zeros(n, n) + 1;
B = zeros(n, n) + 2;
C = A * B;
fprintf('%.1f %.1f\\n', C(:)'*C(:)/(n*n), sum(C(:)));
`,
    expect: [N(128 * 128), N(128 * 64 * 64)],
  },
  {
    name: 'reductions full and by columns',
    source: `
m = 300; n = 5;
A = zeros(m, n) + 3;
s = sum(A);
fprintf('%.1f %.1f %.1f %.1f\\n', s(:)'*s(:)/n, mean(A(:)), max(A(:)), prod(zeros(3,1)+2));
`,
    expect: [N(900 * 900), N(3), N(3), N(8)],
  },
  {
    name: 'indexing is a clear compile error',
    source: `
n = 1000;
x = zeros(n, 1);
for k = 1:50
  x = x + k;
end
fprintf('%.1f\\n', x(1 + 0*x(:)'*x(:)));
`,
    fails: /index|slic|supported/i,
  },
  {
    name: 'loop replay accumulates',
    source: `
n = 1000;
x = zeros(n, 1);
for k = 1:50
  x = x + k;
end
fprintf('%.1f\\n', mean(x));
`,
    expect: [N(1275)],
  },
  {
    name: 'rand statistics and loop-fresh draws',
    source: `
n = 200000;
a = rand(n, 1);
s = zeros(1, 1);
for k = 1:3
  s = s + mean(rand(n, 1));
end
fprintf('%.3f %.3f %.3f %.3f\\n', mean(a), mean(a.^2), s/3, mean(randn(n,1)));
`,
    expect: [
      { value: 0.5, abs: 0.01 },
      { value: 1 / 3, abs: 0.01 },
      { value: 0.5, abs: 0.01 },
      { value: 0, abs: 0.02 },
    ],
  },
  {
    name: 'comparisons, logicals, masks (monte carlo pi)',
    source: `
n = 400000;
x = 2*rand(n, 1) - 1;
y = 2*rand(n, 1) - 1;
inside = (x.^2 + y.^2) <= 1;
fprintf('%.3f\\n', 4*mean(inside));
`,
    expect: [{ value: Math.PI, abs: 0.03 }],
  },
  {
    name: 'transpose and matrix identities',
    source: `
m = 33; n = 17;
A = rand(m, n);
B = A';
d1 = sum(A(:).^2);
d2 = sum(B(:).^2);
E = eye(m);
C = E * A;
d3 = sum(abs(C(:) - A(:)));
fprintf('%.6f %.6f\\n', d1 - d2, d3);
`,
    expect: [{ value: 0, abs: 1e-3 }, { value: 0, abs: 1e-4 }],
  },
  {
    name: 'gemm matrix-vector agrees with itself',
    source: `
n = 48;
A = rand(n, n);
v = rand(n, 1);
w = A * v;
d = sum(w) - sum(A * v);
fprintf('%.6f\\n', d);
`,
    expect: [{ value: 0, abs: 1e-4 }],
  },
  {
    name: 'tic toc segments and echo',
    source: `
n = 100000;
tic;
x = rand(n, 1);
s = mean(x);
t = toc;
tic
y = x + 1;
toc
z = 3.5
`,
    contains: ['Elapsed time is', 'z =', '3.5'],
  },
  {
    name: 'dot and norm',
    source: `
n = 5000;
a = zeros(n,1) + 2;
b = zeros(n,1) + 3;
fprintf('%.1f %.4f\\n', dot(a, b), norm(a) / sqrt(n));
`,
    expect: [N(30000), N(2)],
  },
  {
    name: 'min/max elementwise two-arg',
    source: `
n = 1000;
x = linspace(-1, 1, n);
y = max(x, 0) + min(x, 0);
fprintf('%.6f\\n', sum(abs(y - x)));
`,
    expect: [{ value: 0, abs: 1e-4 }],
  },
  {
    name: 'mod and integer powers',
    source: `
x = linspace(0, 10, 101);
y = mod(x, 3);
fprintf('%.4f %.4f\\n', max(y(:)), sum((0:4).^2));
`,
    expect: [{ value: 2.9, abs: 0.001 }, N(30)],
  },
  {
    name: 'literal row vector uploads',
    source: `
v = [1 2 3 4 5];
fprintf('%.1f\\n', sum(v));
`,
    expect: [N(15)],
  },
  {
    name: 'in-place update (aliased kernel) is correct',
    source: `
n = 100;
u = zeros(n, 1) + 1;
u = u + u.^2;
u = u * 2;
fprintf('%.1f\\n', mean(u));
`,
    expect: [N(4)],
  },
  {
    name: 'while is a clear compile error',
    source: `
x = 1;
while x < 10
  x = x + 1;
end
`,
    fails: /while/i,
  },
  {
    name: 'variable-size rand is a clear compile error',
    source: `
n = rand() * 100;
A = rand(n, 1);
`,
    fails: /compile time/i,
  },
];

/** The %-formatted numbers a run printed (echo/timing lines stripped). */
function printedNumbers(output: string): number[] {
  const out: number[] = [];
  const kept = output
    .split('\n')
    .filter((l) => !/Elapsed time|=/.test(l))
    .join('\n');
  for (const m of kept.matchAll(/-?\d+(?:\.\d+)?(?:e[+-]?\d+)?/gi)) {
    out.push(Number(m[0]));
  }
  return out;
}

const indent = (s: string): string => s.replace(/^/gm, '    ');

/** Run every case; log a line per case; return the failure count. */
export async function runCases(
  device: GPUDevice,
  log: (line: string) => void,
): Promise<number> {
  let failures = 0;
  for (const c of cases) {
    try {
      const run = await runScript(device, c.source);
      if (c.fails) {
        log(`FAIL ${c.name}: expected a compile error, but it ran`);
        failures++;
        continue;
      }
      if (run.result.error) {
        log(`FAIL ${c.name}: runtime error: ${run.result.error}`);
        failures++;
        continue;
      }
      const nums = printedNumbers(run.result.output);
      let ok = true;
      (c.expect ?? []).forEach((e, i) => {
        const got = nums[i];
        const tol = e.abs ?? Math.abs(e.value) * (e.rel ?? 1e-5) + 1e-12;
        if (got === undefined || Math.abs(got - e.value) > tol) {
          log(`FAIL ${c.name}: printed[${i}] = ${got}, want ${e.value} ±${tol}`);
          ok = false;
        }
      });
      for (const s of c.contains ?? []) {
        if (!run.result.output.includes(s)) {
          log(`FAIL ${c.name}: output lacks ${JSON.stringify(s)}`);
          ok = false;
        }
      }
      if (!ok) {
        log(`  output was:\n${indent(run.result.output)}`);
        log(`  plan:\n${indent(run.planDescription.join('\n'))}`);
        failures++;
      } else {
        log(`ok   ${c.name}`);
      }
    } catch (e) {
      if (c.fails) {
        const msg = formatFailure(e, c.source);
        if (c.fails.test(msg)) {
          log(`ok   ${c.name} (declined: ${msg.split('\n')[0].slice(0, 90)})`);
        } else {
          log(`FAIL ${c.name}: wrong error: ${msg}`);
          failures++;
        }
      } else {
        log(`FAIL ${c.name}: ${formatFailure(e, c.source)}`);
        failures++;
      }
    }
  }
  return failures;
}