/**
 * 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. Adapted from acoustic-scattering-3d's renderer; the
 * differences are a rectangular domain instead of a cube, and the body drawn
 * from the wall mask instead of from a speed contrast.
 *
 * The picture is a ray march. Each pixel casts one ray, intersects it with
 * the domain box, 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 body is added as a grey emission where
 * the wall mask says solid, which shows the box and its sound hole without
 * 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 y is provided because a volume render of a wavefield is
 * mostly the outside of a wavefield. Pulling the clip in cuts the picture
 * lengthwise through the string and the cavity, which is the view that
 * actually shows the instrument working.
 */
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: pressure the colormap saturates at
  fwd: vec4f,    // xyz: camera forward, w: contrast exponent
  bg: vec4f,     // rgb: background, w: body strength
  dims: vec4f,   // Lx, Ly, Lz, h
  grid: vec4f,   // nx, ny, nz, ray steps
  misc: vec4f,   // opacity, clip y, (unused), (unused)
  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> wall: 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 = vec3i(i32(V.grid.x), i32(V.grid.y), i32(V.grid.z));
  let h = V.dims.w;
  let g = clamp(vec3i(floor((q + 0.5 * V.dims.xyz) / h)), vec3i(0), n - vec3i(1));
  return u32(g.x + n.x * (g.y + n.y * g.z));
}

// A point on the box'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) -> f32 {
  let d = abs(abs(q) - 0.5 * V.dims.xyz);
  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 * V.dims.x;
  return 1.0 - smoothstep(w, 2.0 * w, mid);
}

@fragment
fn fs(in: VSOut) -> @location(0) vec4f {
  let half = 0.5 * V.dims.xyz;
  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 = (-half - eye) * inv;
  let tb = (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));

  let steps = i32(V.grid.w);
  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 / V.dims.w;
  let scale = max(V.up.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.y > V.misc.y) { continue; }
    let i = voxel(q);
    let v = clamp(p[i] / scale, -1.0, 1.0);
    var a = pow(abs(v), V.fwd.w) * V.misc.x * unit;
    var rgb = textureSampleLevel(cmap, samp, vec2f(0.5 + 0.5 * v, 0.5), 0.0).rgb;
    if (V.bg.w > 0.0) {
      let m = clamp(1.0 - wall[i], 0.0, 1.0);
      let am = m * V.bg.w * 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));

  // The microphone, a dot in a ring at the ray's closest approach to it.
  // Drawn on top rather than composited into the march: 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) / V.dims.x;
      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;
}

/** The camera's orthonormal frame, shared with the overlay so its lines land
 *  on the volume's pixels. */
export interface CameraFrame {
  eye: [number, number, number];
  right: [number, number, number];
  up: [number, number, number];
  fwd: [number, number, number];
  tanHalfFov: number;
  aspect: number;
}

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

const cross = (a: number[], b: number[]): [number, number, 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, number, number]): [number, number, number] => {
  const m = Math.hypot(a[0], a[1], a[2]) || 1;
  return [a[0] / m, a[1] / m, a[2] / m];
};

export function cameraFrame(camera: Camera, aspect: number): CameraFrame {
  const { az, el, dist } = camera;
  const ce = Math.cos(el);
  const toEye: [number, number, number] = [
    ce * Math.cos(az),
    ce * Math.sin(az),
    Math.sin(el),
  ];
  const eye: [number, number, number] = [toEye[0] * dist, toEye[1] * dist, toEye[2] * dist];
  const fwd: [number, number, number] = [-toEye[0], -toEye[1], -toEye[2]];
  const right = norm(cross([0, 0, 1], fwd));
  const up = cross(fwd, right);
  return { eye, right, up, fwd, tanHalfFov: TAN_HALF_FOV, aspect };
}

export interface DrawOptions {
  frame: CameraFrame;
  /** Pressure the colormap saturates at, in both directions. */
  scale: number;
  opacity: number;
  contrast: number;
  /** Samples along each ray. */
  steps: number;
  /** Everything with y above this is not drawn, in metres. */
  clipY: number;
  /** Strength of the body's grey emission, 0 to turn it off. */
  body: number;
  /** Microphone position in metres, or null to not draw it. */
  mic?: { x: number; y: number; z: number } | null;
}

export interface VolumeDims {
  nx: number;
  ny: number;
  nz: number;
  Lx: number;
  Ly: number;
  Lz: number;
  h: number;
}

export class VolumeView {
  readonly canvas: HTMLCanvasElement;

  #device: GPUDevice;
  #context: GPUCanvasContext;
  #pipeline: GPURenderPipeline;
  #layout: GPUBindGroupLayout;
  #uniform: GPUBuffer;
  #host = new ArrayBuffer(9 * 16);
  #sampler: GPUSampler;
  #cmapTexture: GPUTexture;
  #bindGroup: GPUBindGroup | null = null;
  #dims: VolumeDims | null = null;
  /** Background, matched to the CSS so the box 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: the host-owned
   *  pressure state and the wall mask, which is what the body is drawn from. */
  setSource(pressure: GPUBuffer, wall: GPUBuffer, dims: VolumeDims): void {
    this.#dims = dims;
    this.#bindGroup = this.#device.createBindGroup({
      layout: this.#layout,
      entries: [
        { binding: 0, resource: { buffer: this.#uniform } },
        { binding: 1, resource: { buffer: pressure } },
        { binding: 2, resource: { buffer: wall } },
        { 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. */
  draw(o: DrawOptions): void {
    const bg = this.#bindGroup;
    const dims = this.#dims;
    if (!bg || !dims) return;

    const { eye, right, up, fwd } = o.frame;
    const f = new Float32Array(this.#host);
    f.set([eye[0], eye[1], eye[2], o.frame.tanHalfFov], 0);
    f.set([right[0], right[1], right[2], o.frame.aspect], 4);
    f.set([up[0], up[1], up[2], o.scale], 8);
    f.set([fwd[0], fwd[1], fwd[2], o.contrast], 12);
    f.set([this.bg[0], this.bg[1], this.bg[2], o.body], 16);
    f.set([dims.Lx, dims.Ly, dims.Lz, dims.h], 20);
    f.set([dims.nx, dims.ny, dims.nz, o.steps], 24);
    f.set([o.opacity, o.clipY, 0, 0], 28);
    f.set([o.mic?.x ?? 0, o.mic?.y ?? 0, o.mic?.z ?? 0, o.mic ? 1 : 0], 32);
    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();
  }
}