/**
 * The external operations, as GPU dispatches — instantiated per object.
 *
 * Everything a compiled model does is element-wise, so these ops are the only
 * places where a point reads anything but itself: the difference stencils on
 * each object's own grid, the wall-masked Laplacian in the air, and the
 * couplings that carry an object's motion into the air. Where dulcimer had a
 * fixed set of five ops over exactly two grids, here the scene builder
 * instantiates one op per (operation, object) pair — `dxx_s1` is the second
 * difference on string s1's grid, `rad_m1` samples membrane m1 at each air
 * point — each with its own geometry baked into its WGSL. The generated
 * model .m calls them by these names.
 *
 * Layouts live only here and in the renderer; a .m never sees them. An air
 * field is an npts x 1 column vector laid out x-fastest — the point
 * (ix, iy, iz) is element `ix + nx*(iy + ny*iz)` — and an object field is n1
 * (times n2) nodes over its local 0..L1 (x 0..L2) coordinates, first axis
 * fastest. In world coordinates every object is centred (see src/grids.ts),
 * which is what the coupling ops assume; where an object sits in y and z is
 * carried by the scene-built profile fields instead, so it costs no
 * recompile to move one.
 *
 * `lapw` is the wall-masked Laplacian, and it is what makes a body's shell
 * rigid. Each of the six face terms is scaled by the mask at the neighbour it
 * reads, so a face into the wall contributes nothing — the discrete Neumann
 * (zero normal velocity, sound-hard) condition, in the divergence form
 * div(w grad p). The alternative, a wall as a fast material, would cut the
 * global timestep several-fold; a mask costs nothing. Outside the domain the
 * field is taken to be zero (a soft outer boundary); the scene's absorbing
 * layer is meant to have swallowed the wave before it matters.
 */
import { UnsupportedOnGpu, WORKGROUP_SIZE } from './wgsl.ts';
import type { AirGrid, ObjectGrid } from '../grids.ts';

export type OpKind = 'dxx' | 'dxxxx' | 'lap2' | 'lapw' | 'radline' | 'radrect' | 'bridge';

/** One instantiated operation: its name in the .m, its shape contract, and
 *  its WGSL. */
export interface OpInstance {
  name: string;
  kind: OpKind;
  /** Element counts of the arguments, in order. */
  argCounts: number[];
  outCount: number;
  code: string;
}

const lineAt = (n: number): string => `
fn at(j: i32) -> f32 {
  if (j < 0 || j >= ${n}) { return 0.0; }
  return src[u32(j)];
}
`;

const header = (args: number): string =>
  [
    `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;`,
    `@group(0) @binding(1) var<storage, read> src: array<f32>;`,
    ...(args > 1
      ? [`@group(0) @binding(2) var<storage, read> src2: array<f32>;`]
      : []),
  ].join('\n');

/** dxx_<id>: second difference along a line grid, zero beyond the ends. */
export function makeDxx(name: string, g: ObjectGrid): OpInstance {
  const n = g.n1;
  const code = `${header(1)}
${lineAt(n)}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${n}u) { return; }
  let j = i32(i);
  dst[i] = ${1 / (g.h1 * g.h1)}f * (at(j - 1) - 2.0 * at(j) + at(j + 1));
}
`;
  return { name, kind: 'dxx', argCounts: [n], outCount: n, code };
}

/**
 * dxxxx_<id>: fourth difference along a line grid, with *antisymmetric*
 * ghost values (u(-k) = -u(k), and mirrored at the far end). Together with
 * the pinned endpoints this is the simply-supported end condition (u = 0 and
 * u'' = 0), and it makes the discrete operator exactly the square of the
 * pinned second difference — so the modes are sines and a bar's partials
 * land where the dispersion relation f_n = kap·lam_n/2π says, which is what
 * the spectrum overlay draws. Zero ghosts here would quietly clamp the ends
 * instead, and the pitch knob would lie by a factor approaching two.
 */
export function makeDxxxx(name: string, g: ObjectGrid): OpInstance {
  const n = g.n1;
  const inv4 = 1 / (g.h1 * g.h1 * g.h1 * g.h1);
  const code = `${header(1)}
fn at(j: i32) -> f32 {
  var k = j;
  var s = 1.0;
  if (k < 0) { k = -k; s = -1.0; }
  if (k > ${n - 1}) { k = ${2 * (n - 1)} - k; s = -s; }
  if (k < 0 || k > ${n - 1}) { return 0.0; }
  return s * src[u32(k)];
}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${n}u) { return; }
  let j = i32(i);
  dst[i] = ${inv4}f * (at(j - 2) - 4.0 * at(j - 1) + 6.0 * at(j) - 4.0 * at(j + 1) + at(j + 2));
}
`;
  return { name, kind: 'dxxxx', argCounts: [n], outCount: n, code };
}

/**
 * lap_<id>: 5-point Laplacian on a rectangular grid, zero beyond the rim.
 * The rim itself is pinned by the object's `pin` mask, so reading zero
 * outside is consistent with the boundary the scheme imposes.
 */
export function makeLap2(name: string, g: ObjectGrid): OpInstance {
  const { n1, n2, npts, h1, h2 } = g;
  const code = `${header(1)}
fn at(i: i32, j: i32) -> f32 {
  if (i < 0 || i >= ${n1} || j < 0 || j >= ${n2}) { return 0.0; }
  return src[u32(i + ${n1} * j)];
}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let k = gid.x;
  if (k >= ${npts}u) { return; }
  let i = i32(k % ${n1}u);
  let j = i32(k / ${n1}u);
  let c = at(i, j);
  dst[k] = ${1 / (h1 * h1)}f * (at(i - 1, j) - 2.0 * c + at(i + 1, j))
         + ${1 / (h2 * h2)}f * (at(i, j - 1) - 2.0 * c + at(i, j + 1));
}
`;
  return { name, kind: 'lap2', argCounts: [npts], outCount: npts, code };
}

/** lapw: the wall-masked 7-point Laplacian in the air. One instance. */
export function makeLapw(name: string, air: AirGrid): OpInstance {
  const { nx, ny, nz, npts, h } = air;
  const code = `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;
@group(0) @binding(2) var<storage, read> wall: array<f32>;

fn idx(ix: i32, iy: i32, iz: i32) -> i32 {
  return ix + ${nx} * (iy + ${ny} * iz);
}
fn inside(ix: i32, iy: i32, iz: i32) -> bool {
  return ix >= 0 && ix < ${nx} && iy >= 0 && iy < ${ny} && iz >= 0 && iz < ${nz};
}
// One face's contribution: masked by the wall at the neighbour, so a face
// into the wall carries no flux — the Neumann (rigid) condition.
fn face(ix: i32, iy: i32, iz: i32, pc: f32) -> f32 {
  if (!inside(ix, iy, iz)) { return -pc; } // outside the domain: p = 0, open
  let k = u32(idx(ix, iy, iz));
  return wall[k] * (src[k] - pc);
}

@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${npts}u) { return; }
  let ix = i32(i % ${nx}u);
  let iy = i32((i / ${nx}u) % ${ny}u);
  let iz = i32(i / ${nx * ny}u);
  let pc = src[i];
  let s = face(ix - 1, iy, iz, pc) + face(ix + 1, iy, iz, pc)
        + face(ix, iy - 1, iz, pc) + face(ix, iy + 1, iz, pc)
        + face(ix, iy, iz - 1, pc) + face(ix, iy, iz + 1, pc);
  dst[i] = ${1 / (h * h)}f * wall[i] * s;
}
`;
  return { name, kind: 'lapw', argCounts: [npts, npts], outCount: npts, code };
}

/**
 * rad_<id> for a line object: give each air point the object's value at that
 * point's own x (linearly interpolated, zero beyond the ends). Multiplied by
 * a scene-built profile, that is the object radiating from its line.
 */
export function makeRadLine(name: string, g: ObjectGrid, air: AirGrid): OpInstance {
  const n = g.n1;
  const code = `${header(1)}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${air.npts}u) { return; }
  let ix = i % ${air.nx}u;
  let xw = ${-air.Lx / 2}f + (f32(ix) + 0.5) * ${air.h}f;
  let u = (xw - ${-g.L1 / 2}f) / ${g.h1}f;
  let j = i32(floor(u));
  if (j < 0 || j >= ${n - 1}) { dst[i] = 0.0; return; }
  let f = u - f32(j);
  dst[i] = mix(src[u32(j)], src[u32(j) + 1u], f);
}
`;
  return { name, kind: 'radline', argCounts: [n], outCount: air.npts, code };
}

/**
 * rad_<id> for a rectangular object: give each air point the object's value
 * at that point's own (x, y), bilinearly interpolated, zero beyond the rim.
 * Multiplied by a scene-built sheet profile, that is the membrane radiating
 * from its plane.
 */
export function makeRadRect(name: string, g: ObjectGrid, air: AirGrid): OpInstance {
  const { n1, n2 } = g;
  const code = `${header(1)}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${air.npts}u) { return; }
  let ix = i % ${air.nx}u;
  let iy = (i / ${air.nx}u) % ${air.ny}u;
  let xw = ${-air.Lx / 2}f + (f32(ix) + 0.5) * ${air.h}f;
  let yw = ${-air.Ly / 2}f + (f32(iy) + 0.5) * ${air.h}f;
  let u = (xw - ${-g.L1 / 2}f) / ${g.h1}f;
  let v = (yw - ${-g.L2 / 2}f) / ${g.h2}f;
  let ju = i32(floor(u));
  let jv = i32(floor(v));
  if (ju < 0 || ju >= ${n1 - 1} || jv < 0 || jv >= ${n2 - 1}) { dst[i] = 0.0; return; }
  let fu = u - f32(ju);
  let fv = v - f32(jv);
  let a = mix(src[u32(ju + ${n1} * jv)], src[u32(ju + 1 + ${n1} * jv)], fu);
  let b = mix(src[u32(ju + ${n1} * (jv + 1))], src[u32(ju + 1 + ${n1} * (jv + 1))], fu);
  dst[i] = mix(a, b, fv);
}
`;
  return { name, kind: 'radrect', argCounts: [g.npts], outCount: air.npts, code };
}

/**
 * brg_<id>: du/dx at the last node of a line object, broadcast to every air
 * point. The end is pinned, so this is the arriving slope — what the tension
 * pulls the bridge with. Multiplied by a scene-built patch profile, that is
 * the bridge force driving the top plate.
 */
export function makeBridge(name: string, g: ObjectGrid, air: AirGrid): OpInstance {
  const n = g.n1;
  const code = `${header(1)}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${air.npts}u) { return; }
  dst[i] = (src[${n - 1}u] - src[${n - 2}u]) * ${1 / g.h1}f;
}
`;
  return { name, kind: 'bridge', argCounts: [n], outCount: air.npts, code };
}

/** Compiled op pipelines for one scene's set of instances. Shared by every
 *  plan compiled against that scene. */
export class OpPlan {
  #device: GPUDevice;
  #instances = new Map<string, OpInstance>();
  #layouts = new Map<number, GPUBindGroupLayout>();
  #pipelines = new Map<string, GPUComputePipeline>();

  constructor(device: GPUDevice, instances: OpInstance[]) {
    this.#device = device;
    for (const inst of instances) {
      if (this.#instances.has(inst.name)) {
        throw new UnsupportedOnGpu(`duplicate external op '${inst.name}'`);
      }
      this.#instances.set(inst.name, inst);
    }
  }

  /** Every instance, for the type-rule files and the editor's highlighter. */
  get instances(): OpInstance[] {
    return [...this.#instances.values()];
  }

  has(name: string): boolean {
    return this.#instances.has(name);
  }

  /** The op's shape contract, for the planner's checks. */
  spec(name: string): { argCounts: number[]; outCount: number } {
    const inst = this.#instances.get(name);
    if (!inst) throw new UnsupportedOnGpu(`unknown external op '${name}'`);
    return { argCounts: inst.argCounts, outCount: inst.outCount };
  }

  #layout(args: number): GPUBindGroupLayout {
    const existing = this.#layouts.get(args);
    if (existing) return existing;
    const entries: GPUBindGroupLayoutEntry[] = [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
    ];
    for (let i = 0; i < args; i++) {
      entries.push({
        binding: i + 1,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: 'read-only-storage' },
      });
    }
    const layout = this.#device.createBindGroupLayout({ label: `op-${args}`, entries });
    this.#layouts.set(args, layout);
    return layout;
  }

  /** The pipeline for one op, compiled on first use. */
  async pipeline(name: string): Promise<{ pipeline: GPUComputePipeline; outCount: number }> {
    const inst = this.#instances.get(name);
    if (!inst) throw new UnsupportedOnGpu(`unknown external op '${name}'`);
    const existing = this.#pipelines.get(name);
    if (existing) return { pipeline: existing, outCount: inst.outCount };
    const module = this.#device.createShaderModule({ code: inst.code, label: name });
    let pipeline: GPUComputePipeline;
    try {
      pipeline = await this.#device.createComputePipelineAsync({
        layout: this.#device.createPipelineLayout({
          bindGroupLayouts: [this.#layout(inst.argCounts.length)],
        }),
        compute: { module, entryPoint: 'main' },
        label: name,
      });
    } catch (e) {
      // No error scope here either; see makePipeline in plan.ts.
      throw new UnsupportedOnGpu(
        `op '${name}': ${e instanceof Error ? e.message : String(e)}`,
      );
    }
    this.#pipelines.set(name, pipeline);
    return { pipeline, outCount: inst.outCount };
  }

  createBinding(name: string, srcs: GPUBuffer[], dst: GPUBuffer): GPUBindGroup {
    const entries: GPUBindGroupEntry[] = [{ binding: 0, resource: { buffer: dst } }];
    srcs.forEach((s, i) => entries.push({ binding: i + 1, resource: { buffer: s } }));
    return this.#device.createBindGroup({ layout: this.#layout(srcs.length), entries });
  }

  workgroups(outCount: number): number {
    return Math.ceil(outCount / WORKGROUP_SIZE);
  }
}