import { css, drawAxes, fmtNum, PAD, setupCanvas } from './plot';
import { rampLUT } from './spectrogram';
import { timelineBar, type Timeline } from './timeline';

// Displacement viewer for a chain of beads along x: side (y) and top (z)
// views at one frame, plus a space-time image of one component.

export interface ChainGeometry {
  n: number;
  restX: Float64Array; // bead rest x
  leftWall: number | null; // x of left anchor, null if free
  rightWall: number | null;
  slideLeft: boolean;
  slideRight: boolean;
  pinned: Uint8Array;
}

export interface ChainFrames {
  frames: Float32Array; // nFrames x 3n
  times: Float64Array;
}

export class ChainView {
  private g: ChainGeometry | null = null;
  private f: ChainFrames | null = null;
  private comp = 1;
  private maxAbs = [0, 0, 0];
  private exaggeration = 1;
  private stImage: HTMLCanvasElement | null = null;
  private stKey = '';
  private stFrames: Float32Array | null = null;

  private side: HTMLCanvasElement;
  private top: HTMLCanvasElement;
  private st: HTMLCanvasElement;
  private note: HTMLParagraphElement;

  private get frame() {
    return this.tl.frame;
  }

  constructor(root: HTMLElement, private tl: Timeline) {
    root.innerHTML = `
      <div data-r="bar"></div>
      <div class="canvas-wrap short"><canvas data-r="side"></canvas></div>
      <div class="canvas-wrap short"><canvas data-r="top"></canvas></div>
      <p class="note" data-r="note"></p>
      <div class="viewer-bar"><span class="sub">Space-time</span>
        <label class="inline">component <select data-r="comp">
          <option value="1" selected>y (transverse)</option><option value="2">z (transverse)</option><option value="0">x (longitudinal)</option>
        </select></label></div>
      <div class="canvas-wrap"><canvas data-r="st"></canvas></div>`;
    const q = <T extends HTMLElement>(r: string) => root.querySelector(`[data-r="${r}"]`) as T;
    q('bar').replaceWith(timelineBar(tl));
    this.side = q('side');
    this.top = q('top');
    this.st = q('st');
    this.note = q('note');
    q<HTMLSelectElement>('comp').addEventListener('change', (e) => {
      this.comp = Number((e.target as HTMLSelectElement).value);
      this.drawSpaceTime();
    });
    this.st.addEventListener('click', (e) => {
      if (!this.f) return;
      const w = this.st.clientWidth;
      const fx = (e.offsetX - PAD.l) / (w - PAD.l - PAD.r);
      if (fx >= 0 && fx <= 1) tl.seek(Math.round(fx * (this.f.times.length - 1)));
    });
    tl.onChange(() => this.drawAll());
    const ro = new ResizeObserver(() => this.drawAll());
    for (const c of [this.side, this.top, this.st]) ro.observe(c);
  }

  // frames must match the timeline's times
  set(g: ChainGeometry, f: ChainFrames) {
    this.g = g;
    this.f = f;
    this.maxAbs = maxAbsPerComponent(f.frames);
    this.drawAll();
  }

  private drawAll() {
    this.drawFrame();
    this.drawSpaceTime();
  }

  private drawFrame() {
    const g = this.g;
    const f = this.f;
    if (!g || !f) return;
    const xs = [g.leftWall ?? g.restX[0], g.rightWall ?? g.restX[g.n - 1]];
    const L = xs[1] - xs[0] || 1;
    const xr: [number, number] = [xs[0] - 0.03 * L, xs[1] + 0.03 * L];
    const panelH = this.side.clientHeight - PAD.t - PAD.b;
    const pxPerM = (this.side.clientWidth - PAD.l - PAD.r) / (xr[1] - xr[0]);
    const tmax = Math.max(this.maxAbs[1], this.maxAbs[2], 1e-30);
    this.exaggeration = Math.max(1, (0.42 * panelH) / (tmax * pxPerM));
    // displacement range shown; at true scale a large motion may need more than the 1:1 aspect allows
    const half = Math.max((0.5 * panelH) / pxPerM / this.exaggeration, 1.15 * tmax);
    this.drawView(this.side, 1, xr, half);
    this.drawView(this.top, 2, xr, half);
    this.note.textContent =
      `Displacements (all components) drawn ×${fmtNum(this.exaggeration)}. Max |y| ${fmtNum(this.maxAbs[1] * 1000)} mm, ` +
      `max |z| ${fmtNum(this.maxAbs[2] * 1000)} mm, max |x| ${fmtNum(this.maxAbs[0] * 1000)} mm over the run.`;
  }

  private drawView(canvas: HTMLCanvasElement, comp: 1 | 2, xr: [number, number], half: number) {
    const g = this.g!;
    const f = this.f!;
    const { ctx, w, h } = setupCanvas(canvas);
    ctx.clearRect(0, 0, w, h);
    const { X, Y } = drawAxes(ctx, w, h, xr, [-half * 1000, half * 1000], {
      xLabel: 'x (m)',
      yLabel: comp === 1 ? 'y (mm)' : 'z (mm)',
    });
    const E = this.exaggeration;
    const off = this.frame * 3 * g.n;
    const pts: [number, number][] = [];
    const bead = (i: number): [number, number] => [
      X(g.restX[i] + E * f.frames[off + 3 * i]),
      Y(f.frames[off + 3 * i + comp] * 1000),
    ];
    if (g.leftWall !== null) pts.push([X(g.leftWall), g.slideLeft ? bead(0)[1] : Y(0)]);
    for (let i = 0; i < g.n; i++) pts.push(bead(i));
    if (g.rightWall !== null) pts.push([X(g.rightWall), g.slideRight ? bead(g.n - 1)[1] : Y(0)]);
    // walls
    ctx.strokeStyle = css('--text-muted');
    ctx.lineWidth = 2;
    for (const [wx, slide] of [[g.leftWall, g.slideLeft], [g.rightWall, g.slideRight]] as const) {
      if (wx === null) continue;
      ctx.setLineDash(slide ? [4, 3] : []);
      ctx.beginPath();
      ctx.moveTo(X(wx), PAD.t);
      ctx.lineTo(X(wx), h - PAD.b);
      ctx.stroke();
    }
    ctx.setLineDash([]);
    ctx.strokeStyle = css('--series-1');
    ctx.lineWidth = 1.5;
    ctx.beginPath();
    pts.forEach(([x, y], i) => (i ? ctx.lineTo(x, y) : ctx.moveTo(x, y)));
    ctx.stroke();
    const r = g.n > 80 ? 2 : g.n > 30 ? 3 : 4.5;
    const surface = css('--surface-1');
    for (let i = 0; i < g.n; i++) {
      const [x, y] = bead(i);
      ctx.beginPath();
      ctx.arc(x, y, g.pinned[i] ? r + 1.5 : r, 0, 2 * Math.PI);
      ctx.fillStyle = g.pinned[i] ? css('--text-primary') : css('--series-1');
      ctx.fill();
      ctx.lineWidth = 1.5;
      ctx.strokeStyle = surface;
      ctx.stroke();
    }
  }

  private drawSpaceTime() {
    const g = this.g;
    const f = this.f;
    if (!g || !f) return;
    const { ctx, w, h } = setupCanvas(this.st);
    ctx.clearRect(0, 0, w, h);
    const nf = f.times.length;
    const c = this.comp;
    const m = Math.max(this.maxAbs[c], 1e-30);
    const key = `${c}|${f.frames.length}|${css('--diverging')}`;
    if (!this.stImage || this.stKey !== key || this.stFrames !== f.frames) {
      const lut = rampLUT('--diverging');
      const img = new ImageData(nf, g.n);
      for (let k = 0; k < nf; k++) {
        for (let i = 0; i < g.n; i++) {
          const v = f.frames[k * 3 * g.n + 3 * i + c] / m;
          const li = Math.round((0.5 + 0.5 * Math.max(-1, Math.min(1, v))) * 255) * 3;
          const p = 4 * ((g.n - 1 - i) * nf + k);
          img.data[p] = lut[li];
          img.data[p + 1] = lut[li + 1];
          img.data[p + 2] = lut[li + 2];
          img.data[p + 3] = 255;
        }
      }
      const off = document.createElement('canvas');
      off.width = nf;
      off.height = g.n;
      off.getContext('2d')!.putImageData(img, 0, 0);
      this.stImage = off;
      this.stKey = key;
      this.stFrames = f.frames;
    }
    const off = this.stImage;
    ctx.imageSmoothingEnabled = false;
    ctx.drawImage(off, PAD.l, PAD.t, w - PAD.l - PAD.r, h - PAD.t - PAD.b);
    const t1 = f.times[nf - 1];
    const { X } = drawAxes(ctx, w, h, [0, t1], [0.5, g.n + 0.5], {
      xLabel: `time (s); color: ${['x', 'y', 'z'][c]} displacement, ±${fmtNum(m * 1000)} mm`,
      yLabel: 'bead',
      grid: false,
    });
    ctx.strokeStyle = css('--text-primary');
    ctx.lineWidth = 1;
    const cx = Math.round(X(f.times[this.frame])) + 0.5;
    ctx.beginPath();
    ctx.moveTo(cx, PAD.t);
    ctx.lineTo(cx, h - PAD.b);
    ctx.stroke();
  }
}

export function maxAbsPerComponent(frames: Float32Array): number[] {
  const m = [0, 0, 0];
  for (let i = 0; i < frames.length; i++) {
    const c = i % 3;
    m[c] = Math.max(m[c], Math.abs(frames[i]));
  }
  return m;
}