/**
 * IR expression tree -> one WGSL compute kernel.
 *
 * The elementwise core follows turing-surface's emitter: for an `Assign` whose
 * right-hand side is purely element-wise over operands of the target's shape,
 * emit a single kernel that computes one output element per invocation.
 * Because numbl's inline pass has already folded the ANF temps back together,
 * one source line of MATLAB becomes one kernel.
 *
 * The sandbox extends it with:
 *  - comparisons and eager logicals (`<`, `&`, `~`, ...), carried as f32 0/1;
 *  - inline *generators* — `rand`, `randn`, `linspace`, ranges, `zeros`,
 *    `ones`, `eye` — evaluated per element from the linear index, so
 *    `x = 2*rand(n,1) - 1` is one kernel and touches no other buffer;
 *  - runtime scalars read from 1-element storage buffers (`in3[0]`);
 *  - loop variables read from a per-loop dynamic-offset uniform, so a `for`
 *    body compiles once and replays.
 *
 * Everything is f32 — WGSL has no f64. Arrays are column-major linear buffers,
 * matching MATLAB, so `A(:)` and `reshape` are views of the same buffer.
 */
import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
import type {
  IRExpr,
  Assign,
  IndexSlice,
  Span,
} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
import { UnsupportedOnGpu } from './errors.ts';

export const WORKGROUP_SIZE = 64;

const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
export const numel = (t: NumericType): number =>
  (t.shape ?? []).reduce((a, b) => a * b, 1);

/** The compile-time value of a scalar expression, if it has one. */
export const exactValue = (e: IRExpr): number | undefined => {
  if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
  return e.kind === 'NumLit' ? e.value : undefined;
};

/** Element-wise binary builtins -> WGSL infix operator. */
const BINARY_OPS: Record<string, string> = {
  plus: '+',
  minus: '-',
  times: '*',
  rdivide: '/',
  // Degenerate to element-wise when at least one side is a scalar; the
  // both-tensor (true matrix) case never reaches here — the planner routes it
  // to the GEMM kernel (mtimes) or rejects it (mrdivide).
  mtimes: '*',
  mrdivide: '/',
};

/** Comparison builtins -> WGSL comparison; result carried as f32 0/1. */
const COMPARE_OPS: Record<string, string> = {
  lt: '<',
  le: '<=',
  gt: '>',
  ge: '>=',
  eq: '==',
  ne: '!=',
};

/** Element-wise unary builtins -> WGSL prefix operator. */
const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };

/** Element-wise builtin calls -> WGSL builtin of the same arity. */
const CALL_FNS: Record<string, string> = {
  abs: 'abs',
  acos: 'acos',
  asin: 'asin',
  atan: 'atan',
  atan2: 'atan2',
  ceil: 'ceil',
  cos: 'cos',
  cosh: 'cosh',
  exp: 'exp',
  fix: 'trunc',
  floor: 'floor',
  log: 'log',
  log2: 'log2',
  round: 'round',
  sign: 'sign',
  sin: 'sin',
  sinh: 'sinh',
  sqrt: 'sqrt',
  tan: 'tan',
  tanh: 'tanh',
};

/** Reductions the planner materializes before a kernel is built. Their 1-arg
 *  (and for dot, 2-arg) tensor forms never reach the elementwise emitter. */
export const REDUCTIONS = new Set([
  'sum', 'mean', 'prod', 'max', 'min', 'norm', 'dot',
]);

/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
 *  infers AbstractInt and rejects the mixed-type arithmetic. */
function f32Lit(v: number): string {
  if (!Number.isFinite(v)) {
    // WGSL has no NaN/Inf literal, and 0.0/0.0 is a const-eval error;
    // bitcast the IEEE pattern instead.
    if (Number.isNaN(v)) return 'bitcast<f32>(0x7fc00000u)';
    return v > 0 ? 'bitcast<f32>(0x7f800000u)' : 'bitcast<f32>(0xff800000u)';
  }
  return Number.isInteger(v) && Math.abs(v) < 1e21
    ? `${v}.0`
    : String(v).includes('e')
      ? `${v}f`
      : String(v);
}

/** `A(:)` — the only IndexSlice this backend executes. Returns the base Var
 *  expression, which reads the same buffer at the same linear index (a
 *  column-major flatten is the identity on the linear buffer). */
export function fullColonBase(e: IndexSlice): IRExpr | null {
  if (e.index.length !== 1 || e.index[0].kind !== 'Colon') return null;
  return e.base;
}

/** How operands are read inside a kernel. */
export interface KernelInputs {
  /** cName -> storage binding slot, for tensors AND runtime scalars (a
   *  runtime scalar is a 1-element buffer, read as `inN[0]`). */
  buffers: Map<string, number>;
  /** cVar -> dense per-kernel index of enclosing loop variables, each bound
   *  as a dynamic-offset uniform (`lvK`). */
  loopVars: Map<string, number>;
  /** Distinct-per-call-site seeds for `rand`/`randn`. Global to the plan, so
   *  two kernels never share a stream. */
  nextSeed: () => number;
}

interface Ctx {
  io: KernelInputs;
  /** Emitted-helper flags, gathered during emission. */
  usesHash: boolean;
  usedPows: Set<number>;
  usesMod: boolean;
  usesRem: boolean;
}

/** Mix every enclosing loop's iteration counter into a hash lane, so a
 *  generator inside a replayed loop draws fresh values each iteration. */
function seedExpr(seed: number, ctx: Ctx): string {
  let s = `${seed >>> 0}u`;
  for (const [, k] of ctx.io.loopVars) {
    s += ` ^ (lv${k}.it * ${[2654435761, 2246822519, 3266489917, 668265263][k % 4]}u)`;
  }
  return s;
}

/** Is `t` a value with more than one element? (logical/double both count) */
const multi = (t: Type): boolean => isTensor(t);

/**
 * Emit the per-element WGSL expression for `e`. `i` is the element index
 * variable in scope. Comparisons/logicals produce f32 0/1 so any consumer
 * can treat them as numbers, exactly like MATLAB's logicals.
 */
function emitExpr(e: IRExpr, ctx: Ctx): string {
  // Anything numbl constant-folded (pi, 2*pi, n-1, ...) is a literal, no
  // matter what expression kind computed it.
  const exact = exactValue(e);
  if (exact !== undefined) return f32Lit(exact);
  const io = ctx.io;
  switch (e.kind) {
    case 'NumLit':
      return f32Lit(e.value);

    case 'Var': {
      const lv = io.loopVars.get(e.cName);
      if (lv !== undefined) return `lv${lv}.v`;
      if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
        return f32Lit(e.ty.exact);
      }
      const slot = io.buffers.get(e.cName);
      if (slot === undefined) {
        throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
      }
      return multi(e.ty) ? `in${slot}[i]` : `in${slot}[0]`;
    }

    case 'IndexSlice': {
      const base = fullColonBase(e);
      if (!base || base.kind !== 'Var') {
        throw new UnsupportedOnGpu(
          `only the full linearization 'X(:)' of a variable is supported — ` +
            `general indexing/slicing is not implemented on the GPU yet`,
          e.span,
        );
      }
      return emitExpr(base, ctx);
    }

    case 'MakeRange': {
      // start + i*step; the count is already fixed in the node's type.
      const start = emitScalar(e.start, ctx, `range start`);
      const stepV = exactValue(e.step);
      if (stepV === undefined) {
        throw new UnsupportedOnGpu(`a range's step must be a compile-time value`, e.span);
      }
      return `(${start} + f32(i) * ${f32Lit(stepV)})`;
    }

    case 'Binary': {
      if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
          multi(e.left.ty) && multi(e.right.ty)) {
        // The planner materializes tensor mtimes into a GEMM before building
        // the kernel; reaching here means it could not (mrdivide) or a
        // planner bug (mtimes).
        throw new UnsupportedOnGpu(
          e.builtin === 'mrdivide'
            ? `matrix '/' (mrdivide) is not supported; use './' or a factorization`
            : `internal: tensor '*' was not materialized as a GEMM`,
          e.span,
        );
      }
      if (e.builtin === 'power' || e.builtin === 'mpower') {
        return emitPower(e.left, e.right, ctx, e.span);
      }
      const cmp = COMPARE_OPS[e.builtin];
      if (cmp) {
        return `select(0.0, 1.0, ${emitExpr(e.left, ctx)} ${cmp} ${emitExpr(e.right, ctx)})`;
      }
      if (e.builtin === 'and' || e.builtin === 'or' || e.builtin === 'andand' || e.builtin === 'oror') {
        const op = e.builtin === 'and' || e.builtin === 'andand' ? '&&' : '||';
        return `select(0.0, 1.0, (${emitExpr(e.left, ctx)} != 0.0) ${op} (${emitExpr(e.right, ctx)} != 0.0))`;
      }
      const op = BINARY_OPS[e.builtin];
      if (!op) {
        throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
      }
      return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
    }

    case 'Unary': {
      if (e.builtin === 'not') {
        return `select(1.0, 0.0, ${emitExpr(e.operand, ctx)} != 0.0)`;
      }
      if (e.builtin === 'transpose') {
        // A vector transpose changes orientation only — the linear buffer is
        // identical. A matrix transpose is materialized by the planner.
        if (isVectorish(e.operand.ty)) return emitExpr(e.operand, ctx);
        throw new UnsupportedOnGpu(
          `internal: matrix transpose was not materialized`,
          e.span,
        );
      }
      const op = UNARY_OPS[e.builtin];
      if (!op) {
        throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
      }
      return `(${op}${emitExpr(e.operand, ctx)})`;
    }

    case 'Call':
      return emitCall(e, ctx);

    default:
      throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
  }
}

const isVectorish = (t: Type): boolean =>
  isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;

/** A scalar-valued subexpression (range endpoints, linspace args). */
function emitScalar(e: IRExpr, ctx: Ctx, what: string): string {
  if (multi(e.ty)) {
    throw new UnsupportedOnGpu(`${what} must be a scalar`, e.span);
  }
  return emitExpr(e, ctx);
}

function emitCall(e: IRExpr & { kind: 'Call' }, ctx: Ctx): string {
  switch (e.name) {
    case 'zeros':
      return '0.0';
    case 'ones':
      return '1.0';
    case 'eye': {
      const t = e.ty;
      const m = isNumeric(t) && t.shape ? t.shape[0] : undefined;
      if (m === undefined) {
        throw new UnsupportedOnGpu(`'eye' needs a compile-time size`, e.span);
      }
      return `select(0.0, 1.0, (i % ${m}u) == (i / ${m}u))`;
    }
    case 'rand': {
      ctx.usesHash = true;
      return `rand01(i, ${seedExpr(ctx.io.nextSeed(), ctx)})`;
    }
    case 'randn': {
      ctx.usesHash = true;
      const a = seedExpr(ctx.io.nextSeed(), ctx);
      const b = seedExpr(ctx.io.nextSeed(), ctx);
      // Box–Muller; rand01 returns (0,1) so the log is finite.
      return `(sqrt(-2.0 * log(rand01(i, ${a}))) * cos(6.283185307179586 * rand01(i, ${b})))`;
    }
    case 'linspace': {
      if (e.args.length !== 3) {
        throw new UnsupportedOnGpu(`'linspace' needs 3 arguments here`, e.span);
      }
      const n = exactValue(e.args[2]);
      if (n === undefined) {
        throw new UnsupportedOnGpu(`'linspace' count must be a compile-time value`, e.span);
      }
      const a = emitScalar(e.args[0], ctx, `'linspace' start`);
      const b = emitScalar(e.args[1], ctx, `'linspace' end`);
      if (n <= 1) return b; // MATLAB: linspace(a, b, 1) == b
      return `(${a} + f32(i) * ((${b} - ${a}) * ${f32Lit(1 / (n - 1))}))`;
    }
    case 'mod':
      ctx.usesMod = true;
      return `mod_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
    case 'rem':
      ctx.usesRem = true;
      return `rem_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
    case 'max':
    case 'min': {
      // Two-arg elementwise form; the one-arg reduction never reaches here.
      if (e.args.length !== 2) {
        throw new UnsupportedOnGpu(`internal: '${e.name}' reduction was not materialized`, e.span);
      }
      return `${e.name}(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
    }
    case 'xor':
      return `select(0.0, 1.0, (${emitExpr(e.args[0], ctx)} != 0.0) != (${emitExpr(e.args[1], ctx)} != 0.0))`;
    case 'double':
    case 'logical':
      // Representation is f32 either way.
      return emitExpr(e.args[0], ctx);
    default: {
      const fn = CALL_FNS[e.name];
      if (!fn) {
        const isUserFunction = e.cName !== e.name;
        throw new UnsupportedOnGpu(
          isUserFunction
            ? `'${e.name}' is a user-defined function — not supported in the sandbox; inline it`
            : REDUCTIONS.has(e.name)
              ? `internal: '${e.name}' reduction was not materialized`
              : `'${e.name}' cannot be evaluated element-wise on the GPU`,
          e.span,
        );
      }
      return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
    }
  }
}

/**
 * `x.^k`. WGSL's `pow` is undefined for a negative base, so expand literal
 * integer exponents into repeated multiplication — which is also what makes
 * `u.^2` free. Non-integer exponents fall through to `pow`, defined only for
 * a non-negative base (as in MATLAB, where a negative base goes complex —
 * here it is NaN, and the result cross-check will show it).
 */
function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: Span): string {
  const k = exactValue(exponent);
  const b = emitExpr(base, ctx);
  if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 16) {
    if (k === 0) return '1.0';
    ctx.usedPows.add(k);
    return `pow_i${k}(${b})`;
  }
  if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -16) {
    ctx.usedPows.add(-k);
    return `(1.0 / pow_i${-k}(${b}))`;
  }
  return `pow(${b}, ${emitExpr(exponent, ctx)})`;
}

/** Fixed-exponent power helpers, emitted only when used. */
function powHelpers(used: Set<number>): string {
  const out: string[] = [];
  for (const k of [...used].sort((a, b) => a - b)) {
    const body = k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
    out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
  }
  return out.join('\n');
}

/** PCG-style hash -> (0,1). Counter-based: a call site's stream is a pure
 *  function of (element index, seed), so runs are reproducible. */
const HASH_HELPERS = `
fn hash_u(x0: u32) -> u32 {
  var x = x0 * 747796405u + 2891336453u;
  x = ((x >> ((x >> 28u) + 4u)) ^ x) * 277803737u;
  return (x >> 22u) ^ x;
}
fn rand01(i: u32, seed: u32) -> f32 {
  return (f32(hash_u(i ^ (seed * 2654435769u)) & 0x00FFFFFFu) + 0.5) * (1.0 / 16777216.0);
}`;

const MOD_HELPER = `
fn mod_m(a: f32, b: f32) -> f32 { return select(a - b * floor(a / b), a, b == 0.0); }`;
const REM_HELPER = `
fn rem_m(a: f32, b: f32) -> f32 { return select(a - b * trunc(a / b), a, b == 0.0); }`;

/**
 * Reject implicit expansion (broadcasting).
 *
 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
 * expansion semantics — but a kernel that walks one linear index across every
 * operand would quietly compute the wrong thing. So every multi-element
 * operand must have exactly the target's number of elements (a flattened or
 * transposed vector reads the same linear buffer, so only numel must match).
 */
function checkShapes(e: IRExpr, target: NumericType, name: string): void {
  const want = numel(target);
  const walk = (x: IRExpr): void => {
    if (isNumeric(x.ty) && isMultiElement(x.ty)) {
      const got = x.ty.shape ? numel(x.ty) : undefined;
      if (got !== want) {
        throw new UnsupportedOnGpu(
          `'${name}' would need implicit expansion: an operand is ` +
            `${x.ty.shape?.join('x') ?? 'dynamic'} but the result has ${want} ` +
            `elements. Expand it explicitly (the GPU kernel walks one linear ` +
            `index across every operand).`,
          x.span,
        );
      }
      // Same-numel vectors of different orientation share a linear layout;
      // same-numel *matrices* of different shape do not (transpose is not a
      // relayout numbl would insert silently, so shapes agree here).
    }
    switch (x.kind) {
      case 'Binary':
        walk(x.left);
        walk(x.right);
        return;
      case 'Unary':
        walk(x.operand);
        return;
      case 'IndexSlice':
        return; // the base reads through the slice's own (checked) type
      case 'Call':
        // A generator's arguments are sizes/endpoints, not per-element data.
        if (!['zeros', 'ones', 'eye', 'rand', 'randn', 'linspace'].includes(x.name)) {
          x.args.forEach(walk);
        }
        return;
      default:
        return;
    }
  };
  walk(e);
}

export interface Kernel {
  code: string;
  /** Number of output elements. */
  count: number;
  label: string;
  /** Buffer operand cNames in binding order (bindings 1..n). */
  buffers: string[];
  /** Loop-variable cVars in binding order (after the buffers). */
  loops: string[];
}

/** True if the expression draws random numbers anywhere. */
function usesRandom(e: IRExpr): boolean {
  let found = false;
  const walk = (x: IRExpr): void => {
    if (found) return;
    switch (x.kind) {
      case 'Call':
        if (x.name === 'rand' || x.name === 'randn') found = true;
        else x.args.forEach(walk);
        return;
      case 'Binary':
        walk(x.left);
        walk(x.right);
        return;
      case 'Unary':
        walk(x.operand);
        return;
      case 'IndexSlice':
        walk(x.base);
        return;
      default:
        return;
    }
  };
  walk(e);
  return found;
}

/** Every buffer-backed variable the expression reads (tensors and runtime
 *  scalars), and every loop variable. */
export function collectReads(
  e: IRExpr,
  isLoopVar: (cName: string) => boolean,
  isExact: (x: IRExpr) => boolean,
  visitBuffer: (cName: string) => void,
  visitLoop: (cName: string) => void,
): void {
  const walk = (x: IRExpr): void => {
    switch (x.kind) {
      case 'Var':
        if (isLoopVar(x.cName)) visitLoop(x.cName);
        else if (!isExact(x)) visitBuffer(x.cName);
        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(e);
}

/**
 * Can this expression live inside one fused GPU kernel?
 *
 * Wider than numbl's own `isPureElementwiseExpr`: the WGSL emitter fuses
 * transcendental calls, comparisons/logicals, generators, ranges, `X(:)` and
 * vector transposes, all of which numbl's C-side pass declines. The sandbox's
 * fuse pass uses this to fold the temps numbl's inline pass left behind.
 */
export function isGpuFusableExpr(e: IRExpr): boolean {
  if (exactValue(e) !== undefined) return true;
  switch (e.kind) {
    case 'NumLit':
      return true;
    case 'Var':
      return isNumeric(e.ty);
    case 'Binary': {
      if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
          multi(e.left.ty) && multi(e.right.ty)) {
        return false;
      }
      if (e.builtin === 'mpower' && multi(e.left.ty)) return false; // matrix power
      const known =
        e.builtin in BINARY_OPS || e.builtin in COMPARE_OPS ||
        ['and', 'or', 'andand', 'oror', 'power', 'mpower'].includes(e.builtin);
      return known && isGpuFusableExpr(e.left) && isGpuFusableExpr(e.right);
    }
    case 'Unary': {
      if (e.builtin === 'transpose') {
        return isVectorish(e.operand.ty) && isGpuFusableExpr(e.operand);
      }
      return (e.builtin in UNARY_OPS || e.builtin === 'not') && isGpuFusableExpr(e.operand);
    }
    case 'Call': {
      if (['zeros', 'ones', 'eye', 'rand', 'randn'].includes(e.name)) return true;
      if (e.name === 'linspace') return e.args.length === 3;
      if (['mod', 'rem', 'xor', 'atan2'].includes(e.name)) {
        return e.args.every(isGpuFusableExpr);
      }
      if ((e.name === 'max' || e.name === 'min') && e.args.length === 2) {
        return e.args.every(isGpuFusableExpr);
      }
      if (e.name === 'double' || e.name === 'logical') {
        return e.args.length === 1 && isGpuFusableExpr(e.args[0]);
      }
      return e.name in CALL_FNS && e.args.every(isGpuFusableExpr);
    }
    case 'IndexSlice': {
      const base = fullColonBase(e);
      return !!base && base.kind === 'Var';
    }
    case 'MakeRange':
      return exactValue(e.step) !== undefined;
    default:
      return false;
  }
}

/** A fused elementwise subexpression, for embedding inside a non-elementwise
 *  kernel (a reduction's per-element load). `body` reads element `i`. */
export interface FusedLoader {
  body: string;
  helpers: string;
}

/**
 * Emit `e` as a per-element load for a reduction kernel. Same contract as
 * `buildKernel`: `io.buffers` maps operands to binding slots; loop variables
 * are bound on demand (all of them, if the expression draws random numbers).
 */
export function emitLoader(
  e: IRExpr,
  io: KernelInputs,
  enclosingLoops: string[],
): FusedLoader {
  if (usesRandom(e)) {
    for (const cVar of enclosingLoops) {
      if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
    }
  }
  const ctx: Ctx = {
    io,
    usesHash: false,
    usedPows: new Set(),
    usesMod: false,
    usesRem: false,
  };
  const body = emitExpr(e, ctx);
  const helpers = [
    ctx.usesHash ? HASH_HELPERS : '',
    ctx.usesMod ? MOD_HELPER : '',
    ctx.usesRem ? REM_HELPER : '',
    powHelpers(ctx.usedPows),
  ].filter(Boolean).join('\n');
  return { body, helpers };
}

/** Binding declarations shared by every kernel shape: output at 0, operand
 *  buffers after it, loop-variable uniforms after those. */
export function bindingDecls(io: KernelInputs): {
  decls: string[];
  buffers: string[];
  loops: string[];
} {
  const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
  const buffers: string[] = [];
  for (const [cName, slot] of io.buffers) {
    buffers[slot] = cName;
    decls.push(
      `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
    );
  }
  const loops: string[] = [];
  if (io.loopVars.size) {
    decls.push(`struct Lv { v: f32, it: u32 }`);
    for (const [cVar, k] of io.loopVars) {
      loops[k] = cVar;
      decls.push(
        `@group(0) @binding(${io.buffers.size + 1 + k}) var<uniform> lv${k}: Lv;`,
      );
    }
  }
  return { decls, buffers, loops };
}

/**
 * Build the fused elementwise kernel for one `Assign`. `io.buffers` must
 * already map every buffer operand to a binding slot; the output is binding 0
 * and loop-variable uniforms follow the last input.
 *
 * `enclosingLoops` lists the cVars of the loops this statement sits inside
 * (outermost first). Loop variables the expression reads are bound; if the
 * expression draws random numbers, ALL enclosing loop counters are bound and
 * mixed into the stream so each replayed iteration draws fresh values.
 */
export function buildKernel(
  stmt: Pick<Assign, 'name' | 'cName' | 'ty' | 'expr' | 'span'>,
  io: KernelInputs,
  enclosingLoops: string[],
  label: string,
): Kernel {
  if (!isNumeric(stmt.ty)) {
    throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
  }
  if (stmt.ty.isComplex) {
    throw new UnsupportedOnGpu(
      `'${stmt.name}' is complex; the GPU backend is real-only (f32)`,
      stmt.span,
    );
  }
  const count = numel(stmt.ty);
  if (isMultiElement(stmt.ty)) checkShapes(stmt.expr, stmt.ty, stmt.name);

  // Random draws need every enclosing loop counter; make sure they are bound
  // before emission asks for them.
  if (usesRandom(stmt.expr)) {
    for (const cVar of enclosingLoops) {
      if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
    }
  }

  const ctx: Ctx = {
    io,
    usesHash: false,
    usedPows: new Set(),
    usesMod: false,
    usesRem: false,
  };
  const body = emitExpr(stmt.expr, ctx);
  const { decls, buffers, loops } = bindingDecls(io);

  const helpers = [
    ctx.usesHash ? HASH_HELPERS : '',
    ctx.usesMod ? MOD_HELPER : '',
    ctx.usesRem ? REM_HELPER : '',
    powHelpers(ctx.usedPows),
  ].filter(Boolean).join('\n');

  // Dispatch is 2-D so counts past 65535 workgroups still fit: x rows of
  // ELEMENTS_PER_ROW elements each. dispatchFor() picks matching counts.
  const code = `${decls.join('\n')}
${helpers}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x + gid.y * ${ELEMENTS_PER_ROW}u;
  if (i >= ${count}u) { return; }
  out[i] = ${body};
}
`;
  return { code, count, label, buffers, loops };
}

/** Elements covered by one row of the 2-D elementwise dispatch. */
export const ELEMENTS_PER_ROW = 32768 * WORKGROUP_SIZE;

/** Workgroup counts for an elementwise dispatch over `count` elements. */
export function dispatchFor(count: number): [number, number] {
  const rows = Math.ceil(count / ELEMENTS_PER_ROW);
  const x = rows === 1 ? Math.ceil(count / WORKGROUP_SIZE) : 32768;
  return [x, rows];
}