/**
 * The external operations, as GPU dispatches.
 *
 * 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 a string field is ns
 * nodes from x = 0 to x = Ls. Those layouts, and the placement of the string
 * inside the air's coordinates, matter only here and to the renderer; a .m
 * never sees them, which is the reason these are host-provided operations
 * rather than something a model expresses with array slicing.
 *
 * `lapw` is the wall-masked Laplacian, and it is what makes the 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 — which is exactly the
 * discrete Neumann (zero normal velocity, sound-hard) condition, in the
 * divergence form div(w grad p). The alternative the flat siblings use, a
 * wall as a fast material, is unusable here: wood at ~4000 m/s would cut the
 * global timestep twelve-fold, where 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.
 *
 * `spread` and `bridge` are the two couplings from the string into the air.
 * `spread` gives each air point the string value at its own x (linearly
 * interpolated, zero beyond the ends): multiplied by a scene-built line
 * profile, that is the string radiating directly. `bridge` broadcasts the
 * string's slope at its bridge end to every air point: multiplied by a
 * scene-built patch profile, that is the bridge force driving the top plate.
 * Both read tiny buffers and write big ones, so each is one cheap dispatch.
 */
import { UnsupportedOnGpu, WORKGROUP_SIZE } from './wgsl.ts';
import { EXTERNAL_OPS } from './externals.ts';

export type OpKind = 'dxx' | 'dxxxx' | 'lapw' | 'spread' | 'bridge';

export interface OpGeometry {
  /** Air grid. */
  nx: number;
  ny: number;
  nz: number;
  h: number;
  /** String grid. */
  ns: number;
  hs: number;
  /** World x of the string's first node (the nut); the bridge is the last. */
  xs0: number;
  /** World x extent of the air domain, for mapping voxel index to metres. */
  Lx: number;
}

function opWGSL(kind: OpKind, g: OpGeometry): { code: string; outCount: number; args: number } {
  const { nx, ny, nz, h, ns, hs } = g;
  const npts = nx * ny * nz;
  const inv2 = 1 / (hs * hs);
  const inv4 = 1 / (hs * hs * hs * hs);

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

  switch (kind) {
    case 'dxx':
      return {
        args: 1,
        outCount: ns,
        code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;
${stringAt}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${ns}u) { return; }
  let j = i32(i);
  dst[i] = ${inv2}f * (at(j - 1) - 2.0 * at(j) + at(j + 1));
}
`,
      };

    case 'dxxxx':
      return {
        args: 1,
        outCount: ns,
        code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;
${stringAt}
@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${ns}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));
}
`,
      };

    case 'lapw':
      return {
        args: 2,
        outCount: npts,
        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;
}
`,
      };

    case 'spread':
      return {
        args: 1,
        outCount: npts,
        code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;

@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 = i % ${nx}u;
  let xw = ${-g.Lx / 2}f + (f32(ix) + 0.5) * ${h}f;
  let u = (xw - ${g.xs0}f) / ${hs}f;
  let j = i32(floor(u));
  if (j < 0 || j >= ${ns - 1}) { dst[i] = 0.0; return; }
  let f = u - f32(j);
  dst[i] = mix(src[u32(j)], src[u32(j) + 1u], f);
}
`,
      };

    case 'bridge':
      return {
        args: 1,
        outCount: npts,
        code: `@group(0) @binding(0) var<storage, read_write> dst: array<f32>;
@group(0) @binding(1) var<storage, read> src: array<f32>;

@compute @workgroup_size(${WORKGROUP_SIZE})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= ${npts}u) { return; }
  // du/dx at the last node. The end is pinned, so this is the string's
  // arriving slope — what the tension pulls the bridge with.
  dst[i] = (src[${ns - 1}u] - src[${ns - 2}u]) * ${1 / hs}f;
}
`,
      };
  }
}

/** Compiled op pipelines for one pair of grids. Shared by every plan. */
export class OpPlan {
  readonly geometry: OpGeometry;

  #device: GPUDevice;
  #layouts = new Map<number, GPUBindGroupLayout>();
  #pipelines = new Map<OpKind, { pipeline: GPUComputePipeline; outCount: number; args: number }>();

  constructor(device: GPUDevice, geometry: OpGeometry) {
    this.#device = device;
    this.geometry = geometry;
  }

  /** The op's shape contract, for the planner's checks. */
  spec(kind: OpKind): { argCounts: number[]; outCount: number } {
    const s = EXTERNAL_OPS.get(kind);
    if (!s) throw new UnsupportedOnGpu(`unknown external op '${kind}'`);
    const { nx, ny, nz, ns } = this.geometry;
    const size = (k: 'air' | 'string'): number => (k === 'air' ? nx * ny * nz : ns);
    return { argCounts: s.args.map(size), outCount: size(s.out) };
  }

  #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(kind: OpKind): Promise<{ pipeline: GPUComputePipeline; outCount: number; args: number }> {
    const existing = this.#pipelines.get(kind);
    if (existing) return existing;
    const { code, outCount, args } = opWGSL(kind, this.geometry);
    const module = this.#device.createShaderModule({ code, label: kind });
    let pipeline: GPUComputePipeline;
    try {
      pipeline = await this.#device.createComputePipelineAsync({
        layout: this.#device.createPipelineLayout({ bindGroupLayouts: [this.#layout(args)] }),
        compute: { module, entryPoint: 'main' },
        label: kind,
      });
    } catch (e) {
      // No error scope here either; see makePipeline in plan.ts.
      throw new UnsupportedOnGpu(
        `op '${kind}': ${e instanceof Error ? e.message : String(e)}`,
      );
    }
    const built = { pipeline, outCount, args };
    this.#pipelines.set(kind, built);
    return built;
  }

  createBinding(kind: OpKind, 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);
  }
}