/**
 * The scene viewer: the instrument drawn from its own fields.
 *
 * A `visual` section names the passes that compute the medium (normally the
 * same medium shader the precompute runs) and which buffers to draw. The
 * viewer runs those passes at the current parameter values and raymarches
 * the results on the GPU: a 'solid' volume (the body mask) is shaded as an
 * opaque surface where it crosses one half, a 'cloud' volume (a coupling
 * profile) accumulates as a translucent density. Because the fields are the
 * medium shader's own output on the actual air grid, the picture is what
 * the solver sees, grid effects included: coarsen nx and the sound hole
 * staircases, thin a wall under two cells and it visibly stops sealing.
 *
 * Markers (an object's outline, the microphone) are the one declarative
 * part: crisp annotations whose coordinates are expressions over the scope,
 * drawn on a 2D overlay with the same camera.
 */
import { evalExpr } from '../expr.ts';
import type { InstrumentFile, VisualSpec } from '../format.ts';

export interface Camera {
  /** Orbit angles in radians and distance in metres, about the origin. */
  yaw: number;
  pitch: number;
  dist: number;
}

const WG = 64;

interface Vec3 {
  x: number;
  y: number;
  z: number;
}

const basis = (cam: Camera): { eye: Vec3; right: Vec3; up: Vec3; fwd: Vec3 } => {
  const cp = Math.cos(cam.pitch);
  const eye = {
    x: cam.dist * cp * Math.cos(cam.yaw),
    y: cam.dist * cp * Math.sin(cam.yaw),
    z: cam.dist * Math.sin(cam.pitch),
  };
  const norm = (v: Vec3): Vec3 => {
    const l = Math.hypot(v.x, v.y, v.z) || 1;
    return { x: v.x / l, y: v.y / l, z: v.z / l };
  };
  const fwd = norm({ x: -eye.x, y: -eye.y, z: -eye.z });
  const right = norm({ x: -fwd.y, y: fwd.x, z: 0 }); // cross(fwd, z-up)
  const up = {
    x: right.y * fwd.z - right.z * fwd.y,
    y: right.z * fwd.x - right.x * fwd.z,
    z: right.x * fwd.y - right.y * fwd.x,
  };
  return { eye, right, up: { x: -up.x, y: -up.y, z: -up.z }, fwd };
};

const TAN_HALF = Math.tan((35 * Math.PI) / 180 / 2);

/** The raymarch fragment shader, generated for this instrument's volumes.
 *  Exported so the test suite can compile-check it headless. */
export function raymarchWgsl(visual: VisualSpec): string {
  const n = visual.volumes.length;
  const binds = visual.volumes
    .map((_, i) => `@group(0) @binding(${i + 1}) var<storage, read> f${i}: array<f32>;`)
    .join('\n');
  const samplers = visual.volumes
    .map(
      (v, i) => `
fn tri${i}(g: vec3f) -> f32 {
  let gc = clamp(g, vec3f(0.0), vec3f(f32(NX) - 1.001, f32(NY) - 1.001, f32(NZ) - 1.001));
  let i0 = vec3u(floor(gc));
  let fr = gc - floor(gc);
  let s = i0.x + NX * (i0.y + NY * i0.z);
  let v000 = f${i}[s];
  let v100 = f${i}[s + 1u];
  let v010 = f${i}[s + NX];
  let v110 = f${i}[s + NX + 1u];
  let v001 = f${i}[s + NX * NY];
  let v101 = f${i}[s + NX * NY + 1u];
  let v011 = f${i}[s + NX * NY + NX];
  let v111 = f${i}[s + NX * NY + NX + 1u];
  let a = mix(mix(v000, v100, fr.x), mix(v010, v110, fr.x), fr.y);
  let b = mix(mix(v001, v101, fr.x), mix(v011, v111, fr.x), fr.y);
  let v = mix(a, b, fr.z);
  return ${v.invert ? '1.0 - v' : 'v'};
}`,
    )
    .join('\n');
  const marchBody = visual.volumes
    .map((v, i) => {
      const col = `vec3f(${v.color.join(', ')})`;
      if (v.style === 'solid') {
        return `
    {
      let v = tri${i}(g);
      // A mask that never saturates (a wall thinner than the grid can
      // carry) draws as a haze rather than vanishing: seeing the leak is
      // the point of drawing the real fields.
      if (v <= 0.5) {
        let a = clamp(v, 0.0, 0.5) * 0.25;
        accum += trans * a * ${col};
        trans *= 1.0 - a;
      } else {
        let e = 1.0;
        let grad = vec3f(
          tri${i}(g + vec3f(e, 0.0, 0.0)) - tri${i}(g - vec3f(e, 0.0, 0.0)),
          tri${i}(g + vec3f(0.0, e, 0.0)) - tri${i}(g - vec3f(0.0, e, 0.0)),
          tri${i}(g + vec3f(0.0, 0.0, e)) - tri${i}(g - vec3f(0.0, 0.0, e)));
        var nrm = -grad;
        if (dot(nrm, nrm) < 1e-12) { nrm = vec3f(0.0, 0.0, 1.0); }
        nrm = normalize(nrm);
        let light = normalize(vec3f(0.45, -0.6, 0.75));
        let shade = 0.35 + 0.65 * max(dot(nrm, light), 0.0);
        let depth = 1.0 - 0.3 * clamp((t - t0) / (t1 - t0 + 1e-6), 0.0, 1.0);
        accum += trans * ${col} * shade * depth;
        trans = 0.0;
        break;
      }
    }`;
      }
      return `
    {
      let vn = tri${i}(g) * u[6][${i}];
      let a = 1.0 - exp(-vn * 30.0 / u[4].w);
      accum += trans * a * ${col};
      trans *= 1.0 - a;
    }`;
    })
    .join('\n');

  return `// Raymarch the instrument's medium fields. Uniforms u:
// 0: eye.xyz, tanHalf  1: right.xyz, aspect  2: up.xyz  3: fwd.xyz
// 4: halfExtents.xyz, steps  5: (unused)  6: per-cloud value scales
override NX: u32;
override NY: u32;
override NZ: u32;
@group(0) @binding(0) var<uniform> u: array<vec4f, 7>;
${binds}
${samplers}

struct VOut {
  @builtin(position) pos: vec4f,
  @location(0) ndc: vec2f,
}

@vertex
fn vmain(@builtin(vertex_index) vi: u32) -> VOut {
  let xy = vec2f(f32((vi << 1u) & 2u), f32(vi & 2u)) * 2.0 - 1.0;
  return VOut(vec4f(xy, 0.0, 1.0), xy);
}

@fragment
fn fmain(v: VOut) -> @location(0) vec4f {
  let eye = u[0].xyz;
  let dir = normalize(u[3].xyz + v.ndc.x * u[1].w * u[0].w * u[1].xyz + v.ndc.y * u[0].w * u[2].xyz);
  let half = u[4].xyz;
  // Slab intersection with the domain box.
  let inv = 1.0 / dir;
  let tA = (-half - eye) * inv;
  let tB = (half - eye) * inv;
  let tmin = min(tA, tB);
  let tmax = max(tA, tB);
  let t0 = max(max(tmin.x, tmin.y), max(tmin.z, 0.0));
  let t1 = min(min(tmax.x, tmax.y), tmax.z);
  let bg = mix(vec3f(0.10, 0.11, 0.13), vec3f(0.16, 0.17, 0.20), 0.5 * (v.ndc.y + 1.0));
  if (t1 <= t0) { return vec4f(bg, 1.0); }
  var accum = vec3f(0.0);
  var trans = 1.0;
  let steps = u[4].w;
  let dt = (t1 - t0) / steps;
  for (var k = 0.0; k < steps; k += 1.0) {
    let t = t0 + (k + 0.5) * dt;
    let p = eye + t * dir;
    let g = (p + half) / (2.0 * half) * vec3f(f32(NX), f32(NY), f32(NZ)) - 0.5;
${marchBody}
    if (trans < 0.02) { break; }
  }
  return vec4f(accum + trans * bg, 1.0);
}
`;
}

export class Viewer {
  #device: GPUDevice;
  #inst: InstrumentFile;
  #visual: VisualSpec;
  #canvas: HTMLCanvasElement;
  #overlay: HTMLCanvasElement;
  #ctx: GPUCanvasContext;
  #format: GPUTextureFormat;
  #buffers = new Map<string, GPUBuffer>();
  #uniform: GPUBuffer;
  #pipeline: GPURenderPipeline | null = null;
  #bind: GPUBindGroup | null = null;
  #scope: Record<string, number> = {};
  #scales: number[] = [1, 1, 1, 1];
  #dims: [number, number, number] = [1, 1, 1];
  #half: [number, number, number] = [0.5, 0.25, 0.25];

  constructor(device: GPUDevice, inst: InstrumentFile, canvas: HTMLCanvasElement, overlay: HTMLCanvasElement) {
    if (!inst.visual) throw new Error(`instrument '${inst.name}' has no visual section`);
    this.#device = device;
    this.#inst = inst;
    this.#visual = inst.visual;
    this.#canvas = canvas;
    this.#overlay = overlay;
    this.#format = navigator.gpu.getPreferredCanvasFormat();
    const ctx = canvas.getContext('webgpu');
    if (!ctx) throw new Error('no webgpu canvas context');
    this.#ctx = ctx;
    ctx.configure({ device, format: this.#format, alphaMode: 'opaque' });
    this.#uniform = device.createBuffer({
      label: 'view-uniforms',
      size: 7 * 16,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
  }

  /** Run the visual passes at these values and refresh normalization. */
  async update(scope: Record<string, number>): Promise<void> {
    const device = this.#device;
    const v = this.#visual;
    this.#scope = scope;
    this.#dims = v.dims.map((d) => evalExpr(d, scope)) as [number, number, number];
    this.#half = v.extents.map((e) => evalExpr(e, scope) / 2) as [number, number, number];

    for (const b of v.buffers) {
      const bytes = 4 * Math.max(1, evalExpr(b.count, scope));
      const have = this.#buffers.get(b.name);
      if (!have || have.size !== bytes) {
        have?.destroy();
        this.#buffers.set(
          b.name,
          device.createBuffer({
            label: `view-${b.name}`,
            size: bytes,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
          }),
        );
      }
    }

    device.pushErrorScope('validation');
    const enc = device.createCommandEncoder({ label: 'view-update' });
    const cp = enc.beginComputePass();
    for (const pass of v.setup) {
      const s = this.#inst.shaders[pass.shader];
      const constants: Record<string, number> = {};
      for (const c of s.constants) {
        if (!(c in scope)) throw new Error(`visual pass '${pass.shader}': constant '${c}' not in scope`);
        constants[c] = scope[c];
      }
      const pipe = await device.createComputePipelineAsync({
        label: `view-${pass.shader}`,
        layout: 'auto',
        compute: {
          module: device.createShaderModule({ code: s.code, label: `view-${pass.shader}` }),
          entryPoint: 'main',
          constants,
        },
      });
      cp.setPipeline(pipe);
      cp.setBindGroup(
        0,
        device.createBindGroup({
          layout: pipe.getBindGroupLayout(0),
          entries: pass.buffers.map((name, i) => {
            const buf = this.#buffers.get(name);
            if (!buf) throw new Error(`visual pass binds undeclared buffer '${name}'`);
            return { binding: i, resource: { buffer: buf } };
          }),
        }),
      );
      const wg = Math.ceil(evalExpr(pass.dispatch, scope) / WG);
      cp.dispatchWorkgroups(Math.min(wg, 65535), Math.ceil(wg / 65535));
    }
    cp.end();
    device.queue.submit([enc.finish()]);

    // Normalize each cloud volume to its own peak.
    for (let i = 0; i < v.volumes.length; i++) {
      const vol = v.volumes[i];
      if (vol.style !== 'cloud') {
        this.#scales[i] = 1;
        continue;
      }
      const src = this.#buffers.get(vol.buffer)!;
      const rb = device.createBuffer({ size: src.size, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
      const e2 = device.createCommandEncoder();
      e2.copyBufferToBuffer(src, 0, rb, 0, src.size);
      device.queue.submit([e2.finish()]);
      await rb.mapAsync(GPUMapMode.READ);
      const data = new Float32Array(rb.getMappedRange());
      let peak = 0;
      for (const x of data) peak = Math.max(peak, Math.abs(x));
      rb.unmap();
      rb.destroy();
      this.#scales[i] = peak > 0 ? 1 / peak : 1;
    }

    // The raymarch pipeline, rebuilt because the dims are baked in.
    const module = device.createShaderModule({ code: raymarchWgsl(v), label: 'view-raymarch' });
    this.#pipeline = await device.createRenderPipelineAsync({
      label: 'view-raymarch',
      layout: 'auto',
      vertex: { module, entryPoint: 'vmain' },
      fragment: {
        module,
        entryPoint: 'fmain',
        targets: [{ format: this.#format }],
        constants: { NX: this.#dims[0], NY: this.#dims[1], NZ: this.#dims[2] },
      },
      primitive: { topology: 'triangle-list' },
    });
    this.#bind = device.createBindGroup({
      layout: this.#pipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: this.#uniform } },
        ...v.volumes.map((vol, i) => ({ binding: i + 1, resource: { buffer: this.#buffers.get(vol.buffer)! } })),
      ],
    });
    const err = await device.popErrorScope();
    if (err) throw new Error(`viewer: ${err.message}`);
  }

  draw(cam: Camera): void {
    if (!this.#pipeline || !this.#bind) return;
    const device = this.#device;
    const { eye, right, up, fwd } = basis(cam);
    const w = this.#canvas.width;
    const h = this.#canvas.height;
    const steps = 2 * Math.max(...this.#dims);
    const u = new Float32Array(7 * 4);
    u.set([eye.x, eye.y, eye.z, TAN_HALF], 0);
    u.set([right.x, right.y, right.z, w / h], 4);
    u.set([up.x, up.y, up.z, 0], 8);
    u.set([fwd.x, fwd.y, fwd.z, 0], 12);
    u.set([...this.#half, steps], 16);
    u.set(this.#scales, 24);
    device.queue.writeBuffer(this.#uniform, 0, u);
    const enc = device.createCommandEncoder({ label: 'view-draw' });
    const rp = enc.beginRenderPass({
      colorAttachments: [
        {
          view: this.#ctx.getCurrentTexture().createView(),
          loadOp: 'clear',
          storeOp: 'store',
          clearValue: { r: 0, g: 0, b: 0, a: 1 },
        },
      ],
    });
    rp.setPipeline(this.#pipeline);
    rp.setBindGroup(0, this.#bind);
    rp.draw(3);
    rp.end();
    device.queue.submit([enc.finish()]);
    this.#drawOverlay(cam);
  }

  /** Project a world point to overlay pixels; null when behind the eye. */
  #project(cam: ReturnType<typeof basis>, p: [number, number, number]): [number, number] | null {
    const d = { x: p[0] - cam.eye.x, y: p[1] - cam.eye.y, z: p[2] - cam.eye.z };
    const dot = (a: Vec3): number => a.x * d.x + a.y * d.y + a.z * d.z;
    const z = dot(cam.fwd);
    if (z <= 1e-4) return null;
    const w = this.#overlay.width;
    const h = this.#overlay.height;
    const x = dot(cam.right) / (z * TAN_HALF * (w / h));
    const y = dot(cam.up) / (z * TAN_HALF);
    return [(x * 0.5 + 0.5) * w, (1 - (y * 0.5 + 0.5)) * h];
  }

  #drawOverlay(cam: Camera): void {
    const b = basis(cam);
    const g = this.#overlay.getContext('2d');
    if (!g) return;
    const w = this.#overlay.width;
    const h = this.#overlay.height;
    g.clearRect(0, 0, w, h);
    const seg = (p: [number, number, number], q: [number, number, number]): void => {
      const a = this.#project(b, p);
      const c = this.#project(b, q);
      if (!a || !c) return;
      g.beginPath();
      g.moveTo(a[0], a[1]);
      g.lineTo(c[0], c[1]);
      g.stroke();
    };
    // The domain box, faint.
    const [hx, hy, hz] = this.#half;
    g.strokeStyle = 'rgba(255,255,255,0.18)';
    g.lineWidth = 1;
    for (const [i, j] of [[0, 1], [0, 2], [1, 3], [2, 3]] as const) {
      const c = (k: number): [number, number] => [k & 1 ? hx : -hx, k & 2 ? hy : -hy];
      const [xi, yi] = c(i);
      const [xj, yj] = c(j);
      seg([xi, yi, -hz], [xj, yj, -hz]);
      seg([xi, yi, hz], [xj, yj, hz]);
      seg([xi, yi, -hz], [xi, yi, hz]);
    }
    const ev = (e: [string, string, string]): [number, number, number] =>
      e.map((x) => evalExpr(x, this.#scope)) as [number, number, number];
    for (const m of this.#visual.markers) {
      const col = `rgb(${m.color.map((c) => Math.round(255 * c)).join(',')})`;
      g.strokeStyle = col;
      g.fillStyle = col;
      g.lineWidth = 1.5;
      let anchor: [number, number] | null = null;
      if (m.kind === 'point') {
        const p = this.#project(b, ev(m.at));
        if (p) {
          g.beginPath();
          g.arc(p[0], p[1], 4, 0, 2 * Math.PI);
          g.fill();
          anchor = p;
        }
      } else if (m.kind === 'line') {
        const from = ev(m.from);
        const to = ev(m.to);
        seg(from, to);
        anchor = this.#project(b, to);
      } else {
        const c0 = ev(m.corner);
        const du = ev(m.du);
        const dv = ev(m.dv);
        const add = (a: [number, number, number], ...rest: [number, number, number][]): [number, number, number] =>
          rest.reduce((s, r) => [s[0] + r[0], s[1] + r[1], s[2] + r[2]] as [number, number, number], a);
        const c1 = add(c0, du);
        const c2 = add(c0, du, dv);
        const c3 = add(c0, dv);
        seg(c0, c1);
        seg(c1, c2);
        seg(c2, c3);
        seg(c3, c0);
        anchor = this.#project(b, c2);
      }
      if (m.label && anchor) {
        g.font = '11px system-ui, sans-serif';
        g.fillText(m.label, anchor[0] + 6, anchor[1] - 5);
      }
    }
  }

  destroy(): void {
    for (const b of this.#buffers.values()) b.destroy();
    this.#buffers.clear();
    this.#uniform.destroy();
  }
}

/** Predicted partial frequencies read off a probed slots buffer: one entry
 *  per underdamped root, in Hz with a relative weight, deduplicated. */
export function slotPartials(
  slots: Float32Array,
  layout: { stride: number; zRe: number; zIm: number },
  dt: number,
): { freq: number; weight: number }[] {
  const byFreq = new Map<number, number>();
  for (let s = 0; s + layout.stride <= slots.length; s += layout.stride) {
    const zRe = slots[s + layout.zRe];
    const zIm = slots[s + layout.zIm];
    const cAbs = Math.hypot(slots[s], slots[s + 1]);
    if (zIm <= 0 || cAbs === 0) continue;
    const w = Math.abs(Math.atan2(zIm, zRe)) / dt;
    const freq = w / (2 * Math.PI);
    // Weight by rough displacement (the coefficients are acceleration
    // series, which flatter high modes); this only sets tick opacity.
    const key = Math.round(freq * 2) / 2;
    byFreq.set(key, (byFreq.get(key) ?? 0) + cAbs / Math.max(w * w, 1));
  }
  const out = [...byFreq.entries()].map(([freq, weight]) => ({ freq, weight }));
  const max = out.reduce((a, p) => Math.max(a, p.weight), 0);
  return out
    .filter((p) => p.weight > 1e-4 * max)
    .map((p) => ({ freq: p.freq, weight: p.weight / max }))
    .sort((a, b) => a.freq - b.freq);
}