import { forwardRef, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react";
import {
  colormapTable,
  getColormap,
  type FieldFrames,
  type Mount,
  type PlateResult,
  type ProblemSpec,
  type Range,
} from "@plate-dyn/core";

export interface Plate3DViewProps {
  result: PlateResult;
  spec: ProblemSpec;
  frameIndex: number;
  /** Frames that colour the surface (any per-node field) and their colour range */
  colorFrames?: FieldFrames;
  colorRange?: Range;
  colormap?: string;
  /** Exaggeration of the flexural field and the static sag */
  flexScale: number;
  /** Exaggeration of the rigid-body motion (1 is true scale) */
  rigidScale: number;
  width?: number;
  height?: number;
}

type Vec3 = [number, number, number];

const VS = `
attribute vec3 aPos;
attribute vec3 aNormal;
attribute vec3 aColor;
uniform mat4 uProj;
uniform mat4 uView;
varying vec3 vColor;
varying vec3 vNormal;
void main() {
  gl_Position = uProj * uView * vec4(aPos, 1.0);
  vColor = aColor;
  vNormal = (uView * vec4(aNormal, 0.0)).xyz;
}
`;
const FS = `
precision mediump float;
varying vec3 vColor;
varying vec3 vNormal;
uniform float uFlat;
void main() {
  vec3 n = normalize(vNormal);
  float diff = abs(dot(n, normalize(vec3(0.3, 0.6, 1.0))));
  float shade = mix(0.55 + 0.45 * diff, 1.0, uFlat);
  gl_FragColor = vec4(vColor * shade, 1.0);
}
`;

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

function lookAt(eye: Vec3, center: Vec3, up: Vec3): Float32Array {
  const z = norm(sub(eye, center));
  const x = norm(cross(up, z));
  const y = cross(z, x);
  const out = new Float32Array(16);
  out[0] = x[0];
  out[4] = x[1];
  out[8] = x[2];
  out[1] = y[0];
  out[5] = y[1];
  out[9] = y[2];
  out[2] = z[0];
  out[6] = z[1];
  out[10] = z[2];
  out[12] = -dot(x, eye);
  out[13] = -dot(y, eye);
  out[14] = -dot(z, eye);
  out[15] = 1;
  return out;
}
const sub = (a: Vec3, b: Vec3): Vec3 => [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
const add = (a: Vec3, b: Vec3): Vec3 => [a[0] + b[0], a[1] + b[1], a[2] + b[2]];
const scl = (a: Vec3, s: number): Vec3 => [a[0] * s, a[1] * s, a[2] * s];
const dot = (a: Vec3, b: Vec3) => a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
const cross = (a: Vec3, b: Vec3): Vec3 => [
  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: Vec3): Vec3 => {
  const l = Math.hypot(a[0], a[1], a[2]) || 1;
  return [a[0] / l, a[1] / l, a[2] / l];
};
/** Rodrigues rotation of v about the unit axis k by angle a. */
function rotate(v: Vec3, k: Vec3, a: number): Vec3 {
  const c = Math.cos(a);
  const s = Math.sin(a);
  const kv = cross(k, v);
  const kd = dot(k, v);
  return [
    v[0] * c + kv[0] * s + k[0] * kd * (1 - c),
    v[1] * c + kv[1] * s + k[1] * kd * (1 - c),
    v[2] * c + kv[2] * s + k[2] * kd * (1 - c),
  ];
}

/**
 * The plate in its mount, in the laboratory frame: gravity points down the
 * screen, the plate is turned as it hangs or lies, the threads go up to their
 * hooks, the pads sit under it, the jaws of the vise grip it. The surface is
 * displaced by the static sag and the flexural field (exaggerated by
 * flexScale) and moved as a rigid body by the frame's rigid coordinates
 * (applied as exact rotations about the pivot, scaled by rigidScale). Drag to
 * orbit, wheel to zoom.
 */
/** What the movie recorder needs of the scene: its canvas, and a frame drawn on demand. */
export interface Plate3DHandle {
  canvas: () => HTMLCanvasElement | null;
  /** Draw frame k; `size` overrides the backing store, for recording at a larger size */
  drawFrame: (k: number, size?: { width: number; height: number }) => void;
}

export const Plate3DView = forwardRef<Plate3DHandle, Plate3DViewProps>(
  function Plate3DView(props, ref) {
    const {
      result,
      spec,
      frameIndex,
      colorFrames,
      colorRange,
      colormap = "bwr",
      flexScale,
      rigidScale,
      width = 460,
      height = 340,
    } = props;
    const canvasRef = useRef<HTMLCanvasElement>(null);
    const glRef = useRef<{
      gl: WebGLRenderingContext;
      prog: WebGLProgram;
      buffers: { pos: WebGLBuffer; nrm: WebGLBuffer; col: WebGLBuffer; idx: WebGLBuffer };
      lines: { pos: WebGLBuffer; col: WebGLBuffer; nrm: WebGLBuffer };
      loc: Record<string, number | WebGLUniformLocation | null>;
    } | null>(null);
    const [cam, setCam] = useState({ az: 0.6, el: 0.35, zoom: 1 });
    const drag = useRef<{ x: number; y: number } | null>(null);

    const nx = result.grid.x.length;
    const ny = result.grid.y.length;
    const { Lx, Ly } = spec.plate;
    const size = Math.max(Lx, Ly);

    // Lab frame: X horizontal in the plate's plane, Y up, Z the rest
    const frame = useMemo(() => {
      const s = result.static;
      const g: Vec3 = [s.gx, s.gy, s.gz];
      const gl = Math.hypot(...g);
      const up: Vec3 = gl > 0 ? scl(g, -1 / gl) : [0, 0, 1];
      const eh: Vec3 = [result.rigid.eh[0], result.rigid.eh[1], 0];
      const ed: Vec3 = [result.rigid.ed[0], result.rigid.ed[1], 0];
      let X: Vec3 = Math.abs(dot(eh, up)) < 0.99 ? eh : [1, 0, 0];
      X = norm(sub(X, scl(up, dot(X, up))));
      const Z = cross(X, up);
      const toLab = (p: Vec3): Vec3 => [dot(X, p), dot(up, p), dot(Z, p)];
      return { toLab, ed, eh, up };
    }, [result]);

    // Static geometry of the triangles: which cells are on the plate
    const mesh = useMemo(() => {
      const w0 = result.static.w0.data;
      const inside = (i: number, j: number) => Number.isFinite(w0[i + j * ny]);
      const idx: number[] = [];
      for (let j = 0; j < nx - 1; j++)
        for (let i = 0; i < ny - 1; i++) {
          if (!(inside(i, j) && inside(i + 1, j) && inside(i, j + 1) && inside(i + 1, j + 1)))
            continue;
          const a = i + j * ny;
          const b = i + 1 + j * ny;
          const c = i + (j + 1) * ny;
          const d = i + 1 + (j + 1) * ny;
          idx.push(a, b, c, b, d, c);
        }
      return { index: new Uint16Array(idx), inside };
    }, [result, nx, ny]);

    useEffect(() => {
      const canvas = canvasRef.current;
      if (!canvas) return;
      const gl = canvas.getContext("webgl", { antialias: true });
      if (!gl) return;
      const compile = (type: number, src: string) => {
        const sh = gl.createShader(type)!;
        gl.shaderSource(sh, src);
        gl.compileShader(sh);
        return sh;
      };
      const prog = gl.createProgram()!;
      gl.attachShader(prog, compile(gl.VERTEX_SHADER, VS));
      gl.attachShader(prog, compile(gl.FRAGMENT_SHADER, FS));
      gl.linkProgram(prog);
      const mk = () => gl.createBuffer()!;
      glRef.current = {
        gl,
        prog,
        buffers: { pos: mk(), nrm: mk(), col: mk(), idx: mk() },
        lines: { pos: mk(), col: mk(), nrm: mk() },
        loc: {
          aPos: gl.getAttribLocation(prog, "aPos"),
          aNormal: gl.getAttribLocation(prog, "aNormal"),
          aColor: gl.getAttribLocation(prog, "aColor"),
          uProj: gl.getUniformLocation(prog, "uProj"),
          uView: gl.getUniformLocation(prog, "uView"),
          uFlat: gl.getUniformLocation(prog, "uFlat"),
        },
      };
      return () => {
        glRef.current = null;
      };
    }, []);

    /**
     * Draw one frame. `out` overrides the backing store, which the movie
     * recorder uses to render the scene larger than it appears here; the next
     * on-screen draw puts it back.
     */
    const draw = (kIndex: number, out?: { width: number; height: number }) => {
      const r = glRef.current;
      const canvas = canvasRef.current;
      if (!r || !canvas) return;
      const { gl, prog, buffers, lines, loc } = r;
      const vw = out?.width ?? width;
      const vh = out?.height ?? height;
      const dpr = out ? 1 : Math.min(2, window.devicePixelRatio || 1);
      canvas.width = Math.round(vw * dpr);
      canvas.height = Math.round(vh * dpr);
      gl.viewport(0, 0, canvas.width, canvas.height);

      // ---- node positions in the lab frame
      const N = nx * ny;
      const kf = Math.min(kIndex, result.frames.t.length - 1);
      const phi = result.frames.w;
      const w0 = result.static.w0.data;
      const c = result.rigid.frames.map((row) => (kf >= 0 ? row[kf] : 0));
      const pivot: Vec3 = [result.rigid.pivot.x, result.rigid.pivot.y, 0];
      const { ed, eh, toLab } = frame;
      const flexOf = (n: number) => {
        const z0 = Number.isFinite(w0[n]) ? w0[n] : 0;
        const f = kf >= 0 ? phi.data[n + kf * N] : 0;
        return flexScale * (z0 + (Number.isFinite(f) ? f : 0));
      };
      const rigidPos = (p: Vec3): Vec3 => {
        if (result.rigid.nr === 0) return p;
        let q = sub(p, pivot);
        q = rotate(q, eh, -c[1] * rigidScale);
        q = rotate(q, ed, c[2] * rigidScale);
        q = add(q, pivot);
        return [q[0], q[1], q[2] + c[0] * rigidScale];
      };
      const pos = new Float32Array(3 * N);
      const labPos: Vec3[] = new Array(N);
      for (let j = 0; j < nx; j++)
        for (let i = 0; i < ny; i++) {
          const n = i + j * ny;
          const p: Vec3 = [result.grid.x[j], result.grid.y[i], flexOf(n)];
          const lp = toLab(rigidPos(p));
          labPos[n] = lp;
          pos[3 * n] = lp[0];
          pos[3 * n + 1] = lp[1];
          pos[3 * n + 2] = lp[2];
        }
      // normals from neighbouring nodes
      const nrm = new Float32Array(3 * N);
      for (let j = 0; j < nx; j++)
        for (let i = 0; i < ny; i++) {
          const n = i + j * ny;
          const jl = Math.max(0, j - 1);
          const jr = Math.min(nx - 1, j + 1);
          const il = Math.max(0, i - 1);
          const ir = Math.min(ny - 1, i + 1);
          const tx = sub(labPos[i + jr * ny], labPos[i + jl * ny]);
          const ty = sub(labPos[ir + j * ny], labPos[il + j * ny]);
          const nn = norm(cross(tx, ty));
          nrm[3 * n] = nn[0];
          nrm[3 * n + 1] = nn[1];
          nrm[3 * n + 2] = nn[2];
        }
      // colours
      const col = new Float32Array(3 * N);
      const table = colormapTable(getColormap(colormap));
      const range = colorRange ?? { min: -1, max: 1 };
      const span = range.max - range.min || 1;
      for (let n = 0; n < N; n++) {
        let v = colorFrames && kf >= 0 ? colorFrames.data[n + kf * N] : NaN;
        if (!Number.isFinite(v)) v = 0.5 * (range.min + range.max);
        let t = (v - range.min) / span;
        t = t < 0 ? 0 : t > 1 ? 1 : t;
        const k = Math.round(t * 255) * 3;
        col[3 * n] = table[k] / 255;
        col[3 * n + 1] = table[k + 1] / 255;
        col[3 * n + 2] = table[k + 2] / 255;
      }

      // ---- the mount, as lines and boxes in the lab frame
      const L: number[] = [];
      const LC: number[] = [];
      const line = (a: Vec3, b: Vec3, rgb: Vec3) => {
        L.push(...a, ...b);
        LC.push(...rgb, ...rgb);
      };
      const grey: Vec3 = [0.35, 0.35, 0.38];
      const dark: Vec3 = [0.15, 0.15, 0.18];
      const red: Vec3 = [0.82, 0.29, 0.36];
      const box = (lo: Vec3, hi: Vec3, rgb: Vec3, toLabFn = toLab) => {
        const cs: Vec3[] = [];
        for (let k = 0; k < 8; k++)
          cs.push(toLabFn([k & 1 ? hi[0] : lo[0], k & 2 ? hi[1] : lo[1], k & 4 ? hi[2] : lo[2]]));
        const e = [
          [0, 1],
          [2, 3],
          [4, 5],
          [6, 7],
          [0, 2],
          [1, 3],
          [4, 6],
          [5, 7],
          [0, 4],
          [1, 5],
          [2, 6],
          [3, 7],
        ];
        for (const [a, b] of e) line(cs[a], cs[b], rgb);
      };
      const circle = (centre: Vec3, r: number, rgb: Vec3) => {
        const n = 32;
        for (let k = 0; k < n; k++) {
          const a0 = (2 * Math.PI * k) / n;
          const a1 = (2 * Math.PI * (k + 1)) / n;
          line(
            toLab([centre[0] + r * Math.cos(a0), centre[1] + r * Math.sin(a0), centre[2]]),
            toLab([centre[0] + r * Math.cos(a1), centre[1] + r * Math.sin(a1), centre[2]]),
            rgb,
          );
        }
      };
      const m: Mount = spec.mount;
      const nodeAt = (x: number, y: number): number => {
        let best = 0;
        let bd = Infinity;
        for (let j = 0; j < nx; j++)
          for (let i = 0; i < ny; i++) {
            const d = (result.grid.x[j] - x) ** 2 + (result.grid.y[i] - y) ** 2;
            if (d < bd && mesh.inside(i, j)) {
              bd = d;
              best = i + j * ny;
            }
          }
        return best;
      };
      const t = 0.02 * size; // a nominal thickness for the drawn parts
      if (m.type === "thread") {
        const n = nodeAt(m.x, m.y);
        const hole = labPos[n];
        const hook = add(
          toLab([result.grid.x[Math.floor(n / ny)], result.grid.y[n % ny], 0]),
          scl(frame.up, m.length),
        );
        line(hole, hook, dark);
        line(add(hook, [-t, 0, 0]), add(hook, [t, 0, 0]), dark);
      } else if (m.type === "threads") {
        for (const [x, y] of [
          [m.x1, m.y1],
          [m.x2, m.y2],
        ]) {
          const n = nodeAt(x, y);
          const hole = labPos[n];
          const hook = add(
            toLab([result.grid.x[Math.floor(n / ny)], result.grid.y[n % ny], 0]),
            scl(frame.up, m.length),
          );
          line(hole, hook, dark);
          line(add(hook, [-t, 0, 0]), add(hook, [t, 0, 0]), dark);
        }
      } else if (m.type === "pads") {
        for (const p of m.pads)
          box([p.x - p.rx, p.y - p.ry, -2 * t], [p.x + p.rx, p.y + p.ry, 0], grey);
      } else if (m.type === "cymbal") {
        const cx = Lx / 2;
        const cy = Ly / 2;
        circle([cx, cy, -t], m.r_felt, grey);
        circle([cx, cy, t], m.r_felt, grey);
        circle([cx, cy, -t], m.r_rod, dark);
        line(toLab([cx, cy, -6 * t]), toLab([cx, cy, 3 * t]), dark);
      } else if (m.type === "vise") {
        const d = m.depth;
        const lo: Vec3 =
          m.edge === "x0"
            ? [-t, -t, -t]
            : m.edge === "x1"
              ? [Lx - d, -t, -t]
              : m.edge === "y0"
                ? [-t, -t, -t]
                : [-t, Ly - d, -t];
        const hi: Vec3 =
          m.edge === "x0"
            ? [d, Ly + t, t]
            : m.edge === "x1"
              ? [Lx + t, Ly + t, t]
              : m.edge === "y0"
                ? [Lx + t, d, t]
                : [Lx + t, Ly + t, t];
        box(lo, hi, dark);
      } else if (m.type === "pinched") {
        const n = nodeAt(m.x, m.y);
        const p = labPos[n];
        const nrmL = norm(
          cross(
            sub(labPos[nodeAt(m.x + Lx / 10, m.y)], p),
            sub(labPos[nodeAt(m.x, m.y + Ly / 10)], p),
          ),
        );
        line(add(p, scl(nrmL, t)), add(p, scl(nrmL, 4 * t)), dark);
        line(add(p, scl(nrmL, -t)), add(p, scl(nrmL, -4 * t)), dark);
      }
      // gravity arrow beside the plate
      {
        const corner = add(toLab([Lx * 1.15, Ly * 1.05, 0]), [0, 0, 0]);
        const tip = add(corner, scl(frame.up, -0.25 * size));
        line(corner, tip, red);
        const side = scl(cross(frame.up, [0, 0, 1]), 0.03 * size);
        line(tip, add(add(tip, scl(frame.up, 0.06 * size)), side), red);
        line(tip, sub(add(tip, scl(frame.up, 0.06 * size)), side), red);
      }
      // outline of the undeformed plate at rest, faint
      {
        const pale: Vec3 = [0.75, 0.75, 0.78];
        if (spec.plate.shape === "ellipse") circle([Lx / 2, Ly / 2, 0], Lx / 2, pale);
        else box([0, 0, 0], [Lx, Ly, 0], pale);
      }

      // ---- camera
      const centre = toLab([Lx / 2, Ly / 2, 0]);
      const dist = (2.2 * size) / cam.zoom;
      const eye: Vec3 = add(centre, [
        dist * Math.cos(cam.el) * Math.sin(cam.az),
        dist * Math.sin(cam.el),
        dist * Math.cos(cam.el) * Math.cos(cam.az),
      ]);
      const view = lookAt(eye, centre, [0, 1, 0]);
      const proj = perspective(0.6, vw / vh, 0.01 * size, 20 * size);

      gl.clearColor(0.97, 0.97, 0.975, 1);
      gl.enable(gl.DEPTH_TEST);
      gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);
      gl.useProgram(prog);
      gl.uniformMatrix4fv(loc.uProj as WebGLUniformLocation, false, proj);
      gl.uniformMatrix4fv(loc.uView as WebGLUniformLocation, false, view);
      const attr = (buf: WebGLBuffer, data: Float32Array, a: number) => {
        gl.bindBuffer(gl.ARRAY_BUFFER, buf);
        gl.bufferData(gl.ARRAY_BUFFER, data, gl.DYNAMIC_DRAW);
        gl.enableVertexAttribArray(a);
        gl.vertexAttribPointer(a, 3, gl.FLOAT, false, 0, 0);
      };
      gl.uniform1f(loc.uFlat as WebGLUniformLocation, 0);
      attr(buffers.pos, pos, loc.aPos as number);
      attr(buffers.nrm, nrm, loc.aNormal as number);
      attr(buffers.col, col, loc.aColor as number);
      gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, buffers.idx);
      gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, mesh.index, gl.STATIC_DRAW);
      gl.drawElements(gl.TRIANGLES, mesh.index.length, gl.UNSIGNED_SHORT, 0);
      // lines
      gl.uniform1f(loc.uFlat as WebGLUniformLocation, 1);
      const lp = new Float32Array(L);
      attr(lines.pos, lp, loc.aPos as number);
      attr(lines.nrm, new Float32Array(L.length), loc.aNormal as number);
      attr(lines.col, new Float32Array(LC), loc.aColor as number);
      gl.drawArrays(gl.LINES, 0, lp.length / 3);
    };

    // The latest draw, so that the recorder can call it between renders; a
    // re-render with unchanged props leaves the effect alone, so the two never
    // fight over the canvas.
    const drawRef = useRef(draw);
    drawRef.current = draw;
    useImperativeHandle(
      ref,
      () => ({
        canvas: () => canvasRef.current,
        drawFrame: (k: number, out?: { width: number; height: number }) => drawRef.current(k, out),
      }),
      [],
    );

    useEffect(() => {
      drawRef.current(frameIndex);
    }, [
      result,
      spec,
      frameIndex,
      colorFrames,
      colorRange,
      colormap,
      flexScale,
      rigidScale,
      width,
      height,
      cam,
      frame,
      mesh,
      nx,
      ny,
      Lx,
      Ly,
      size,
    ]);

    return (
      <canvas
        ref={canvasRef}
        style={{ width, height, display: "block", cursor: "grab", borderRadius: 4 }}
        onPointerDown={(e) => {
          drag.current = { x: e.clientX, y: e.clientY };
          (e.target as HTMLElement).setPointerCapture(e.pointerId);
        }}
        onPointerMove={(e) => {
          if (!drag.current) return;
          const dx = e.clientX - drag.current.x;
          const dy = e.clientY - drag.current.y;
          drag.current = { x: e.clientX, y: e.clientY };
          setCam((c) => ({
            ...c,
            az: c.az - dx * 0.01,
            el: Math.max(-1.4, Math.min(1.4, c.el + dy * 0.01)),
          }));
        }}
        onPointerUp={() => {
          drag.current = null;
        }}
        onWheel={(e) => {
          e.preventDefault();
          setCam((c) => ({
            ...c,
            zoom: Math.max(0.3, Math.min(5, c.zoom * Math.exp(-e.deltaY * 0.001))),
          }));
        }}
        aria-label="the plate in its mount"
      />
    );
  },
);