/**
 * Sandbox-side fusion: fold the single-use ANF temps numbl's inline pass
 * left behind.
 *
 * numbl's pass (`inlinePass`) only folds producers its C backend can fuse —
 * no tensor-producing Calls (`sin(x)`), no logical results, no ranges. The
 * WGSL emitter fuses all of those, so without this pass a line like
 * `y = y + 0.1*sin(x + k) .* exp(-x)` becomes five kernels instead of one.
 *
 * Same shape as numbl's pass, deliberately narrower where it matters:
 *  - only compiler temps (`_mtoc2_*`) are folded, so every user-named
 *    variable still materializes — one source line stays one kernel, and
 *    anything the user might echo or reuse keeps its buffer;
 *  - the producer's RHS must be fusable on the GPU (isGpuFusableExpr);
 *  - the single use must be in a following Assign/ExprStmt at the same body
 *    level, with no intervening write to the producer's operands and no
 *    control-flow statement in between.
 *
 * `for` bodies are processed as their own levels; nothing folds across the
 * loop boundary.
 */
import type { Assign, ExprStmt, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import { isMultiElement, type Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
import { isGpuFusableExpr } from './wgsl.ts';

const isMulti = (t: Type): boolean => t.kind === 'Numeric' && isMultiElement(t);

export function fuseTemps(stmts: IRStmt[]): void {
  for (const s of stmts) {
    if (s.kind === 'For') fuseTemps(s.body);
  }
  for (let iter = 0; iter < 32; iter++) {
    if (!fuseOnePass(stmts)) break;
  }
}

const isTemp = (cName: string): boolean => cName.startsWith('_mtoc2_');

function fuseOnePass(stmts: IRStmt[]): boolean {
  const uses = useCounts(stmts);
  for (let i = 0; i < stmts.length; i++) {
    const p = stmts[i];
    if (p.kind !== 'Assign') continue;
    if (!isTemp(p.cName)) continue;
    if (uses.get(p.cName) !== 1) continue;
    if (!isGpuFusableExpr(p.expr)) continue;

    const reads = new Set<string>();
    walkVars(p.expr, (c) => reads.add(c));

    for (let j = i + 1; j < stmts.length; j++) {
      const c = stmts[j];
      if (c.kind !== 'Assign' && c.kind !== 'ExprStmt') break; // control flow
      const holder = c as Assign | ExprStmt;
      if (c.kind === 'ExprStmt' && isMulti(p.ty) && countIn(holder.expr, p.cName) === 1) {
        break; // disp/fprintf read tensors by name, not as expressions
      }
      if (countIn(holder.expr, p.cName) === 1) {
        holder.expr = substitute(holder.expr, p.cName, p.expr);
        stmts.splice(i, 1);
        return true;
      }
      if (c.kind === 'Assign' && (c.cName === p.cName || reads.has(c.cName))) {
        break; // intervening write invalidates the fold window
      }
    }
  }
  return false;
}

/** Var-read counts by cName, across this level AND nested bodies (a use
 *  inside a nested loop must keep the temp alive at this level). */
function useCounts(stmts: IRStmt[]): Map<string, number> {
  const counts = new Map<string, number>();
  const bump = (c: string): void => {
    counts.set(c, (counts.get(c) ?? 0) + 1);
  };
  const walkStmts = (list: IRStmt[]): void => {
    for (const s of list) {
      if (s.kind === 'Assign' || s.kind === 'ExprStmt') walkVars(s.expr, bump);
      else if (s.kind === 'For') {
        walkVars(s.start, bump);
        walkVars(s.end, bump);
        walkStmts(s.body);
      }
    }
  };
  walkStmts(stmts);
  return counts;
}

function walkVars(e: IRExpr, visit: (cName: string) => void): void {
  const walk = (x: IRExpr): void => {
    switch (x.kind) {
      case 'Var':
        visit(x.cName);
        return;
      case 'Binary':
        walk(x.left);
        walk(x.right);
        return;
      case 'Unary':
        walk(x.operand);
        return;
      case 'Call':
        x.args.forEach(walk);
        return;
      case 'IndexSlice':
        walk(x.base);
        x.index.forEach((a) => {
          const inner = (a as { expr?: IRExpr }).expr;
          if (inner) walk(inner);
        });
        return;
      case 'MakeRange':
        walk(x.start);
        walk(x.step);
        walk(x.end);
        return;
      default:
        return;
    }
  };
  walk(e);
}

function countIn(e: IRExpr, cName: string): number {
  let n = 0;
  walkVars(e, (c) => {
    if (c === cName) n++;
  });
  return n;
}

/** Replace the (single) `Var` read of `cName` with `replacement`. */
function substitute(e: IRExpr, cName: string, replacement: IRExpr): IRExpr {
  const sub = (x: IRExpr): IRExpr => {
    if (x.kind === 'Var' && x.cName === cName) return replacement;
    switch (x.kind) {
      case 'Binary':
        x.left = sub(x.left);
        x.right = sub(x.right);
        return x;
      case 'Unary':
        x.operand = sub(x.operand);
        return x;
      case 'Call':
        for (let i = 0; i < x.args.length; i++) x.args[i] = sub(x.args[i]);
        return x;
      case 'IndexSlice':
        x.base = sub(x.base);
        return x;
      case 'MakeRange':
        x.start = sub(x.start);
        x.step = sub(x.step);
        x.end = sub(x.end);
        return x;
      default:
        return x;
    }
  };
  return sub(e);
}