/**
 * Lowered script -> a replayable sequence of GPU + host operations.
 *
 * Everything expensive happens once, here: pipeline compilation, buffer
 * allocation, bind-group construction. Because numbl fixes every type and
 * shape at lowering time, the op sequence is fully static; executing it is
 * pure command recording plus the host ops (tic/toc/printing) the script asked
 * for, which are the only synchronization points.
 *
 * Statement routing:
 *  - elementwise trees          -> one fused kernel (buildKernel)
 *  - `A * B` (both tensors)     -> tiled GEMM; non-variable operands are
 *                                  materialized into scratch first
 *  - matrix `A'`                -> tiled transpose kernel
 *  - vector `x'`, `X(:)`,       -> a view of the same buffer when the source
 *    `reshape`                     is never reassigned, else a plain copy
 *  - sum/mean/prod/max/min/     -> reduction kernels over a fused loader
 *    norm/dot
 *  - tic/toc/disp/fprintf and   -> host ops: the executor flushes GPU work,
 *    unsuppressed echoes           then times/reads/prints
 *  - `for` with exact bounds    -> body planned ONCE; the loop variable lives
 *                                  in a dynamic-offset uniform (one 256-byte
 *                                  slot per iteration) and the executor
 *                                  re-encodes the body per iteration
 */
import type {
  Assign,
  Call,
  ExprStmt,
  For,
  IRExpr,
  IRStmt,
  Span,
} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import {
  isMultiElement,
  type NumericType,
  type Type,
} from 'numbl-src/numbl-core/jit/lowering/types.ts';
import type { CompiledScript } from './compile.ts';
import { UnsupportedOnGpu } from './errors.ts';
import {
  buildKernel,
  bindingDecls,
  dispatchFor,
  emitLoader,
  exactValue,
  fullColonBase,
  numel,
  REDUCTIONS,
  type KernelInputs,
} from './wgsl.ts';
import {
  gemmDispatch,
  gemmKernel,
  reduceColumns,
  reduceFullPass1,
  reduceFullPass2,
  reducePartials,
  transposeDispatch,
  transposeKernel,
  type Combine,
  type Epilogue,
  type MapKind,
} from './kernels.ts';

const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
const isVectorish = (t: Type): boolean =>
  isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;

/** Iterations a single `for` may replay. Each costs one 256-byte uniform slot
 *  and a re-encode of the body's dispatches. */
const MAX_TRIPS = 65536;
/** Total dispatches one run may encode, across all loops. */
const MAX_DISPATCHES = 2_000_000;
/** Loop-variable uniform slot stride (minUniformBufferOffsetAlignment). */
const LV_STRIDE = 256;

export interface Slot {
  buffer: GPUBuffer;
  count: number;
}

/** What a printed/displayed value reads from. */
export type ValueRef =
  | { kind: 'literal'; value: number }
  | { kind: 'buffer'; slot: Slot; count: number }
  | { kind: 'host'; cName: string };

/** One piece of an fprintf: fixed text or a formatted value. */
export type EmitPart =
  | { kind: 'text'; text: string }
  | { kind: 'value'; ref: ValueRef; spec: string };

export type Op =
  | {
      kind: 'kernel';
      pipeline: GPUComputePipeline;
      bindGroup: GPUBindGroup;
      dispatch: [number, number];
      /** cVars whose dynamic offsets must be passed, in binding order. */
      loops: string[];
      label: string;
      copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
    }
  | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string }
  | { kind: 'write'; slot: Slot; data: Float32Array; label: string }
  | { kind: 'loop'; cVar: string; trips: number; uniform: GPUBuffer; body: Op[]; label: string }
  | { kind: 'tic'; assignTo?: { cName: string; slot: Slot } }
  | {
      kind: 'toc';
      print: boolean;
      sinceCName?: string;
      assignTo?: { cName: string; slot: Slot };
      /** 1-based tic..toc pair index, for the timing table. */
      seq: number;
    }
  | { kind: 'emit'; parts: EmitPart[] }
  | {
      kind: 'display';
      /** Variable name, or null for disp() (which prints no name). */
      label: string | null;
      ref: ValueRef;
      shape: number[];
    };

/** One generated compute shader, for the "show me the kernels" pane. */
export interface KernelSource {
  /** The label of the statement that first created it. */
  label: string;
  code: string;
}

export interface ScriptPlan {
  ops: Op[];
  /** Human-readable op sequence — what the script actually compiled to. */
  describe(): string[];
  /** Every distinct WGSL kernel, in creation order (reused kernels appear
   *  once, under their first label). */
  kernels: KernelSource[];
  destroy(): void;
}

interface PlannerVar {
  slot: Slot;
  shape: number[];
}

export async function planScript(
  device: GPUDevice,
  compiled: CompiledScript,
): Promise<ScriptPlan> {
  const owned: GPUBuffer[] = [];
  /** cName -> buffer-backed variable (tensors and runtime scalars). */
  const vars = new Map<string, PlannerVar>();
  /** cName -> cName it is a view of. */
  const aliases = new Map<string, string>();
  /** cName -> char value (format strings). */
  const chars = new Map<string, string>();
  /** cNames whose value the executor knows on the host (tic/toc results). */
  const hostScalars = new Set<string>();
  /** Loop-variable uniform buffers, by cVar. */
  const loopUniforms = new Map<string, GPUBuffer>();
  const pipelines = new Map<string, GPUComputePipeline>();
  const kernelSources: KernelSource[] = [];
  const describeLines: string[] = [];

  let seedCounter = 1;
  let tempCounter = 0;
  let tocCounter = 0;
  let dispatchBudget = MAX_DISPATCHES;

  // How many times each cName is assigned, anywhere. Decides when a variable
  // may be a view of another's buffer, and when an exact scalar still needs a
  // real buffer (a later assignment makes its uses flow-dependent).
  const assignCounts = new Map<string, number>();
  {
    const walkCounts = (stmts: IRStmt[]): void => {
      for (const s of stmts) {
        if (s.kind === 'Assign') {
          assignCounts.set(s.cName, (assignCounts.get(s.cName) ?? 0) + 1);
        } else if (s.kind === 'For') {
          walkCounts(s.body);
        }
      }
    };
    walkCounts(compiled.stmts);
  }

  const resolve = (cName: string): string => {
    let c = cName;
    while (aliases.has(c)) c = aliases.get(c)!;
    return c;
  };

  const maxBytes = device.limits.maxStorageBufferBindingSize;

  const makeSlot = (label: string, count: number): Slot => {
    const bytes = Math.max(4, 4 * count);
    if (bytes > maxBytes) {
      throw new UnsupportedOnGpu(
        `'${label}' needs ${(bytes / 1e6).toFixed(0)} MB, over this device's ` +
          `storage-buffer limit of ${(maxBytes / 1e6).toFixed(0)} MB`,
      );
    }
    const buffer = device.createBuffer({
      label: `mgpu-${label}`,
      size: bytes,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
    });
    owned.push(buffer);
    return { buffer, count };
  };

  /** The variable's slot, allocating on first assignment. */
  const slotFor = (cName: string, name: string, ty: NumericType, span: Span): Slot => {
    const key = resolve(cName);
    const count = numel(ty);
    const existing = vars.get(key);
    if (existing) {
      if (existing.slot.count !== count) {
        throw new UnsupportedOnGpu(
          `'${name}' changes size between assignments (${existing.slot.count} ` +
            `-> ${count} elements); the sandbox fixes each variable's storage once`,
          span,
        );
      }
      existing.shape = ty.shape ?? [count, 1];
      return existing.slot;
    }
    const slot = makeSlot(name, count);
    vars.set(key, { slot, shape: ty.shape ?? [count, 1] });
    return slot;
  };

  const readSlot = (cName: string, name: string, span: Span): Slot => {
    const v = vars.get(resolve(cName));
    if (!v) {
      throw new UnsupportedOnGpu(`'${name}' is read before it has a value`, span);
    }
    return v.slot;
  };

  async function pipeline(code: string, label: string, layout: GPUBindGroupLayout): Promise<GPUComputePipeline> {
    const hit = pipelines.get(code);
    if (hit) return hit;
    device.pushErrorScope('validation');
    const module = device.createShaderModule({ code, label });
    const info = await module.getCompilationInfo();
    const errors = info.messages.filter((m) => m.type === 'error');
    if (errors.length) {
      throw new UnsupportedOnGpu(
        `generated WGSL failed to compile for '${label}':\n` +
          errors.map((e) => `  ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
          `\n--- shader ---\n${code}`,
      );
    }
    const p = await device.createComputePipelineAsync({
      layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
      compute: { module, entryPoint: 'main' },
      label,
    });
    const err = await device.popErrorScope();
    if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
    pipelines.set(code, p);
    kernelSources.push({ label, code });
    return p;
  }

  /** Bind group layout: out at 0, `inputs` read-only buffers, then `loops`
   *  dynamic-offset uniforms. Explicit so unused bindings still match. */
  function kernelLayout(inputs: number, loops: number): GPUBindGroupLayout {
    const entries: GPUBindGroupLayoutEntry[] = [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
    ];
    for (let i = 0; i < inputs; i++) {
      entries.push({
        binding: i + 1,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: 'read-only-storage' },
      });
    }
    for (let k = 0; k < loops; k++) {
      entries.push({
        binding: inputs + 1 + k,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: 'uniform', hasDynamicOffset: true },
      });
    }
    return device.createBindGroupLayout({ entries });
  }

  /** Ops are appended to the innermost list; loops nest via this stack. */
  const opStack: Op[][] = [[]];
  const ops = (): Op[] => opStack[opStack.length - 1];
  const enclosingLoops: string[] = [];
  const inLoop = (): boolean => enclosingLoops.length > 0;

  const spendDispatches = (n: number, span: Span): void => {
    // Inside loops the body is re-encoded per iteration; multiply out.
    let mult = 1;
    for (const f of loopTrips) mult *= f;
    dispatchBudget -= n * mult;
    if (dispatchBudget < 0) {
      throw new UnsupportedOnGpu(
        `this script would encode more than ${MAX_DISPATCHES.toLocaleString()} ` +
          `GPU dispatches; shrink loop counts`,
        span,
      );
    }
  };
  const loopTrips: number[] = [];

  /** KernelInputs seeded with fresh maps; buffers registered on demand. */
  const freshIo = (): KernelInputs => ({
    buffers: new Map(),
    loopVars: new Map(),
    nextSeed: () => seedCounter++,
  });

  /** Register every buffer-backed read in `expr` into `io.buffers`. */
  function collectBuffers(expr: IRExpr, io: KernelInputs, span: Span): void {
    const walk = (x: IRExpr): void => {
      // A constant-folded subtree emits as a literal; nothing under it is read.
      if (exactValue(x) !== undefined) return;
      switch (x.kind) {
        case 'Var': {
          if (enclosingLoops.includes(x.cName)) {
            if (!io.loopVars.has(x.cName)) io.loopVars.set(x.cName, io.loopVars.size);
            return;
          }
          if (!isNumeric(x.ty)) {
            throw new UnsupportedOnGpu(`'${x.name}' is not numeric`, x.span);
          }
          if (typeof x.ty.exact === 'number') return; // folds to a literal
          const key = resolve(x.cName);
          if (!io.buffers.has(key)) io.buffers.set(key, io.buffers.size);
          // Rewrite so the emitter sees the resolved name.
          x.cName = key;
          return;
        }
        case 'Binary':
          walk(x.left);
          walk(x.right);
          return;
        case 'Unary':
          walk(x.operand);
          return;
        case 'IndexSlice':
          walk(x.base);
          return;
        case 'MakeRange':
          walk(x.start);
          walk(x.step);
          return;
        case 'Call':
          if (!['zeros', 'ones', 'eye', 'rand', 'randn'].includes(x.name)) {
            x.args.forEach(walk);
          }
          return;
        default:
          return;
      }
    };
    walk(expr);
    void span;
  }

  /** Bind-group entries for a kernel built from `io`, writing `target`. */
  function kernelBindGroup(
    io: KernelInputs,
    bufferOrder: string[],
    loopOrder: string[],
    target: GPUBuffer,
    layout: GPUBindGroupLayout,
    span: Span,
  ): GPUBindGroup {
    const entries: GPUBindGroupEntry[] = [{ binding: 0, resource: { buffer: target } }];
    bufferOrder.forEach((cName, i) => {
      const v = vars.get(resolve(cName));
      if (!v) throw new UnsupportedOnGpu(`a value read here has no buffer`, span);
      entries.push({ binding: i + 1, resource: { buffer: v.slot.buffer } });
    });
    loopOrder.forEach((cVar, k) => {
      const u = loopUniforms.get(cVar);
      if (!u) throw new UnsupportedOnGpu(`internal: no uniform for loop '${cVar}'`, span);
      entries.push({
        binding: bufferOrder.length + 1 + k,
        resource: { buffer: u, offset: 0, size: 8 },
      });
    });
    return device.createBindGroup({ layout, entries });
  }

  /** Materialize an arbitrary tensor expression into a slot, planning
   *  whatever ops that takes. A plain Var is returned as-is. */
  async function materialize(expr: IRExpr): Promise<{ cName: string; slot: Slot }> {
    if (expr.kind === 'Var' && isTensor(expr.ty) && typeof (expr.ty as NumericType).exact !== 'object') {
      return { cName: resolve(expr.cName), slot: readSlot(expr.cName, expr.name, expr.span) };
    }
    if (expr.kind === 'IndexSlice') {
      const base = fullColonBase(expr);
      if (base && base.kind === 'Var') {
        return { cName: resolve(base.cName), slot: readSlot(base.cName, base.name, base.span) };
      }
    }
    if (!isNumeric(expr.ty)) {
      throw new UnsupportedOnGpu(`expression is not numeric`, expr.span);
    }
    const cName = `%tmp${tempCounter++}`;
    await planValue(cName, cName, expr.ty, expr, expr.span);
    return { cName, slot: vars.get(resolve(cName))!.slot };
  }

  /** Replace non-elementwise subtrees (GEMM, matrix transpose, reductions)
   *  with materialized temps, so what remains is one fused kernel. */
  async function hoistNonElementwise(expr: IRExpr): Promise<IRExpr> {
    const hoist = async (x: IRExpr): Promise<IRExpr> => {
      if (isHoistRoot(x)) {
        const { cName, slot } = await materialize(x);
        void slot;
        return {
          kind: 'Var',
          name: cName,
          cName,
          ty: x.ty,
          span: x.span,
        };
      }
      switch (x.kind) {
        case 'Binary':
          x.left = await hoist(x.left);
          x.right = await hoist(x.right);
          return x;
        case 'Unary':
          x.operand = await hoist(x.operand);
          return x;
        case 'Call':
          for (let i = 0; i < x.args.length; i++) x.args[i] = await hoist(x.args[i]);
          return x;
        case 'IndexSlice':
          x.base = await hoist(x.base);
          return x;
        default:
          return x;
      }
    };
    // The root itself was already routed by planValue; only hoist children.
    switch (expr.kind) {
      case 'Binary':
        expr.left = await hoist(expr.left);
        expr.right = await hoist(expr.right);
        return expr;
      case 'Unary':
        expr.operand = await hoist(expr.operand);
        return expr;
      case 'Call':
        for (let i = 0; i < expr.args.length; i++) expr.args[i] = await hoist(expr.args[i]);
        return expr;
      case 'IndexSlice':
        expr.base = await hoist(expr.base);
        return expr;
      default:
        return expr;
    }
  }

  function isHoistRoot(x: IRExpr): boolean {
    if (x.kind === 'Binary' && x.builtin === 'mtimes' && isTensor(x.left.ty) && isTensor(x.right.ty)) {
      return true;
    }
    if (x.kind === 'Unary' && x.builtin === 'transpose' && isTensor(x.operand.ty) && !isVectorish(x.operand.ty)) {
      return true;
    }
    if (x.kind === 'Call' && REDUCTIONS.has(x.name)) {
      // Two-arg max/min is elementwise, not a reduction.
      if ((x.name === 'max' || x.name === 'min') && x.args.length === 2) return false;
      return x.args.some((a) => isTensor(a.ty));
    }
    return false;
  }

  /** Plan `dest = expr`, routing to the right kind of op. */
  async function planValue(
    cName: string,
    name: string,
    ty: NumericType,
    expr: IRExpr,
    span: Span,
  ): Promise<void> {
    // View-or-copy forms: X(:), reshape, vector transpose, plain `B = A`.
    const viewOf = viewSource(expr);
    if (viewOf) {
      const srcSlot = readSlot(viewOf.cName, viewOf.name, span);
      const srcKey = resolve(viewOf.cName);
      if (numel(ty) !== srcSlot.count) {
        throw new UnsupportedOnGpu(`'${name}' and '${viewOf.name}' differ in size`, span);
      }
      const destAssigns = assignCounts.get(cName) ?? 1;
      const srcAssigns = assignCounts.get(viewOf.cName) ?? 1;
      if (destAssigns === 1 && srcAssigns === 1 && !inLoop() && !vars.has(cName)) {
        aliases.set(cName, srcKey);
        // Track the new orientation under the alias's own name via vars of
        // the base: display uses the Assign's ty directly, so nothing to do.
        describeLines.push(`view    ${name} -> ${viewOf.name}`);
        return;
      }
      const dest = slotFor(cName, name, ty, span);
      if (dest.buffer !== srcSlot.buffer) {
        ops().push({
          kind: 'copy',
          from: srcSlot.buffer,
          to: dest.buffer,
          bytes: 4 * srcSlot.count,
          label: `${name} = ${viewOf.name}`,
        });
        describeLines.push(`copy    ${name} = ${viewOf.name}`);
      }
      return;
    }

    // GEMM: A * B with both sides tensors.
    if (expr.kind === 'Binary' && expr.builtin === 'mtimes' && isTensor(expr.left.ty) && isTensor(expr.right.ty)) {
      return planGemm(cName, name, ty, expr, span);
    }
    // Matrix transpose.
    if (expr.kind === 'Unary' && expr.builtin === 'transpose' && isTensor(expr.operand.ty) && !isVectorish(expr.operand.ty)) {
      return planTranspose(cName, name, ty, expr, span);
    }
    // Reductions.
    if (expr.kind === 'Call' && REDUCTIONS.has(expr.name) &&
        !((expr.name === 'max' || expr.name === 'min') && expr.args.length === 2) &&
        expr.args.some((a) => isTensor(a.ty))) {
      return planReduce(cName, name, ty, expr as Call, span);
    }

    // Fused elementwise kernel (with any non-elementwise subtrees hoisted).
    const hoisted = await hoistNonElementwise(expr);
    const io = freshIo();
    collectBuffers(hoisted, io, span);
    const label = `${name} = <${numel(ty)} elem>`;
    const kernel = buildKernel({ name, cName, ty, expr: hoisted, span }, io, [...enclosingLoops], label);
    const layout = kernelLayout(kernel.buffers.length, kernel.loops.length);
    const pipe = await pipeline(kernel.code, label, layout);

    const dest = slotFor(cName, name, ty, span);
    // WebGPU forbids aliasing a writable binding with a readable one, so an
    // in-place update (`u = u + 1`) writes scratch and copies back.
    const aliased = kernel.buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
    const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;

    const bindGroup = kernelBindGroup(io, kernel.buffers, kernel.loops, target.buffer, layout, span);
    spendDispatches(1, span);
    ops().push({
      kind: 'kernel',
      pipeline: pipe,
      bindGroup,
      dispatch: dispatchFor(kernel.count),
      loops: kernel.loops,
      label,
      copyBack: aliased
        ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
        : undefined,
    });
    describeLines.push(`kernel  ${label}${aliased ? ' (+copy-back)' : ''}`);
  }

  /** `B = A`, `X(:)`, `reshape(A, ...)`, vector `x'` — pure views/copies. */
  function viewSource(expr: IRExpr): { cName: string; name: string } | null {
    if (expr.kind === 'Var' && isTensor(expr.ty)) {
      return { cName: expr.cName, name: expr.name };
    }
    if (expr.kind === 'IndexSlice') {
      const base = fullColonBase(expr);
      if (base && base.kind === 'Var' && isTensor(base.ty)) {
        return { cName: base.cName, name: base.name };
      }
      return null;
    }
    if (expr.kind === 'Call' && expr.name === 'reshape' && expr.args.length >= 1 &&
        expr.args[0].kind === 'Var' && isTensor(expr.args[0].ty)) {
      return { cName: expr.args[0].cName, name: expr.args[0].name };
    }
    if (expr.kind === 'Unary' && expr.builtin === 'transpose' &&
        isVectorish(expr.operand.ty) && expr.operand.kind === 'Var' && isTensor(expr.operand.ty)) {
      return { cName: expr.operand.cName, name: expr.operand.name };
    }
    return null;
  }

  async function planGemm(
    cName: string,
    name: string,
    ty: NumericType,
    expr: IRExpr & { kind: 'Binary' },
    span: Span,
  ): Promise<void> {
    const a = await materialize(expr.left);
    const b = await materialize(expr.right);
    const [m, k] = shapeOf(expr.left.ty, span);
    const [, n] = shapeOf(expr.right.ty, span);
    const dest = slotFor(cName, name, ty, span);
    const aliased = dest.buffer === a.slot.buffer || dest.buffer === b.slot.buffer;
    const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;

    const label = `${name} = ${m}x${k} * ${k}x${n}`;
    const layout = kernelLayout(2, 0);
    const pipe = await pipeline(gemmKernel(m, k, n), label, layout);
    const bindGroup = device.createBindGroup({
      layout,
      entries: [
        { binding: 0, resource: { buffer: target.buffer } },
        { binding: 1, resource: { buffer: a.slot.buffer } },
        { binding: 2, resource: { buffer: b.slot.buffer } },
      ],
    });
    spendDispatches(1, span);
    ops().push({
      kind: 'kernel',
      pipeline: pipe,
      bindGroup,
      dispatch: gemmDispatch(m, n),
      loops: [],
      label,
      copyBack: aliased
        ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
        : undefined,
    });
    describeLines.push(`gemm    ${label}`);
  }

  async function planTranspose(
    cName: string,
    name: string,
    ty: NumericType,
    expr: IRExpr & { kind: 'Unary' },
    span: Span,
  ): Promise<void> {
    const src = await materialize(expr.operand);
    const [m, n] = shapeOf(expr.operand.ty, span);
    const dest = slotFor(cName, name, ty, span);
    const aliased = dest.buffer === src.slot.buffer;
    const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;

    const label = `${name} = (${m}x${n})'`;
    const layout = kernelLayout(1, 0);
    const pipe = await pipeline(transposeKernel(m, n), label, layout);
    const bindGroup = device.createBindGroup({
      layout,
      entries: [
        { binding: 0, resource: { buffer: target.buffer } },
        { binding: 1, resource: { buffer: src.slot.buffer } },
      ],
    });
    spendDispatches(1, span);
    ops().push({
      kind: 'kernel',
      pipeline: pipe,
      bindGroup,
      dispatch: transposeDispatch(m, n),
      loops: [],
      label,
      copyBack: aliased
        ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
        : undefined,
    });
    describeLines.push(`transp  ${label}`);
  }

  async function planReduce(
    cName: string,
    name: string,
    ty: NumericType,
    call: Call,
    span: Span,
  ): Promise<void> {
    const fn = call.name;
    if (fn !== 'dot' && call.args.length !== 1) {
      throw new UnsupportedOnGpu(
        `'${fn}' supports only the one-argument form here (no dim/'all' ` +
          `arguments — use ${fn}(X(:)) for the whole array)`,
        span,
      );
    }

    // Build the loader expression: the (fused) element the reduction eats.
    let loaderExpr: IRExpr;
    let map: MapKind = 'id';
    if (fn === 'dot') {
      if (call.args.length !== 2) {
        throw new UnsupportedOnGpu(`'dot' takes two vectors`, span);
      }
      loaderExpr = {
        kind: 'Binary',
        builtin: 'times',
        left: call.args[0],
        right: call.args[1],
        ty: call.args[0].ty,
        span,
      };
    } else {
      loaderExpr = call.args[0];
      if (fn === 'norm') {
        if (!isVectorish(loaderExpr.ty)) {
          throw new UnsupportedOnGpu(
            `'norm' of a matrix is the spectral norm, which the sandbox does ` +
              `not compute; norm(v) for vectors only`,
            span,
          );
        }
        map = 'sq';
      }
    }
    // See through X(:) so the loader reads the base buffer directly.
    if (loaderExpr.kind === 'IndexSlice') {
      const base = fullColonBase(loaderExpr);
      if (base) loaderExpr = { ...base, ty: loaderExpr.ty } as IRExpr;
    }

    const inputTy = loaderExpr.ty as NumericType;
    const inCount = numel(inputTy);
    const outCount = numel(ty);
    const full = outCount === 1;
    if (!full) {
      const [m, n] = shapeOf(inputTy, span);
      const [om, on] = shapeOf(ty, span);
      if (om !== 1 || on !== n) {
        throw new UnsupportedOnGpu(
          `'${fn}' along that dimension is not supported — transpose first, ` +
            `or reduce the whole array with ${fn}(X(:))`,
          span,
        );
      }
      void m;
    }

    const combine: Combine =
      fn === 'prod' ? 'mul' : fn === 'max' ? 'max' : fn === 'min' ? 'min' : 'add';
    const epilogue: Epilogue =
      fn === 'mean'
        ? { kind: 'scale', by: 1 / (full ? inCount : shapeOf(inputTy, span)[0]) }
        : fn === 'norm'
          ? { kind: 'sqrt' }
          : { kind: 'none' };

    // Hoist nested non-elementwise pieces, then emit the fused loader.
    loaderExpr = await hoistNonElementwise(
      loaderExpr.kind === 'Binary' || loaderExpr.kind === 'Unary' || loaderExpr.kind === 'Call' || loaderExpr.kind === 'IndexSlice'
        ? loaderExpr
        : loaderExpr,
    );
    if (isHoistRoot(loaderExpr)) {
      const { cName: mc } = await materialize(loaderExpr);
      loaderExpr = { kind: 'Var', name: mc, cName: mc, ty: loaderExpr.ty, span };
    }
    const io = freshIo();
    collectBuffers(loaderExpr, io, span);
    const loader = emitLoader(loaderExpr, io, [...enclosingLoops]);
    const { decls, buffers, loops } = bindingDecls(io);

    const dest = slotFor(cName, name, ty, span);
    const aliased = buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
    const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;

    if (full) {
      const numWg = reducePartials(inCount);
      const partials = makeSlot(`${name}-partials`, numWg);
      const label1 = `${name} = ${fn}(<${inCount} elem>) pass1`;
      const layout1 = kernelLayout(buffers.length, loops.length);
      const pipe1 = await pipeline(
        reduceFullPass1(decls, loader, inCount, numWg, combine, map),
        label1,
        layout1,
      );
      const bg1 = kernelBindGroup(io, buffers, loops, partials.buffer, layout1, span);
      const label2 = `${name} = ${fn}(...) pass2`;
      const layout2 = kernelLayout(1, 0);
      const pipe2 = await pipeline(reduceFullPass2(numWg, combine, epilogue), label2, layout2);
      const bg2 = device.createBindGroup({
        layout: layout2,
        entries: [
          { binding: 0, resource: { buffer: target.buffer } },
          { binding: 1, resource: { buffer: partials.buffer } },
        ],
      });
      spendDispatches(2, span);
      ops().push({
        kind: 'kernel', pipeline: pipe1, bindGroup: bg1,
        dispatch: [numWg, 1], loops, label: label1,
      });
      ops().push({
        kind: 'kernel', pipeline: pipe2, bindGroup: bg2,
        dispatch: [1, 1], loops: [], label: label2,
        copyBack: aliased
          ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
          : undefined,
      });
      describeLines.push(`reduce  ${name} = ${fn}(<${inCount} elem>)`);
    } else {
      const [m, n] = shapeOf(inputTy, span);
      if (n > 65535) {
        throw new UnsupportedOnGpu(`'${fn}' over ${n} columns exceeds the dispatch limit`, span);
      }
      const label = `${name} = ${fn}(${m}x${n} by columns)`;
      const layout = kernelLayout(buffers.length, loops.length);
      const pipe = await pipeline(
        reduceColumns(decls, loader, m, combine, map, epilogue),
        label,
        layout,
      );
      const bg = kernelBindGroup(io, buffers, loops, target.buffer, layout, span);
      spendDispatches(1, span);
      ops().push({
        kind: 'kernel', pipeline: pipe, bindGroup: bg,
        dispatch: [n, 1], loops, label,
        copyBack: aliased
          ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
          : undefined,
      });
      describeLines.push(`reduce  ${label}`);
    }
  }

  function shapeOf(t: Type, span: Span): [number, number] {
    if (!isNumeric(t) || !t.shape) {
      throw new UnsupportedOnGpu(`shape is not known at compile time`, span);
    }
    if (t.shape.length !== 2) {
      throw new UnsupportedOnGpu(`only 2-D arrays are supported (got ${t.shape.length}-D)`, span);
    }
    return [t.shape[0], t.shape[1]];
  }

  // ── Host-op planning ──────────────────────────────────────────────────

  const noHostOpsInLoops = (what: string, span: Span): void => {
    if (inLoop()) {
      throw new UnsupportedOnGpu(
        `'${what}' inside a for loop is not supported: the loop body is ` +
          `compiled once and replayed on the GPU, so per-iteration host I/O ` +
          `has nowhere to run. Hoist it out of the loop (or time the whole loop).`,
        span,
      );
    }
  };

  /** A ValueRef for a scalar-valued expression the host wants to print. The
   *  inline pass folds argument temps into the call, so this routinely sees
   *  whole expressions — they get computed into a 1-element buffer. */
  async function scalarRef(e: IRExpr, span: Span): Promise<ValueRef> {
    if (isNumeric(e.ty) && isTensor(e.ty)) {
      throw new UnsupportedOnGpu(
        `printing an array here is not supported (MATLAB would recycle the ` +
          `format); print a scalar, or use disp`,
        span,
      );
    }
    const exact = exactValue(e);
    if (exact !== undefined) return { kind: 'literal', value: exact };
    if (e.kind === 'Var') {
      if (hostScalars.has(e.cName)) return { kind: 'host', cName: e.cName };
      const slot = readSlot(e.cName, e.name, span);
      return { kind: 'buffer', slot, count: 1 };
    }
    if (!isNumeric(e.ty)) {
      throw new UnsupportedOnGpu(`only numeric values can be printed`, span);
    }
    const cName = `%tmp${tempCounter++}`;
    await planValue(cName, cName, e.ty, e, span);
    return { kind: 'buffer', slot: vars.get(resolve(cName))!.slot, count: 1 };
  }

  async function planDisp(call: Call, span: Span): Promise<void> {
    noHostOpsInLoops('disp', span);
    if (call.args.length !== 1) {
      throw new UnsupportedOnGpu(`'disp' takes one argument`, span);
    }
    const a = call.args[0];
    if (a.kind === 'StringLit') {
      ops().push({ kind: 'emit', parts: [{ kind: 'text', text: a.value + '\n' }] });
      return;
    }
    if (a.kind === 'Var' && chars.has(a.cName)) {
      ops().push({ kind: 'emit', parts: [{ kind: 'text', text: chars.get(a.cName)! + '\n' }] });
      return;
    }
    if (!isNumeric(a.ty)) throw new UnsupportedOnGpu(`'disp' argument is not numeric`, span);
    if (isTensor(a.ty)) {
      if (a.kind !== 'Var') {
        throw new UnsupportedOnGpu(`'disp' of an expression — give it a name first`, span);
      }
      const slot = readSlot(a.cName, a.name, span);
      ops().push({
        kind: 'display',
        label: null,
        ref: { kind: 'buffer', slot, count: slot.count },
        shape: a.ty.shape ?? [slot.count, 1],
      });
      return;
    }
    ops().push({ kind: 'display', label: null, ref: await scalarRef(a, span), shape: [1, 1] });
  }

  async function planFprintf(call: Call, span: Span): Promise<void> {
    noHostOpsInLoops('fprintf', span);
    let args = call.args;
    // fprintf(1, fmt, ...) — MATLAB's stdout file id.
    if (args.length >= 2 && args[0].kind === 'NumLit' && args[0].value === 1) {
      args = args.slice(1);
    }
    if (args.length === 0) throw new UnsupportedOnGpu(`'fprintf' needs a format string`, span);
    const fmtArg = args[0];
    const fmt =
      fmtArg.kind === 'StringLit'
        ? fmtArg.value
        : fmtArg.kind === 'Var' && chars.has(fmtArg.cName)
          ? chars.get(fmtArg.cName)!
          : null;
    if (fmt === null) {
      throw new UnsupportedOnGpu(`'fprintf' format must be a literal string`, span);
    }
    const parts = parseFormat(fmt, span);
    const specs = parts.filter((p) => p.kind === 'spec');
    const values = args.slice(1);
    if (specs.length !== values.length) {
      throw new UnsupportedOnGpu(
        `'fprintf' format has ${specs.length} conversion(s) but ${values.length} ` +
          `value(s); the sandbox does not recycle the format over arrays`,
        span,
      );
    }
    let vi = 0;
    const emitParts: EmitPart[] = [];
    for (const p of parts) {
      if (p.kind === 'text') emitParts.push({ kind: 'text', text: p.text });
      else emitParts.push({ kind: 'value', ref: await scalarRef(values[vi++], span), spec: p.spec });
    }
    ops().push({ kind: 'emit', parts: emitParts });
  }

  /** printf-format split: text runs (escapes decoded) and % conversions. */
  function parseFormat(
    fmt: string,
    span: Span,
  ): ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] {
    const out: ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] = [];
    let text = '';
    for (let i = 0; i < fmt.length; i++) {
      const c = fmt[i];
      if (c === '\\') {
        const n = fmt[i + 1];
        if (n === 'n') { text += '\n'; i++; }
        else if (n === 't') { text += '\t'; i++; }
        else if (n === '\\') { text += '\\'; i++; }
        else text += c;
      } else if (c === '%') {
        if (fmt[i + 1] === '%') { text += '%'; i++; continue; }
        const m = /^%[-+ 0#]*\d*(?:\.\d+)?[diufeEgGs]/.exec(fmt.slice(i));
        if (!m) {
          throw new UnsupportedOnGpu(
            `'fprintf': unsupported conversion at "${fmt.slice(i, i + 6)}"`,
            span,
          );
        }
        if (text) { out.push({ kind: 'text', text }); text = ''; }
        out.push({ kind: 'spec', spec: m[0] });
        i += m[0].length - 1;
      } else {
        text += c;
      }
    }
    if (text) out.push({ kind: 'text', text });
    return out;
  }

  // ── Statement walk ────────────────────────────────────────────────────

  async function planStmt(stmt: IRStmt): Promise<void> {
    switch (stmt.kind) {
      case 'Assign':
        return planAssign(stmt);
      case 'ExprStmt':
        return planExprStmt(stmt);
      case 'For':
        return planFor(stmt);
      default:
        throw new UnsupportedOnGpu(
          `'${stmtName(stmt.kind)}' is not supported in the sandbox`,
          stmt.span,
        );
    }
  }

  function stmtName(kind: string): string {
    return (
      {
        If: 'if', While: 'while', Break: 'break', Continue: 'continue',
        IndexStore: 'indexed assignment', IndexSliceStore: 'indexed assignment',
        MultiAssignCall: 'multiple assignment', MemberStore: 'struct assignment',
        CellIndexStore: 'cell assignment',
      }[kind] ?? kind
    );
  }

  const isTemp = (name: string): boolean => name.startsWith('_mtoc2_') || name.startsWith('%tmp');

  async function planAssign(stmt: Assign): Promise<void> {
    // Char values (format strings) ride along on the host. Their type kind
    // is numbl's 'Char'/'String', not Numeric.
    if (stmt.expr.kind === 'StringLit') {
      chars.set(stmt.cName, stmt.expr.value);
      return;
    }
    if (!isNumeric(stmt.ty)) {
      throw new UnsupportedOnGpu(
        `'${stmt.name}' is not a numeric value (only numbers, strings for ` +
          `printing, and numeric arrays exist in the sandbox)`,
        stmt.span,
      );
    }

    // tic/toc as values.
    if (stmt.expr.kind === 'Call' && (stmt.expr.name === 'tic' || stmt.expr.name === 'toc')) {
      noHostOpsInLoops(stmt.expr.name, stmt.span);
      const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
      hostScalars.add(stmt.cName);
      if (stmt.expr.name === 'tic') {
        ops().push({ kind: 'tic', assignTo: { cName: stmt.cName, slot } });
      } else {
        const since = tocSince(stmt.expr, stmt.span);
        ops().push({
          kind: 'toc', print: false, sinceCName: since,
          assignTo: { cName: stmt.cName, slot }, seq: ++tocCounter,
        });
      }
      maybeEcho(stmt);
      return;
    }

    // Exact scalar: folds into every consumer. It only needs real storage if
    // the variable is reassigned elsewhere (later uses then read the buffer).
    if (typeof stmt.ty.exact === 'number') {
      if ((assignCounts.get(stmt.cName) ?? 1) > 1) {
        const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
        ops().push({
          kind: 'write', slot,
          data: new Float32Array([stmt.ty.exact]),
          label: `${stmt.name} = ${stmt.ty.exact}`,
        });
      }
      maybeEcho(stmt);
      return;
    }
    // Exact tensor (a literal like [1 2 3]): upload the data.
    if (stmt.ty.exact instanceof Float64Array) {
      const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
      ops().push({
        kind: 'write', slot,
        data: Float32Array.from(stmt.ty.exact),
        label: `${stmt.name} = <literal ${slot.count} elem>`,
      });
      maybeEcho(stmt);
      return;
    }
    if (stmt.ty.exact !== undefined) {
      throw new UnsupportedOnGpu(`complex values are not supported (f32 backend)`, stmt.span);
    }

    await planValue(stmt.cName, stmt.name, stmt.ty, stmt.expr, stmt.span);
    maybeEcho(stmt);
  }

  /** MATLAB-style echo for statements without a trailing semicolon. */
  function maybeEcho(stmt: Assign): void {
    if (isTemp(stmt.name)) return;
    if (!compiled.isEchoed(stmt.span.start)) return;
    noHostOpsInLoops(`echo of '${stmt.name}' (add a semicolon)`, stmt.span);
    if (!isNumeric(stmt.ty)) return;
    const exact = typeof stmt.ty.exact === 'number' ? stmt.ty.exact : undefined;
    const shape = stmt.ty.shape ?? [1, 1];
    if (exact !== undefined) {
      ops().push({
        kind: 'display', label: stmt.name,
        ref: { kind: 'literal', value: exact }, shape,
      });
      return;
    }
    if (hostScalars.has(stmt.cName)) {
      ops().push({
        kind: 'display', label: stmt.name,
        ref: { kind: 'host', cName: stmt.cName }, shape,
      });
      return;
    }
    const v = vars.get(resolve(stmt.cName));
    if (!v) return;
    ops().push({
      kind: 'display', label: stmt.name,
      ref: { kind: 'buffer', slot: v.slot, count: v.slot.count }, shape,
    });
  }

  function tocSince(call: IRExpr & { kind: 'Call' }, span: Span): string | undefined {
    if (call.args.length === 0) return undefined;
    const a = call.args[0];
    if (a.kind === 'Var' && hostScalars.has(a.cName)) return a.cName;
    throw new UnsupportedOnGpu(
      `'toc(t)' needs a value produced by 't = tic'`,
      span,
    );
  }

  async function planExprStmt(stmt: ExprStmt): Promise<void> {
    const e = stmt.expr;
    if (e.kind === 'Call') {
      switch (e.name) {
        case 'tic':
          noHostOpsInLoops('tic', stmt.span);
          ops().push({ kind: 'tic' });
          return;
        case 'toc':
        case 'toc_print': // numbl's lowering of a bare `toc` statement
          noHostOpsInLoops('toc', stmt.span);
          ops().push({
            kind: 'toc', print: true,
            sinceCName: tocSince(e, stmt.span), seq: ++tocCounter,
          });
          return;
        case 'disp':
          return planDisp(e, stmt.span);
        case 'fprintf':
          return planFprintf(e, stmt.span);
        case 'rng':
          throw new UnsupportedOnGpu(
            `'rng' is not supported: the sandbox's rand/randn streams are ` +
              `deterministic per run already`,
            stmt.span,
          );
        default:
          break;
      }
    }
    // A bare expression with a value: numbl assigns display-relevant results
    // to `ans` as an Assign, so a leftover ExprStmt is side-effect-free.
    if (e.kind === 'Call') {
      throw new UnsupportedOnGpu(`'${e.name}' is not supported in the sandbox`, stmt.span);
    }
  }

  async function planFor(stmt: For): Promise<void> {
    const from = exactValue(stmt.start);
    const to = exactValue(stmt.end);
    if (from === undefined || to === undefined) {
      throw new UnsupportedOnGpu(
        `a for loop's bounds must be known when the script is compiled ` +
          `(assign them from literals)`,
        stmt.span,
      );
    }
    const trips = Math.floor((to - from) / stmt.step + 1e-9) + 1;
    if (trips <= 0) return; // never executes
    if (trips > MAX_TRIPS) {
      throw new UnsupportedOnGpu(
        `'for ${stmt.varName}' runs ${trips} iterations, over the sandbox ` +
          `limit of ${MAX_TRIPS}`,
        stmt.span,
      );
    }

    // One 256-byte uniform slot per iteration: { v: f32, it: u32 }.
    const uniform = device.createBuffer({
      label: `mgpu-loop-${stmt.varName}`,
      size: trips * LV_STRIDE,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    owned.push(uniform);
    const data = new ArrayBuffer(trips * LV_STRIDE);
    const f32 = new Float32Array(data);
    const u32 = new Uint32Array(data);
    for (let i = 0; i < trips; i++) {
      f32[(i * LV_STRIDE) / 4] = from + i * stmt.step;
      u32[(i * LV_STRIDE) / 4 + 1] = i;
    }
    device.queue.writeBuffer(uniform, 0, data);
    loopUniforms.set(stmt.cVar, uniform);

    const body: Op[] = [];
    opStack.push(body);
    enclosingLoops.push(stmt.cVar);
    loopTrips.push(trips);
    try {
      for (const s of stmt.body) await planStmt(s);
    } finally {
      loopTrips.pop();
      enclosingLoops.pop();
      opStack.pop();
    }
    ops().push({
      kind: 'loop', cVar: stmt.cVar, trips, uniform, body,
      label: `for ${stmt.varName} = ${from}:${stmt.step}:${to}`,
    });
    describeLines.push(`loop    for ${stmt.varName} (${trips} iterations, body above)`);

    // MATLAB leaves the loop variable holding its final value; make that
    // readable afterwards.
    const finalValue = from + (trips - 1) * stmt.step;
    const slot = slotFor(
      stmt.cVar, stmt.varName,
      { kind: 'Numeric', elem: 'double', isComplex: false,
        dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
        shape: [1, 1], sign: 'unknown' },
      stmt.span,
    );
    ops().push({
      kind: 'write', slot, data: new Float32Array([finalValue]),
      label: `${stmt.varName} = ${finalValue} (final)`,
    });
  }

  for (const stmt of compiled.stmts) {
    await planStmt(stmt);
  }

  return {
    ops: opStack[0],
    describe: () => describeLines,
    kernels: kernelSources,
    destroy: () => {
      for (const b of owned) b.destroy();
      owned.length = 0;
    },
  };
}