/**
 * Drawing the pressure field, straight out of the buffer the solver wrote.
 *
 * There is no readback in the display path: the fragment shader reads the
 * solver's storage buffer directly, so a frame costs one draw call and no
 * GPU->CPU round trip. (The host does read the field back occasionally, but
 * only to decide the colour scale — see the autoscale in main.ts.)
 *
 * Two things are drawn at once. The pressure goes through a diverging
 * colormap about zero, symmetric because the field is signed and a wave is
 * not more interesting on one side of zero than the other. The medium is
 * blended in underneath as a grey wash proportional to how far the local
 * sound speed departs from the background, which shows a hard scatterer as a
 * distinct shape and a smooth one as a soft cloud without either needing its
 * own kind of drawing.
 *
 * Sampling is bilinear rather than nearest, so a 256-point grid on a 700-pixel
 * canvas looks like a wave rather than like a grid.
 */
import type { ColormapFunc } from './colormaps.ts';

const SHADER = `
struct View {
  nx: u32,
  ny: u32,
  scale: f32,
  cref: f32,
  cdev: f32,
  medium: f32,
  micX: f32,
  micY: f32,
  micOn: f32,
};

@group(0) @binding(0) var<uniform> view: View;
@group(0) @binding(1) var<storage, read> field: array<f32>;
@group(0) @binding(2) var<storage, read> speed: array<f32>;
@group(0) @binding(3) var cmap: texture_2d<f32>;
@group(0) @binding(4) var cmapSampler: sampler;

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) uv: vec2f,
};

@vertex
fn vs(@builtin(vertex_index) vi: u32) -> VSOut {
  // One oversized triangle covering the viewport.
  var xy = array<vec2f, 3>(vec2f(-1.0, -3.0), vec2f(-1.0, 1.0), vec2f(3.0, 1.0));
  var out: VSOut;
  let p = xy[vi];
  out.pos = vec4f(p, 0.0, 1.0);
  // uv (0,0) at the top-left corner of the domain.
  out.uv = vec2f((p.x + 1.0) * 0.5, (1.0 - p.y) * 0.5);
  return out;
}

fn idx(ix: i32, iy: i32) -> u32 {
  let cx = clamp(ix, 0, i32(view.nx) - 1);
  let cy = clamp(iy, 0, i32(view.ny) - 1);
  return u32(cx + i32(view.nx) * cy);
}

fn bilinear(gx: f32, gy: f32, isField: bool) -> f32 {
  let x0 = i32(floor(gx));
  let y0 = i32(floor(gy));
  let fx = gx - f32(x0);
  let fy = gy - f32(y0);
  var v00: f32; var v10: f32; var v01: f32; var v11: f32;
  if (isField) {
    v00 = field[idx(x0, y0)];
    v10 = field[idx(x0 + 1, y0)];
    v01 = field[idx(x0, y0 + 1)];
    v11 = field[idx(x0 + 1, y0 + 1)];
  } else {
    v00 = speed[idx(x0, y0)];
    v10 = speed[idx(x0 + 1, y0)];
    v01 = speed[idx(x0, y0 + 1)];
    v11 = speed[idx(x0 + 1, y0 + 1)];
  }
  return mix(mix(v00, v10, fx), mix(v01, v11, fx), fy);
}

@fragment
fn fs(in: VSOut) -> @location(0) vec4f {
  // Row 0 of the buffer is the bottom of the picture: y increases upward in
  // the grid, downward on the screen.
  let gx = in.uv.x * f32(view.nx) - 0.5;
  let gy = (1.0 - in.uv.y) * f32(view.ny) - 0.5;

  let p = bilinear(gx, gy, true);
  let t = clamp(0.5 + 0.5 * p / max(view.scale, 1e-20), 0.0, 1.0);
  var rgb = textureSample(cmap, cmapSampler, vec2f(t, 0.5)).rgb;

  if (view.medium > 0.0) {
    let c = bilinear(gx, gy, false);
    let m = clamp(abs(c - view.cref) / max(view.cdev, 1e-20), 0.0, 1.0);
    rgb = mix(rgb, vec3f(0.35, 0.35, 0.38), view.medium * m);
  }

  // The microphone, as a ring sized in grid cells so it stays the same size
  // on screen whatever the canvas is scaled to.
  if (view.micOn > 0.0) {
    let d = length(vec2f(gx - view.micX, gy - view.micY));
    let r = 0.022 * f32(view.nx);
    let w = 0.005 * f32(view.nx);
    if (abs(d - r) < w) {
      rgb = mix(rgb, vec3f(1.0, 1.0, 1.0), 0.9);
    } else if (d < 0.006 * f32(view.nx)) {
      rgb = mix(rgb, vec3f(1.0, 1.0, 1.0), 0.9);
    }
  }
  return vec4f(rgb, 1.0);
}
`;

export interface FieldViewOptions {
  device: GPUDevice;
  canvas: HTMLCanvasElement;
  nx: number;
  ny: number;
}

export class FieldView {
  readonly canvas: HTMLCanvasElement;

  #device: GPUDevice;
  #context: GPUCanvasContext;
  #pipeline: GPURenderPipeline;
  #layout: GPUBindGroupLayout;
  #uniform: GPUBuffer;
  #uniformData = new ArrayBuffer(48);
  #sampler: GPUSampler;
  #cmapTexture: GPUTexture;
  #bindGroup: GPUBindGroup | null = null;
  #nx: number;
  #ny: number;

  constructor(opts: FieldViewOptions) {
    const { device, canvas } = opts;
    this.#device = device;
    this.canvas = canvas;
    this.#nx = opts.nx;
    this.#ny = opts.ny;

    const context = canvas.getContext('webgpu');
    if (!context) throw new Error('this canvas has no WebGPU context');
    this.#context = context;
    const format = navigator.gpu.getPreferredCanvasFormat();
    context.configure({ device, format, alphaMode: 'opaque' });

    this.#layout = device.createBindGroupLayout({
      label: 'field-view',
      entries: [
        { binding: 0, visibility: GPUShaderStage.FRAGMENT, buffer: { type: 'uniform' } },
        {
          binding: 1,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: { type: 'read-only-storage' },
        },
        {
          binding: 2,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: { type: 'read-only-storage' },
        },
        { binding: 3, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'float' } },
        { binding: 4, visibility: GPUShaderStage.FRAGMENT, sampler: { type: 'filtering' } },
      ],
    });

    const module = device.createShaderModule({ code: SHADER, label: 'field-view' });
    this.#pipeline = device.createRenderPipeline({
      label: 'field-view',
      layout: device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
      vertex: { module, entryPoint: 'vs' },
      fragment: { module, entryPoint: 'fs', targets: [{ format }] },
      primitive: { topology: 'triangle-list' },
    });

    this.#uniform = device.createBuffer({
      label: 'field-view-uniform',
      size: this.#uniformData.byteLength,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    this.#sampler = device.createSampler({ magFilter: 'linear', minFilter: 'linear' });
    this.#cmapTexture = device.createTexture({
      label: 'colormap',
      size: [256, 1],
      format: 'rgba8unorm',
      usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
    });
  }

  /** Point the view at the buffers of a (new) simulation. */
  setSource(pressure: GPUBuffer, speed: GPUBuffer, nx: number, ny: number): void {
    this.#nx = nx;
    this.#ny = ny;
    this.#bindGroup = this.#device.createBindGroup({
      layout: this.#layout,
      entries: [
        { binding: 0, resource: { buffer: this.#uniform } },
        { binding: 1, resource: { buffer: pressure } },
        { binding: 2, resource: { buffer: speed } },
        { binding: 3, resource: this.#cmapTexture.createView() },
        { binding: 4, resource: this.#sampler },
      ],
    });
  }

  setColormap(cmap: ColormapFunc): void {
    const data = new Uint8Array(256 * 4);
    for (let i = 0; i < 256; i++) {
      const [r, g, b] = cmap(i / 255);
      data[4 * i] = r;
      data[4 * i + 1] = g;
      data[4 * i + 2] = b;
      data[4 * i + 3] = 255;
    }
    this.#device.queue.writeTexture(
      { texture: this.#cmapTexture },
      data,
      { bytesPerRow: 256 * 4 },
      { width: 256, height: 1 },
    );
  }

  /**
   * Draw one frame. `scale` is the pressure the colormap saturates at, in
   * both directions; `cref`/`cdev` describe the medium wash, and `medium` is
   * how strongly to apply it (0 turns it off).
   */
  draw(opts: {
    scale: number;
    cref: number;
    cdev: number;
    medium: number;
    /** Microphone position, in grid-index coordinates. */
    mic?: { ix: number; iy: number } | null;
  }): void {
    if (!this.#bindGroup) return;
    const u32 = new Uint32Array(this.#uniformData);
    const f32 = new Float32Array(this.#uniformData);
    u32[0] = this.#nx;
    u32[1] = this.#ny;
    f32[2] = opts.scale;
    f32[3] = opts.cref;
    f32[4] = opts.cdev;
    f32[5] = opts.medium;
    f32[6] = opts.mic?.ix ?? 0;
    f32[7] = opts.mic?.iy ?? 0;
    f32[8] = opts.mic ? 1 : 0;
    this.#device.queue.writeBuffer(this.#uniform, 0, this.#uniformData);

    const encoder = this.#device.createCommandEncoder({ label: 'field-view' });
    const pass = encoder.beginRenderPass({
      colorAttachments: [
        {
          view: this.#context.getCurrentTexture().createView(),
          clearValue: { r: 0, g: 0, b: 0, a: 1 },
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    });
    pass.setPipeline(this.#pipeline);
    pass.setBindGroup(0, this.#bindGroup);
    pass.draw(3);
    pass.end();
    this.#device.queue.submit([encoder.finish()]);
  }

  /** Match the canvas's backing store to its CSS size. */
  resize(): void {
    const dpr = Math.min(window.devicePixelRatio || 1, 2);
    const rect = this.canvas.getBoundingClientRect();
    const w = Math.max(1, Math.round(rect.width * dpr));
    const h = Math.max(1, Math.round(rect.height * dpr));
    if (this.canvas.width !== w || this.canvas.height !== h) {
      this.canvas.width = w;
      this.canvas.height = h;
    }
  }

  destroy(): void {
    this.#uniform.destroy();
    this.#cmapTexture.destroy();
  }
}