/**
 * The Laplacian stencils, as external operations the .m can call: `lap2` (the
 * 5-point second-order stencil) and `lap4` (the 9-point fourth-order one).
 *
 * Everything else a model does is element-wise, so these are the only places
 * where a grid point reads its neighbours — the one thing the element-wise
 * kernel emitter cannot express, since it walks a single linear index across
 * every operand. Keeping them as named operations rather than as array slicing
 * (`p(2:end-1, :)` and friends) means the .m never has to know how the grid is
 * laid out in the buffer, and the host is free to implement the stencil as one
 * dispatch (src/mgpu/stencil.ts).
 *
 * numbl needs only their *type rule* in order to lower a call site. It gets
 * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
 * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
 * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
 * imports available, so `transfer` builds numbl `Type` objects as plain
 * literals, and the grid size is baked in by the generator below (a grid change
 * recompiles anyway).
 *
 * The `emit`/`cBody` exports exist only because the loader's contract requires
 * them; we never emit C. The actual implementation is supplied by the WGSL
 * backend.
 */

export interface GridSizes {
  /** Grid points, nx*ny. Grid fields are npts x 1 column vectors. */
  npts: number;
}

const numericType = (rows: number, cols: number): string =>
  `{ kind: "Numeric", elem: "double", isComplex: false, ` +
  `dims: [${dim(rows)}, ${dim(cols)}], shape: [${rows}, ${cols}], sign: "unknown" }`;

// numbl's tensorDouble() canonicalizes an extent of 1 to its shared DIM_ONE
// singleton; mirror that so types compare equal to host-built ones.
const dim = (n: number): string =>
  n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;

/** Source for one stencil's `.mtoc2.js`: grid field in, grid field out. */
function stencilSource(name: string, npts: number): string {
  return `
exports.name = ${JSON.stringify(name)};

exports.transfer = function (argTypes, nargout) {
  if (argTypes.length !== 1) {
    throw new Error("${name} takes exactly one argument, got " + argTypes.length);
  }
  if (nargout > 1) {
    throw new Error("${name} returns one value, but " + nargout + " were requested");
  }
  var a = argTypes[0];
  if (!a || a.kind !== "Numeric" || a.isComplex) {
    throw new Error("${name} requires a real numeric array");
  }
  var s = a.shape;
  if (!s || s.length !== 2 || s[0] !== ${npts} || s[1] !== 1) {
    throw new Error(
      "${name} works on grid fields, so its argument must be ${npts}x1, not " +
        (s ? s.join("x") : "unknown shape")
    );
  }
  return [${numericType(npts, 1)}];
};

// Never called: this project executes the IR on WebGPU and emits no C.
exports.emit = function () {
  throw new Error("${name}: no C backend (this stencil runs on WebGPU)");
};
exports.cBody = function () {
  return "";
};
`;
}

/** Workspace files that make `lap2` / `lap4` resolvable during lowering. */
export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
  return [
    { name: 'lap2.mtoc2.js', source: stencilSource('lap2', g.npts) },
    { name: 'lap4.mtoc2.js', source: stencilSource('lap4', g.npts) },
  ];
}

/** Names the WGSL backend must implement as stencil dispatches rather than
 *  element-wise kernels. */
export const EXTERNAL_OPS = new Set(['lap2', 'lap4']);