/**
 * The non-elementwise kernels: tiled GEMM, tiled transpose, and reductions.
 *
 * Everything is column-major f32, matching MATLAB's layout, so `A(:)` walks
 * the same linear buffer the kernels index. All shapes are compile-time
 * constants baked into the WGSL — nothing here reads a dims uniform, which is
 * what lets the planner prebuild every pipeline and bind group.
 *
 * The GEMM is the classic 16x16 shared-memory tile (adapted to column-major
 * from matmul-bench's row-major sgemm). It will not beat a tuned BLAS, but it
 * is the honest baseline for "what does A*B cost in a browser".
 *
 * Reductions take a *fused loader*: the per-element expression emitted by
 * wgsl.ts, so `sum(a .* b + c)` reads its operands exactly once, in one pass.
 */
import type { FusedLoader } from './wgsl.ts';

const TILE = 16;

/** C(m x n) = A(m x k) * B(k x n), column-major. Bindings: 0=C, 1=A, 2=B. */
export function gemmKernel(m: number, k: number, n: number): string {
  return `
@group(0) @binding(0) var<storage, read_write> C: array<f32>;
@group(0) @binding(1) var<storage, read> A: array<f32>;
@group(0) @binding(2) var<storage, read> B: array<f32>;

var<workgroup> tileA: array<array<f32, ${TILE}>, ${TILE}>;
var<workgroup> tileB: array<array<f32, ${TILE}>, ${TILE}>;

@compute @workgroup_size(${TILE}, ${TILE})
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let row = gid.x;
  let col = gid.y;
  let lx = lid.x;
  let ly = lid.y;
  var acc: f32 = 0.0;
  let numTiles = ${Math.ceil(k / TILE)}u;
  for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
    // Consecutive lx reads consecutive addresses in both loads (column-major).
    let aCol = t * ${TILE}u + ly;
    let bRow = t * ${TILE}u + lx;
    tileA[lx][ly] = select(0.0, A[row + aCol * ${m}u], row < ${m}u && aCol < ${k}u);
    tileB[lx][ly] = select(0.0, B[bRow + col * ${k}u], bRow < ${k}u && col < ${n}u);
    workgroupBarrier();
    for (var p: u32 = 0u; p < ${TILE}u; p = p + 1u) {
      acc = acc + tileA[lx][p] * tileB[p][ly];
    }
    workgroupBarrier();
  }
  if (row < ${m}u && col < ${n}u) {
    C[row + col * ${m}u] = acc;
  }
}
`;
}

export const gemmDispatch = (m: number, n: number): [number, number] => [
  Math.ceil(m / TILE),
  Math.ceil(n / TILE),
];

/** out(n x m) = in(m x n)', column-major, staged through a tile so both the
 *  read and the write are coalesced. Bindings: 0=out, 1=in. */
export function transposeKernel(m: number, n: number): string {
  return `
@group(0) @binding(0) var<storage, read_write> out: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;

var<workgroup> tile: array<array<f32, ${TILE}>, ${TILE}>;

@compute @workgroup_size(${TILE}, ${TILE})
fn main(
  @builtin(workgroup_id) wg: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let lx = lid.x;
  let ly = lid.y;
  // Read block (wg.x, wg.y) of src: rows wg.x*T.., cols wg.y*T..
  let sr = wg.x * ${TILE}u + lx;
  let sc = wg.y * ${TILE}u + ly;
  if (sr < ${m}u && sc < ${n}u) {
    tile[ly][lx] = src[sr + sc * ${m}u];
  }
  workgroupBarrier();
  // Write the transposed block: rows of out are cols of src.
  let dr = wg.y * ${TILE}u + lx;
  let dc = wg.x * ${TILE}u + ly;
  if (dr < ${n}u && dc < ${m}u) {
    out[dr + dc * ${n}u] = tile[lx][ly];
  }
}
`;
}

export const transposeDispatch = (m: number, n: number): [number, number] => [
  Math.ceil(m / TILE),
  Math.ceil(n / TILE),
];

// ── Reductions ──────────────────────────────────────────────────────────

export type Combine = 'add' | 'mul' | 'max' | 'min';
/** Applied to each loaded element before combining. `sq` serves norm/dot. */
export type MapKind = 'id' | 'sq';
/** Applied to the final value. `scale` divides (mean); `sqrt` closes norm. */
export type Epilogue = { kind: 'none' } | { kind: 'scale'; by: number } | { kind: 'sqrt' };

const REDUCE_WG = 256;

const IDENT: Record<Combine, string> = {
  add: '0.0',
  mul: '1.0',
  max: '-3.4028234663852886e+38',
  min: '3.4028234663852886e+38',
};

const comb = (c: Combine, a: string, b: string): string =>
  c === 'add' ? `${a} + ${b}`
  : c === 'mul' ? `${a} * ${b}`
  : `${c}(${a}, ${b})`;

const mapped = (m: MapKind, v: string): string => (m === 'sq' ? `${v} * ${v}` : v);

const epilogued = (e: Epilogue, v: string): string =>
  e.kind === 'scale' ? `(${v}) * ${e.by}` : e.kind === 'sqrt' ? `sqrt(${v})` : v;

const treeReduce = (c: Combine): string => `
  sdata[li] = acc;
  workgroupBarrier();
  var stride = ${REDUCE_WG / 2}u;
  while (stride > 0u) {
    if (li < stride) {
      sdata[li] = ${comb(c, 'sdata[li]', 'sdata[li + stride]')};
    }
    workgroupBarrier();
    stride = stride / 2u;
  }`;

/** Number of pass-1 partials for a full reduction over `count` elements. */
export function reducePartials(count: number): number {
  return Math.max(1, Math.min(1024, Math.ceil(count / (REDUCE_WG * 8))));
}

/**
 * Full reduction, pass 1: `numWg` workgroups grid-stride over `count`
 * elements of the fused loader, leaving one partial each. Loader binding
 * declarations (`decls`) come from wgsl.ts's bindingDecls, whose binding 0
 * ("out") is the partials buffer here.
 */
export function reduceFullPass1(
  decls: string[],
  loader: FusedLoader,
  count: number,
  numWg: number,
  combine: Combine,
  map: MapKind,
): string {
  return `${decls.join('\n')}
${loader.helpers}
var<workgroup> sdata: array<f32, ${REDUCE_WG}>;

@compute @workgroup_size(${REDUCE_WG})
fn main(
  @builtin(workgroup_id) wg: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let li = lid.x;
  var acc: f32 = ${IDENT[combine]};
  var i = wg.x * ${REDUCE_WG}u + li;
  while (i < ${count}u) {
    let v = ${loader.body};
    acc = ${comb(combine, 'acc', mapped(map, 'v'))};
    i = i + ${numWg * REDUCE_WG}u;
  }
${treeReduce(combine)}
  if (li == 0u) { out[wg.x] = sdata[0]; }
}
`;
}

/** Full reduction, pass 2: one workgroup folds the partials and applies the
 *  epilogue. Bindings: 0=result (1 element), 1=partials. */
export function reduceFullPass2(
  numPartials: number,
  combine: Combine,
  epilogue: Epilogue,
): string {
  return `
@group(0) @binding(0) var<storage, read_write> out: array<f32>;
@group(0) @binding(1) var<storage, read> partials: array<f32>;
var<workgroup> sdata: array<f32, ${REDUCE_WG}>;

@compute @workgroup_size(${REDUCE_WG})
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
  let li = lid.x;
  var acc: f32 = ${IDENT[combine]};
  var i = li;
  while (i < ${numPartials}u) {
    acc = ${comb(combine, 'acc', 'partials[i]')};
    i = i + ${REDUCE_WG}u;
  }
${treeReduce(combine)}
  if (li == 0u) { out[0] = ${epilogued(epilogue, 'sdata[0]')}; }
}
`;
}

/**
 * Column-wise reduction of an (m x n) input: one workgroup per column,
 * threads grid-stride down the column (contiguous in column-major), leaving
 * out[column]. Dispatch n workgroups.
 */
export function reduceColumns(
  decls: string[],
  loader: FusedLoader,
  m: number,
  combine: Combine,
  map: MapKind,
  epilogue: Epilogue,
): string {
  return `${decls.join('\n')}
${loader.helpers}
var<workgroup> sdata: array<f32, ${REDUCE_WG}>;

@compute @workgroup_size(${REDUCE_WG})
fn main(
  @builtin(workgroup_id) wg: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let li = lid.x;
  let col = wg.x;
  var acc: f32 = ${IDENT[combine]};
  var r = li;
  while (r < ${m}u) {
    let i = r + col * ${m}u;
    let v = ${loader.body};
    acc = ${comb(combine, 'acc', mapped(map, 'v'))};
    r = r + ${REDUCE_WG}u;
  }
${treeReduce(combine)}
  if (li == 0u) { out[col] = ${epilogued(epilogue, 'sdata[0]')}; }
}
`;
}