/**
 * 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-to-CPU round trip. (The host does read a reduction back occasionally, to
 * decide the colour scale.)
 *
 * The picture is a ray march. Each pixel casts one ray, intersects it with the
 * cube, and steps along it accumulating emission front to back: the pressure
 * goes through a diverging colormap about zero, and the opacity goes as a
 * power of |p|, so quiet regions are transparent and the wavefronts are what
 * you see. The medium is added as a grey emission 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.
 *
 * Two honest limitations. Sampling is nearest-neighbour, not trilinear: the
 * field lives in a storage buffer rather than a filterable 3D texture, so
 * trilinear would be eight fetches per sample and the march takes tens of
 * millions of samples a frame. With the ray step set near the cell size the
 * difference is visible mainly as a faint stippling on strong wavefronts. And
 * the compositing is emission only, with no lighting and no shadowing, so what
 * is behind a strong feature is dimmed but never occluded correctly.
 *
 * A clip plane on x is provided because a volume render of a wavefield is
 * mostly the outside of a wavefield. Pulling the clip in is how you see the
 * interior, and it is the closest thing here to the flat sibling's picture.
 */
import type { ColormapFunc } from './colormaps.ts';

const SHADER = `
struct View {
  eye: vec4f,    // xyz: eye position, w: tan(fov/2)
  right: vec4f,  // xyz: camera right,  w: aspect ratio
  up: vec4f,     // xyz: camera up,     w: domain side L
  fwd: vec4f,    // xyz: camera forward, w: pressure the colormap saturates at
  bg: vec4f,     // rgb: background
  a: vec4f,      // opacity, medium strength, reference speed, speed spread
  b: vec4f,      // ray steps, clip plane x, grid side n, contrast exponent
  m: vec4f,      // xyz: microphone position, w: whether to draw it
};

@group(0) @binding(0) var<uniform> V: View;
@group(0) @binding(1) var<storage, read> p: array<f32>;
@group(0) @binding(2) var<storage, read> cs: array<f32>;
@group(0) @binding(3) var cmap: texture_2d<f32>;
@group(0) @binding(4) var samp: sampler;

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) ndc: 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 q = xy[vi];
  out.pos = vec4f(q, 0.0, 1.0);
  out.ndc = q;
  return out;
}

fn voxel(q: vec3f) -> u32 {
  let n = i32(V.b.z);
  let h = V.up.w / f32(n);
  let g = clamp(vec3i(floor((q + vec3f(0.5 * V.up.w)) / h)), vec3i(0), vec3i(n - 1));
  return u32(g.x + n * g.y + n * n * g.z);
}

// A point on the cube's surface is on an edge when two of its three distances
// to the bounding planes vanish, so the test is on the median of the three.
fn edge(q: vec3f, half: f32, L: f32) -> f32 {
  let d = abs(abs(q) - vec3f(half));
  let lo = min(d.x, min(d.y, d.z));
  let hi = max(d.x, max(d.y, d.z));
  let mid = d.x + d.y + d.z - lo - hi;
  let w = 0.004 * L;
  return 1.0 - smoothstep(w, 2.0 * w, mid);
}

@fragment
fn fs(in: VSOut) -> @location(0) vec4f {
  let L = V.up.w;
  let half = 0.5 * L;
  let eye = V.eye.xyz;
  let dir = normalize(V.fwd.xyz
      + in.ndc.x * V.right.w * V.eye.w * V.right.xyz
      + in.ndc.y * V.eye.w * V.up.xyz);

  let inv = 1.0 / dir;
  let ta = (vec3f(-half) - eye) * inv;
  let tb = (vec3f(half) - eye) * inv;
  let lo = min(ta, tb);
  let hi = max(ta, tb);
  let t1 = min(min(hi.x, hi.y), hi.z);
  var t0 = max(max(lo.x, lo.y), lo.z);

  var col = V.bg.rgb;
  if (t1 <= max(t0, 0.0)) { return vec4f(col, 1.0); }
  t0 = max(t0, 0.0);

  let lineCol = vec3f(0.42, 0.47, 0.55);
  col = mix(col, lineCol, 0.5 * edge(eye + dir * t1, half, L));

  let steps = i32(V.b.x);
  let dl = (t1 - t0) / f32(steps);
  // Opacity is quoted per cell, so a longer ray step is proportionally more
  // opaque and the picture does not change brightness with the quality knob.
  let unit = dl / (L / V.b.z);
  let scale = max(V.fwd.w, 1e-20);
  let grey = vec3f(0.55, 0.58, 0.63);

  var acc = vec3f(0.0);
  var alpha = 0.0;
  for (var k = 0; k < steps; k = k + 1) {
    if (alpha > 0.995) { break; }
    let q = eye + dir * (t0 + (f32(k) + 0.5) * dl);
    if (q.x > V.b.y) { continue; }
    let i = voxel(q);
    let v = clamp(p[i] / scale, -1.0, 1.0);
    var a = pow(abs(v), V.b.w) * V.a.x * unit;
    var rgb = textureSampleLevel(cmap, samp, vec2f(0.5 + 0.5 * v, 0.5), 0.0).rgb;
    if (V.a.y > 0.0) {
      let m = clamp(abs(cs[i] - V.a.z) / max(V.a.w, 1e-20), 0.0, 1.0);
      let am = m * V.a.y * unit;
      let tot = a + am;
      if (tot > 1e-12) { rgb = (a * rgb + am * grey) / tot; }
      a = tot;
    }
    a = clamp(a, 0.0, 1.0);
    acc = acc + (1.0 - alpha) * a * rgb;
    alpha = alpha + (1.0 - alpha) * a;
  }
  col = acc + (1.0 - alpha) * col;
  col = mix(col, lineCol, 0.75 * edge(eye + dir * t0, half, L));

  // The microphone, a dot in a ring at the ray's closest approach to it.
  // Drawn on top rather than composited into the march, like the 2D app's
  // white ring: a marker, not a thing in the scene.
  if (V.m.w > 0.0) {
    let toM = V.m.xyz - eye;
    let tm = dot(toM, dir);
    if (tm > 0.0) {
      let d = length(toM - tm * dir) / L;
      let dotm = 1.0 - smoothstep(0.004, 0.007, d);
      let ring = 1.0 - smoothstep(0.0015, 0.004, abs(d - 0.016));
      col = mix(col, vec3f(1.0), 0.9 * max(dotm, ring));
    }
  }
  return vec4f(col, 1.0);
}
`;

/** Where the camera is looking from. Angles in radians, distance in metres. */
export interface Camera {
  az: number;
  el: number;
  dist: number;
}

export interface DrawOptions {
  camera: Camera;
  /** Pressure the colormap saturates at, in both directions. */
  scale: number;
  opacity: number;
  contrast: number;
  /** Samples along each ray. */
  steps: number;
  /** Everything with x above this is not drawn, in metres. */
  clipX: number;
  /** Strength of the medium wash, 0 to turn it off. */
  medium: number;
  cref: number;
  cdev: number;
  /** Microphone position in metres, or null to not draw it. */
  mic?: { x: number; y: number; z: number } | null;
}

const TAN_HALF_FOV = Math.tan((32 * Math.PI) / 360);

const cross = (a: number[], b: number[]) => [
  a[1] * b[2] - a[2] * b[1],
  a[2] * b[0] - a[0] * b[2],
  a[0] * b[1] - a[1] * b[0],
];
const norm = (a: number[]) => {
  const m = Math.hypot(a[0], a[1], a[2]) || 1;
  return [a[0] / m, a[1] / m, a[2] / m];
};

export class VolumeView {
  readonly canvas: HTMLCanvasElement;

  #device: GPUDevice;
  #context: GPUCanvasContext;
  #pipeline: GPURenderPipeline;
  #layout: GPUBindGroupLayout;
  #uniform: GPUBuffer;
  #host = new ArrayBuffer(8 * 16);
  #sampler: GPUSampler;
  #cmapTexture: GPUTexture;
  #bindGroups: GPUBindGroup[] = [];
  #n = 0;
  #L = 1;
  /** Background, matched to the CSS so the cube sits on the page. */
  bg: [number, number, number] = [0.043, 0.055, 0.071];

  constructor(device: GPUDevice, canvas: HTMLCanvasElement) {
    this.#device = device;
    this.canvas = canvas;

    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: 'volume-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: 'volume-view' });
    this.#pipeline = device.createRenderPipeline({
      label: 'volume-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: 'volume-view-uniform',
      size: this.#host.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. Two pressure buffers
   * are given because the solver alternates between them; `draw` is told which
   * one is current.
   */
  setSource(pressures: [GPUBuffer, GPUBuffer], speed: GPUBuffer, n: number, L: number): void {
    this.#n = n;
    this.#L = L;
    this.#bindGroups = pressures.map((p) =>
      this.#device.createBindGroup({
        layout: this.#layout,
        entries: [
          { binding: 0, resource: { buffer: this.#uniform } },
          { binding: 1, resource: { buffer: p } },
          { 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 from the pressure buffer with the given index. */
  draw(which: number, o: DrawOptions): void {
    const bg = this.#bindGroups[which];
    if (!bg) return;

    const { az, el, dist } = o.camera;
    const ce = Math.cos(el);
    const toEye = [ce * Math.cos(az), ce * Math.sin(az), Math.sin(el)];
    const eye = toEye.map((v) => v * dist);
    const fwd = toEye.map((v) => -v);
    const right = norm(cross([0, 0, 1], fwd));
    const up = cross(fwd, right);

    const rect = this.canvas.width / Math.max(this.canvas.height, 1);
    const f = new Float32Array(this.#host);
    f.set([eye[0], eye[1], eye[2], TAN_HALF_FOV], 0);
    f.set([right[0], right[1], right[2], rect], 4);
    f.set([up[0], up[1], up[2], this.#L], 8);
    f.set([fwd[0], fwd[1], fwd[2], o.scale], 12);
    f.set([this.bg[0], this.bg[1], this.bg[2], 1], 16);
    f.set([o.opacity, o.medium, o.cref, o.cdev], 20);
    f.set([o.steps, o.clipX, this.#n, o.contrast], 24);
    f.set([o.mic?.x ?? 0, o.mic?.y ?? 0, o.mic?.z ?? 0, o.mic ? 1 : 0], 28);
    this.#device.queue.writeBuffer(this.#uniform, 0, this.#host);

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

  /**
   * Match the canvas's backing store to its CSS size. Device pixel ratio is
   * ignored: the march is the whole cost of a frame and it scales with pixels,
   * so a retina display would pay four times over for a picture that is
   * already smooth.
   */
  resize(): void {
    const rect = this.canvas.getBoundingClientRect();
    const w = Math.max(1, Math.round(rect.width));
    const h = Math.max(1, Math.round(rect.height));
    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();
  }
}