/**
 * A WebGPU view of the scene: the surface colored by the surface field, the
 * three planes, and the source. For the room, walls whose outward normal faces
 * the camera are cut away in the fragment shader, so the view looks into the
 * room through the nearer walls; the cymbal is drawn whole. Drag to orbit,
 * scroll or pinch to zoom.
 */
import { basis2d, type Surface } from '../surface.ts';
import { DIVERGING, SEQUENTIAL } from './colormaps.ts';
import type { Field, Mode } from './field.ts';

/** Multisampling; 1 when the query string has ?nomsaa (SwiftShader's GPU process does not survive the resolve in headless Chrome). */
const SAMPLES = typeof location !== 'undefined' && location.search.includes('nomsaa') ? 1 : 4;

const SHARED = /* wgsl */ `
struct View {
  viewProj: mat4x4f,
  eye: vec4f,
  p1: vec4f,   // cos wt, sin wt, umax, 0
  p2: vec4f,   // mode (0 pressure, 1 level), db min, db max, level colormap (0 sequential, 1 diverging)
  p3: vec4f,   // cut away surfaces facing the camera (0, 1), 0, 0, 0
}
@group(0) @binding(0) var<uniform> V: View;
@group(0) @binding(1) var cmapD: texture_2d<f32>;
@group(0) @binding(2) var cmapS: texture_2d<f32>;
@group(0) @binding(3) var samp: sampler;

// lref: the reference of the level at this point ('ref' is reserved in WGSL)
fn colorOf(re: f32, im: f32, lref: f32) -> vec3f {
  if (V.p2.x < 0.5) {
    let t = clamp(0.5 + 0.5 * (re * V.p1.x + im * V.p1.y) / V.p1.z, 0.0, 1.0);
    return textureSampleLevel(cmapD, samp, vec2f(t, 0.5), 0.0).rgb;
  }
  let db = 20.0 * log2(max(length(vec2f(re, im)) / lref, 1e-30)) * 0.30102999566;
  let t = clamp((db - V.p2.y) / (V.p2.z - V.p2.y), 0.0, 1.0);
  if (V.p2.w > 0.5) { return textureSampleLevel(cmapD, samp, vec2f(t, 0.5), 0.0).rgb; }
  return textureSampleLevel(cmapS, samp, vec2f(t, 0.5), 0.0).rgb;
}
`;

const WALL = /* wgsl */ `
${SHARED}
struct VOut {
  @builtin(position) pos: vec4f,
  @location(0) world: vec3f,
  @location(1) nrm: vec3f,
  @location(2) u: vec2f,
  @location(3) lref: f32,
}
@vertex fn vs(@location(0) p: vec3f, @location(1) n: vec3f, @location(2) u: vec2f, @location(3) lref: f32) -> VOut {
  var o: VOut;
  o.pos = V.viewProj * vec4f(p, 1.0);
  o.world = p;
  o.nrm = n;
  o.u = u;
  o.lref = lref;
  return o;
}
@fragment fn fs(i: VOut) -> @location(0) vec4f {
  let toEye = normalize(V.eye.xyz - i.world);
  let n = normalize(i.nrm);
  if (V.p3.x > 0.5 && dot(n, toEye) > 0.0) { discard; }
  let shade = 0.6 + 0.4 * abs(dot(n, toEye));
  return vec4f(colorOf(i.u.x, i.u.y, i.lref) * shade, 1.0);
}
`;

const PLANE = /* wgsl */ `
${SHARED}
@group(0) @binding(4) var field: texture_2d<f32>;
struct VOut {
  @builtin(position) pos: vec4f,
  @location(0) uv: vec2f,
}
@vertex fn vs(@location(0) p: vec3f, @location(1) uv: vec2f) -> VOut {
  var o: VOut;
  o.pos = V.viewProj * vec4f(p, 1.0);
  o.uv = uv;
  return o;
}
@fragment fn fs(i: VOut) -> @location(0) vec4f {
  // bilinear by hand: r32float textures are not filterable everywhere
  let dim = vec2f(textureDimensions(field));
  let x = clamp(i.uv * (dim - 1.0), vec2f(0.0), dim - 1.0);
  let x0 = vec2i(floor(x));
  let x1 = min(x0 + 1, vec2i(dim) - 1);
  let f = x - vec2f(x0);
  // (Re u, Im u, level reference)
  let a = textureLoad(field, x0, 0).xyz;
  let b = textureLoad(field, vec2i(x1.x, x0.y), 0).xyz;
  let c = textureLoad(field, vec2i(x0.x, x1.y), 0).xyz;
  let d = textureLoad(field, x1, 0).xyz;
  // NaN marks points where the field is not computed
  if (any(a != a) || any(b != b) || any(c != c) || any(d != d)) { discard; }
  let u = mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
  return vec4f(colorOf(u.x, u.y, u.z), 1.0);
}
`;

const MARKER = /* wgsl */ `
${SHARED}
@vertex fn vs(@location(0) p: vec3f) -> @builtin(position) vec4f {
  return V.viewProj * vec4f(p, 1.0);
}
@fragment fn fs() -> @location(0) vec4f {
  return vec4f(0.05, 0.05, 0.05, 1.0);
}
`;

type Mat4 = Float32Array;
const BG = { r: 0.043, g: 0.055, b: 0.071, a: 1 };

function perspective(fovy: number, aspect: number, near: number, far: number): Mat4 {
  const f = 1 / Math.tan(fovy / 2);
  const m = new Float32Array(16);
  m[0] = f / aspect;
  m[5] = f;
  m[10] = far / (near - far);
  m[11] = -1;
  m[14] = (near * far) / (near - far);
  return m;
}

function lookAt(eye: number[], target: number[], up: number[]): Mat4 {
  const z = norm(sub(eye, target));
  const x = norm(cross(up, z));
  const y = cross(z, x);
  const m = new Float32Array(16);
  m.set([x[0], y[0], z[0], 0, x[1], y[1], z[1], 0, x[2], y[2], z[2], 0, -dot(x, eye), -dot(y, eye), -dot(z, eye), 1]);
  return m;
}

function mul(a: Mat4, b: Mat4): Mat4 {
  const m = new Float32Array(16);
  for (let c = 0; c < 4; c++) for (let r = 0; r < 4; r++) {
    let s = 0;
    for (let k = 0; k < 4; k++) s += a[k * 4 + r] * b[c * 4 + k];
    m[c * 4 + r] = s;
  }
  return m;
}

const sub = (a: number[], b: number[]) => [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
const dot = (a: number[], b: number[]) => a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
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 l = Math.hypot(a[0], a[1], a[2]);
  return [a[0] / l, a[1] / l, a[2] / l];
};

export interface Show {
  walls: boolean;
  planes: [boolean, boolean, boolean];
}

export class RoomView {
  private ctx: GPUCanvasContext;
  private format: GPUTextureFormat;
  private wallPipe: GPURenderPipeline;
  private planePipe: GPURenderPipeline;
  private markerPipe: GPURenderPipeline;
  private ubuf: GPUBuffer;
  private cmapD: GPUTexture;
  private cmapS: GPUTexture;
  private sampler: GPUSampler;
  private baseGroup: GPUBindGroup;
  private wall: { vbuf: GPUBuffer; ibuf: GPUBuffer; count: number } | null = null;
  private planes: { vbuf: GPUBuffer; tex: GPUTexture; group: GPUBindGroup }[] = [];
  private marker: { vbuf: GPUBuffer; count: number } | null = null;
  private msaa: GPUTexture | null = null;
  private depth: GPUTexture | null = null;
  private theta = 0.75;
  private phi = 0.95;
  private dist = 5.2;
  private field: Field | null = null;
  private dirty = true;

  constructor(
    private device: GPUDevice,
    private canvas: HTMLCanvasElement,
  ) {
    this.ctx = canvas.getContext('webgpu') as GPUCanvasContext;
    this.format = navigator.gpu.getPreferredCanvasFormat();
    this.ctx.configure({ device, format: this.format, alphaMode: 'opaque' });
    this.ubuf = device.createBuffer({ size: 128, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    const mkmap = (lut: Uint8Array<ArrayBuffer>) => {
      const t = device.createTexture({ size: [256, 1], format: 'rgba8unorm', usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST });
      device.queue.writeTexture({ texture: t }, lut, { bytesPerRow: 1024 }, [256, 1]);
      return t;
    };
    this.cmapD = mkmap(DIVERGING);
    this.cmapS = mkmap(SEQUENTIAL);
    this.sampler = device.createSampler({ magFilter: 'linear', minFilter: 'linear', addressModeU: 'clamp-to-edge' });

    const baseEntries: GPUBindGroupLayoutEntry[] = [
      { binding: 0, visibility: GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, buffer: { type: 'uniform' } },
      { binding: 1, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'float' } },
      { binding: 2, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'float' } },
      { binding: 3, visibility: GPUShaderStage.FRAGMENT, sampler: { type: 'filtering' } },
    ];
    const baseLayout = device.createBindGroupLayout({ entries: baseEntries });
    const planeLayout = device.createBindGroupLayout({
      entries: [...baseEntries, { binding: 4, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'unfilterable-float' } }],
    });
    this.baseGroup = device.createBindGroup({ layout: baseLayout, entries: this.baseEntries() });
    const common = (code: string, layout: GPUBindGroupLayout, buffers: GPUVertexBufferLayout[], topology: GPUPrimitiveTopology = 'triangle-list'): GPURenderPipeline => {
      const module = device.createShaderModule({ code });
      return device.createRenderPipeline({
        layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
        vertex: { module, entryPoint: 'vs', buffers },
        fragment: { module, entryPoint: 'fs', targets: [{ format: this.format }] },
        primitive: { topology, cullMode: 'none' },
        depthStencil: { format: 'depth24plus', depthWriteEnabled: true, depthCompare: 'less' },
        multisample: { count: SAMPLES },
      });
    };
    this.wallPipe = common(WALL, baseLayout, [
      {
        arrayStride: 36,
        attributes: [
          { shaderLocation: 0, offset: 0, format: 'float32x3' },
          { shaderLocation: 1, offset: 12, format: 'float32x3' },
          { shaderLocation: 2, offset: 24, format: 'float32x2' },
          { shaderLocation: 3, offset: 32, format: 'float32' },
        ],
      },
    ]);
    this.planePipe = common(PLANE, planeLayout, [
      {
        arrayStride: 20,
        attributes: [
          { shaderLocation: 0, offset: 0, format: 'float32x3' },
          { shaderLocation: 1, offset: 12, format: 'float32x2' },
        ],
      },
    ]);
    this.planeLayout = planeLayout;
    this.markerPipe = common(MARKER, baseLayout, [{ arrayStride: 12, attributes: [{ shaderLocation: 0, offset: 0, format: 'float32x3' }] }]);
    this.attachControls();
    new ResizeObserver(() => {
      this.dirty = true;
    }).observe(canvas);
  }

  private planeLayout: GPUBindGroupLayout;

  private baseEntries(): GPUBindGroupEntry[] {
    return [
      { binding: 0, resource: { buffer: this.ubuf } },
      { binding: 1, resource: this.cmapD.createView() },
      { binding: 2, resource: this.cmapS.createView() },
      { binding: 3, resource: this.sampler },
    ];
  }

  private attachControls() {
    const c = this.canvas;
    const pointers = new Map<number, { x: number; y: number }>();
    let pinch = 0;
    c.addEventListener('pointerdown', (e) => {
      c.setPointerCapture(e.pointerId);
      pointers.set(e.pointerId, { x: e.clientX, y: e.clientY });
      if (pointers.size === 2) {
        const [a, b] = [...pointers.values()];
        pinch = Math.hypot(a.x - b.x, a.y - b.y);
      }
    });
    c.addEventListener('pointermove', (e) => {
      const p = pointers.get(e.pointerId);
      if (!p) return;
      if (pointers.size === 1) {
        this.theta -= (e.clientX - p.x) * 0.008;
        this.phi = Math.min(Math.PI - 0.05, Math.max(0.05, this.phi - (e.clientY - p.y) * 0.008));
      }
      p.x = e.clientX;
      p.y = e.clientY;
      if (pointers.size === 2) {
        const [a, b] = [...pointers.values()];
        const d = Math.hypot(a.x - b.x, a.y - b.y);
        if (pinch > 0) this.zoom(pinch / d);
        pinch = d;
      }
      this.dirty = true;
    });
    const up = (e: PointerEvent) => {
      pointers.delete(e.pointerId);
      pinch = 0;
    };
    c.addEventListener('pointerup', up);
    c.addEventListener('pointercancel', up);
    c.addEventListener(
      'wheel',
      (e) => {
        e.preventDefault();
        this.zoom(Math.exp(e.deltaY * 0.001));
      },
      { passive: false },
    );
  }

  /** the scene's radius: 1 for the room */
  private scale = 1;
  private center: [number, number, number] = [0, 0, 0];

  private zoom(f: number) {
    this.dist = Math.min(12 * this.scale, Math.max(2.2 * this.scale, this.dist * f));
    this.dirty = true;
  }

  /** Drop the displayed field, for while a new one is computed. */
  clear() {
    this.wall?.vbuf.destroy();
    this.wall?.ibuf.destroy();
    this.wall = null;
    for (const p of this.planes) {
      p.vbuf.destroy();
      p.tex.destroy();
    }
    this.planes = [];
    this.marker?.vbuf.destroy();
    this.marker = null;
    this.field = null;
    this.dirty = true;
  }

  /** The surface mesh: every patch sampled on an nv x nv grid, with the surface field interpolated by the patch's Legendre expansion. */
  setField(surf: Surface, f: Field, nv = 9) {
    const { device } = this;
    this.field = f;
    if (f.scene.radius !== this.scale) {
      this.dist *= f.scene.radius / this.scale;
      this.scale = f.scene.radius;
    }
    this.center = f.scene.center;
    const npols = surf.npols;
    const npatches = surf.npatches;
    const npts = surf.npts;
    const uv = new Float64Array(2 * nv * nv);
    for (let a = 0; a < nv; a++) for (let b = 0; b < nv; b++) {
      uv[a * nv + b] = -1 + (2 * b) / (nv - 1);
      uv[nv * nv + a * nv + b] = -1 + (2 * a) / (nv - 1);
    }
    const B = basis2d(uv, surf.norder); // (npols, nv^2)
    const ng = nv * nv;
    const verts = new Float32Array(9 * ng * npatches);
    const idx = new Uint32Array(6 * (nv - 1) * (nv - 1) * npatches);
    const cr = new Float64Array(npols);
    const ci = new Float64Array(npols);
    let ii = 0;
    for (let j = 0; j < npatches; j++) {
      const g = surf.evalGeometry(j, uv);
      for (let n = 0; n < npols; n++) {
        let sr = 0;
        let si = 0;
        for (let l = 0; l < npols; l++) {
          const U = surf.umat[n * npols + l];
          sr += U * f.uWall[2 * (j * npols + l)];
          si += U * f.uWall[2 * (j * npols + l) + 1];
        }
        cr[n] = sr;
        ci[n] = si;
      }
      for (let q = 0; q < ng; q++) {
        const o = 9 * (j * ng + q);
        const ru = [g[3 * ng + q], g[4 * ng + q], g[5 * ng + q]];
        const rv = [g[6 * ng + q], g[7 * ng + q], g[8 * ng + q]];
        const nn = norm(cross(ru, rv));
        let ur = 0;
        let ui = 0;
        for (let n = 0; n < npols; n++) {
          ur += cr[n] * B[n * ng + q];
          ui += ci[n] * B[n * ng + q];
        }
        const ref = f.freeFieldLevel ? 1 / (4 * Math.PI * Math.hypot(g[q] - f.x0[0], g[ng + q] - f.x0[1], g[2 * ng + q] - f.x0[2])) : f.ref;
        verts.set([g[q], g[ng + q], g[2 * ng + q], nn[0], nn[1], nn[2], ur, ui, ref], o);
      }
      for (let a = 0; a < nv - 1; a++) for (let b = 0; b < nv - 1; b++) {
        const v0 = j * ng + a * nv + b;
        idx.set([v0, v0 + 1, v0 + nv + 1, v0, v0 + nv + 1, v0 + nv], ii);
        ii += 6;
      }
    }
    void npts;
    this.wall?.vbuf.destroy();
    this.wall?.ibuf.destroy();
    const vbuf = device.createBuffer({ size: verts.byteLength, usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(vbuf, 0, verts);
    const ibuf = device.createBuffer({ size: idx.byteLength, usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(ibuf, 0, idx);
    this.wall = { vbuf, ibuf, count: idx.length };

    // the three planes, as textures of (Re u, Im u, level reference) with NaN where there is no field
    for (const p of this.planes) {
      p.vbuf.destroy();
      p.tex.destroy();
    }
    this.planes = f.slices.map(({ slice, u, ref }) => {
      const n = f.ngrid;
      const tex = device.createTexture({ size: [n, n], format: 'rgba32float', usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST });
      const data = new Float32Array(4 * n * n);
      for (let q = 0; q < n * n; q++) {
        data[4 * q] = slice.mask[q] ? u[2 * q] : NaN;
        data[4 * q + 1] = slice.mask[q] ? u[2 * q + 1] : NaN;
        data[4 * q + 2] = ref ? ref[q] : f.ref;
      }
      device.queue.writeTexture({ texture: tex }, data, { bytesPerRow: 16 * n }, [n, n]);
      const [i, jx] = slice.axes;
      const [ea, eb] = slice.extent;
      const corner = (s: number, t: number) => {
        const x = [0, 0, 0];
        x[i] = ea[0] + (ea[1] - ea[0]) * s;
        x[jx] = eb[0] + (eb[1] - eb[0]) * t;
        x[slice.fixed] = slice.value;
        return [...x, s, t];
      };
      const quad = new Float32Array([...corner(0, 0), ...corner(1, 0), ...corner(1, 1), ...corner(0, 0), ...corner(1, 1), ...corner(0, 1)]);
      const vb = device.createBuffer({ size: quad.byteLength, usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST });
      device.queue.writeBuffer(vb, 0, quad);
      const group = device.createBindGroup({ layout: this.planeLayout, entries: [...this.baseEntries(), { binding: 4, resource: tex.createView() }] });
      return { vbuf: vb, tex, group };
    });

    // the source: a small octahedron
    const r = 0.035 * this.scale;
    const [x, y, z] = f.x0;
    const P = [
      [x + r, y, z], [x - r, y, z], [x, y + r, z], [x, y - r, z], [x, y, z + r], [x, y, z - r],
    ];
    const faces = [[0, 2, 4], [2, 1, 4], [1, 3, 4], [3, 0, 4], [2, 0, 5], [1, 2, 5], [3, 1, 5], [0, 3, 5]];
    const mv = new Float32Array(faces.flat().flatMap((k) => P[k]));
    this.marker?.vbuf.destroy();
    const mb = device.createBuffer({ size: mv.byteLength, usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(mb, 0, mv);
    this.marker = { vbuf: mb, count: faces.length * 3 };
    this.dirty = true;
  }

  private targets() {
    const dpr = Math.min(2, window.devicePixelRatio || 1);
    const w = Math.max(1, Math.round(this.canvas.clientWidth * dpr));
    const h = Math.max(1, Math.round(this.canvas.clientHeight * dpr));
    if (this.canvas.width !== w || this.canvas.height !== h || !this.depth) {
      this.canvas.width = w;
      this.canvas.height = h;
      this.msaa?.destroy();
      this.depth?.destroy();
      this.msaa = this.device.createTexture({ size: [w, h], format: this.format, sampleCount: SAMPLES, usage: GPUTextureUsage.RENDER_ATTACHMENT });
      this.depth = this.device.createTexture({ size: [w, h], format: 'depth24plus', sampleCount: SAMPLES, usage: GPUTextureUsage.RENDER_ATTACHMENT });
    }
    return { w, h };
  }

  draw(mode: Mode, phase: number, show: Show, animating: boolean) {
    if (!animating && !this.dirty) return;
    this.dirty = false;
    const { device } = this;
    const { w, h } = this.targets();
    const c = this.center;
    const eye = [
      c[0] + this.dist * Math.sin(this.phi) * Math.cos(this.theta),
      c[1] + this.dist * Math.sin(this.phi) * Math.sin(this.theta),
      c[2] + this.dist * Math.cos(this.phi),
    ];
    const vp = mul(perspective((35 * Math.PI) / 180, w / h, 0.01 * this.scale, 50 * this.scale), lookAt(eye, c, [0, 0, 1]));
    const f = this.field;
    const u = new Float32Array(32);
    u.set(vp, 0);
    u.set([eye[0], eye[1], eye[2], 1], 16);
    u.set([Math.cos(phase), Math.sin(phase), f?.umax ?? 1, 0], 20);
    u.set([mode === 'pressure' ? 0 : 1, f?.level.min ?? -30, f?.level.max ?? 15, f?.level.diverging ? 1 : 0], 24);
    u.set([f?.scene.cutaway === false ? 0 : 1, 0, 0, 0], 28);
    device.queue.writeBuffer(this.ubuf, 0, u);
    const enc = device.createCommandEncoder();
    const pass = enc.beginRenderPass({
      colorAttachments: [
        SAMPLES > 1
          ? { view: this.msaa!.createView(), resolveTarget: this.ctx.getCurrentTexture().createView(), clearValue: BG, loadOp: 'clear', storeOp: 'discard' }
          : { view: this.ctx.getCurrentTexture().createView(), clearValue: BG, loadOp: 'clear', storeOp: 'store' },
      ],
      depthStencilAttachment: { view: this.depth!.createView(), depthClearValue: 1, depthLoadOp: 'clear', depthStoreOp: 'discard' },
    });
    if (this.wall && show.walls) {
      pass.setPipeline(this.wallPipe);
      pass.setBindGroup(0, this.baseGroup);
      pass.setVertexBuffer(0, this.wall.vbuf);
      pass.setIndexBuffer(this.wall.ibuf, 'uint32');
      pass.drawIndexed(this.wall.count);
    }
    if (this.planes.length) {
      pass.setPipeline(this.planePipe);
      this.planes.forEach((p, s) => {
        if (!show.planes[s]) return;
        pass.setBindGroup(0, p.group);
        pass.setVertexBuffer(0, p.vbuf);
        pass.draw(6);
      });
    }
    if (this.marker) {
      pass.setPipeline(this.markerPipe);
      pass.setBindGroup(0, this.baseGroup);
      pass.setVertexBuffer(0, this.marker.vbuf);
      pass.draw(this.marker.count);
    }
    pass.end();
    device.queue.submit([enc.finish()]);
  }

  invalidate() {
    this.dirty = true;
  }
}