// Incremental STFT spectrogram drawn on a 2D canvas. The signal is kept as
// it arrives and turned into dB magnitude columns (one FFT per hop) as soon
// as a complete window is available. The plot area is an ImageData at
// device-pixel resolution whose pixel columns are max-pooled over the STFT
// columns and frequency bins they cover; on append only the pixel columns
// touched by the new STFT columns are recomputed, so draw() is a putImageData
// plus the axes and the cursor. The dB scale is relative to the largest
// magnitude seen so far (0 dB at the peak, down to -dbRange); when the peak
// rises the whole image is re-rendered from the stored columns.

import { fft, hann } from "./fft.ts";

export interface SpectrogramOptions {
  fs: number;
  window: number;
  hop: number;
  maxHz: number;
  dbRange: number;
  logFreq: boolean;
}

const LOG_FLOOR_HZ = 20;
const EMPTY_RGB = [26, 26, 32];
const FRAME_RGB = "#121216";
const AXIS_COLOR = "rgba(215, 215, 228, 0.85)";
const TICK_COLOR = "rgba(110, 225, 255, 0.95)";
const CURSOR_COLOR = "rgba(255, 255, 255, 0.85)";
// Re-render everything when the peak has risen by more than this since the last full render.
const PEAK_TOLERANCE_DB = 0.5;

export class Spectrogram {
  private readonly canvas: HTMLCanvasElement;
  private readonly ctx: CanvasRenderingContext2D;
  private opts: SpectrogramOptions;

  // Signal and STFT state.
  private signal = new Float32Array(1 << 16);
  private nSamples = 0;
  private consumed = 0;
  private columns: Float32Array[] = [];
  private peakDb = -Infinity;
  private win: Float32Array;
  private re: Float32Array;
  private im: Float32Array;

  // Display state.
  private duration = 1;
  private cursor: number | null = null;
  private ticks: number[] | null = null;

  // Geometry in device pixels.
  private dpr = 1;
  private W = 1;
  private H = 1;
  private left = 0;
  private top = 0;
  private pw = 1;
  private ph = 1;
  private fmin = 0;
  private fmax = 1;
  private image: ImageData | null = null;
  private renderedPeak = -Infinity;
  /** Per pixel row, the bin range [2*py, 2*py+1) pooled into it. */
  private rowBins = new Int32Array(0);
  /** Per pixel column, the STFT column range [2*px, 2*px+1) pooled into it. */
  private pixCols = new Int32Array(0);

  private raf = 0;
  private observer: ResizeObserver | null = null;
  private destroyed = false;

  constructor(canvas: HTMLCanvasElement, opts: SpectrogramOptions) {
    const ctx = canvas.getContext("2d");
    if (!ctx) throw new Error("Spectrogram: could not get a 2D canvas context");
    this.canvas = canvas;
    this.ctx = ctx;
    this.opts = { ...opts };
    validate(this.opts);
    this.win = hann(this.opts.window);
    this.re = new Float32Array(this.opts.window);
    this.im = new Float32Array(this.opts.window);
    if (typeof ResizeObserver !== "undefined") {
      this.observer = new ResizeObserver(() => this.resize());
      this.observer.observe(canvas);
    }
    this.resize();
  }

  /** Changes options; the columns are recomputed from the signal seen so far when the STFT parameters change. */
  setOptions(opts: Partial<SpectrogramOptions>): void {
    const next = { ...this.opts, ...opts };
    validate(next);
    const prev = this.opts;
    this.opts = next;
    if (next.window !== prev.window || next.hop !== prev.hop || next.fs !== prev.fs) {
      this.win = hann(next.window);
      this.re = new Float32Array(next.window);
      this.im = new Float32Array(next.window);
      this.columns = [];
      this.consumed = 0;
      this.peakDb = -Infinity;
      this.computeColumns();
    }
    this.rebuildGeometry();
    this.scheduleDraw();
  }

  /** Clears the signal and the image. */
  reset(): void {
    this.nSamples = 0;
    this.consumed = 0;
    this.columns = [];
    this.peakDb = -Infinity;
    this.renderAll();
    this.scheduleDraw();
  }

  /** Appends signal; any new complete columns are computed and rendered incrementally. */
  append(samples: Float32Array): void {
    if (samples.length === 0) return;
    const need = this.nSamples + samples.length;
    if (need > this.signal.length) {
      let cap = this.signal.length;
      while (cap < need) cap *= 2;
      const grown = new Float32Array(cap);
      grown.set(this.signal.subarray(0, this.nSamples));
      this.signal = grown;
    }
    this.signal.set(samples, this.nSamples);
    this.nSamples = need;
    const before = this.columns.length;
    this.computeColumns();
    if (this.columns.length > before) this.renderNew(before);
    this.scheduleDraw();
  }

  /** Replaces the whole signal. */
  setSignal(samples: Float32Array): void {
    this.nSamples = 0;
    this.consumed = 0;
    this.columns = [];
    this.peakDb = -Infinity;
    if (samples.length === 0) {
      this.renderAll();
      this.scheduleDraw();
      return;
    }
    this.append(samples);
  }

  /** Playback cursor [s], drawn as a vertical line; null hides it. */
  setCursor(seconds: number | null): void {
    if (seconds === this.cursor) return;
    this.cursor = seconds;
    this.scheduleDraw();
  }

  /** Short horizontal ticks at these frequencies [Hz] on the left edge of the plot (e.g. modal frequencies). */
  setTicks(freqs: number[] | null): void {
    this.ticks = freqs ? freqs.slice() : null;
    this.scheduleDraw();
  }

  /** The time axis spans this duration [s], so the image does not rescale as the signal grows. */
  setDuration(seconds: number): void {
    const d = seconds > 0 && Number.isFinite(seconds) ? seconds : 1;
    if (d === this.duration) return;
    this.duration = d;
    this.rebuildGeometry();
    this.scheduleDraw();
  }

  /** Matches the canvas pixel size to its CSS size and the device pixel ratio, then re-renders. */
  resize(): void {
    if (this.destroyed) return;
    const dpr = (typeof window !== "undefined" && window.devicePixelRatio) || 1;
    const cssW = this.canvas.clientWidth || this.canvas.width / dpr || 1;
    const cssH = this.canvas.clientHeight || this.canvas.height / dpr || 1;
    const W = Math.max(1, Math.round(cssW * dpr));
    const H = Math.max(1, Math.round(cssH * dpr));
    if (this.canvas.width !== W) this.canvas.width = W;
    if (this.canvas.height !== H) this.canvas.height = H;
    if (W === this.W && H === this.H && dpr === this.dpr && this.image) {
      this.scheduleDraw();
      return;
    }
    this.dpr = dpr;
    this.W = W;
    this.H = H;
    this.rebuildGeometry();
    this.scheduleDraw();
  }

  /** Draws the current image, axes, ticks and cursor immediately. */
  draw(): void {
    if (this.raf) {
      cancelAnimationFrame(this.raf);
      this.raf = 0;
    }
    if (this.destroyed || !this.image) return;
    const ctx = this.ctx;
    const { dpr, left, top, pw, ph } = this;
    ctx.setTransform(1, 0, 0, 1, 0, 0);
    ctx.fillStyle = FRAME_RGB;
    ctx.fillRect(0, 0, this.W, this.H);
    ctx.putImageData(this.image, left, top);
    this.drawAxes();
    if (this.ticks) {
      ctx.fillStyle = TICK_COLOR;
      const len = Math.round(6 * dpr);
      const thick = Math.max(1, Math.round(dpr));
      for (const f of this.ticks) {
        if (!(f >= this.fmin && f <= this.fmax)) continue;
        const y = Math.round(top + (1 - this.freqFrac(f)) * ph - thick / 2);
        ctx.fillRect(left, y, len, thick);
      }
    }
    const c = this.cursor;
    if (c !== null && c >= 0 && c <= this.duration) {
      const x = Math.round(left + (c / this.duration) * pw);
      ctx.fillStyle = CURSOR_COLOR;
      ctx.fillRect(x, top, Math.max(1, Math.round(dpr)), ph);
    }
  }

  destroy(): void {
    this.destroyed = true;
    if (this.raf) {
      cancelAnimationFrame(this.raf);
      this.raf = 0;
    }
    this.observer?.disconnect();
    this.observer = null;
    this.columns = [];
    this.signal = new Float32Array(0);
    this.image = null;
  }

  // ---- STFT ----

  private computeColumns(): void {
    const { window: n, hop } = this.opts;
    const nBins = (n >> 1) + 1;
    const re = this.re;
    const im = this.im;
    const win = this.win;
    while (this.consumed + n <= this.nSamples) {
      const off = this.consumed;
      for (let k = 0; k < n; k++) {
        re[k] = this.signal[off + k] * win[k];
        im[k] = 0;
      }
      fft(re, im);
      const col = new Float32Array(nBins);
      let peak = this.peakDb;
      for (let k = 0; k < nBins; k++) {
        const p = re[k] * re[k] + im[k] * im[k];
        const db = 10 * Math.log10(p + 1e-30);
        col[k] = db;
        if (db > peak) peak = db;
      }
      this.peakDb = peak;
      this.columns.push(col);
      this.consumed += hop;
    }
  }

  // ---- Geometry and rendering of the plot image ----

  private freqFrac(f: number): number {
    if (this.opts.logFreq) {
      return (Math.log(f) - Math.log(this.fmin)) / (Math.log(this.fmax) - Math.log(this.fmin));
    }
    return (f - this.fmin) / (this.fmax - this.fmin);
  }

  private fracFreq(frac: number): number {
    if (this.opts.logFreq) {
      const lo = Math.log(this.fmin);
      const hi = Math.log(this.fmax);
      return Math.exp(lo + frac * (hi - lo));
    }
    return this.fmin + frac * (this.fmax - this.fmin);
  }

  private rebuildGeometry(): void {
    const { dpr, W, H } = this;
    const o = this.opts;
    this.left = Math.round(50 * dpr);
    this.top = Math.round(6 * dpr);
    const bottom = Math.round(20 * dpr);
    const right = Math.round(8 * dpr);
    this.pw = Math.max(1, W - this.left - right);
    this.ph = Math.max(1, H - this.top - bottom);
    const nyquist = o.fs / 2;
    if (o.logFreq) {
      this.fmin = LOG_FLOOR_HZ;
      this.fmax = Math.max(LOG_FLOOR_HZ * 2, Math.min(o.maxHz, nyquist));
    } else {
      this.fmin = 0;
      this.fmax = Math.max(1, Math.min(o.maxHz, nyquist));
    }

    // Bin range per pixel row (row 0 at the top).
    const binHz = o.fs / o.window;
    const nBins = (o.window >> 1) + 1;
    const ph = this.ph;
    this.rowBins = new Int32Array(2 * ph);
    for (let py = 0; py < ph; py++) {
      const fHi = this.fracFreq(1 - py / ph);
      const fLo = this.fracFreq(1 - (py + 1) / ph);
      let b0 = Math.round(fLo / binHz);
      let b1 = Math.round(fHi / binHz);
      b0 = Math.max(0, Math.min(nBins - 1, b0));
      b1 = Math.max(b0 + 1, Math.min(nBins, b1));
      this.rowBins[2 * py] = b0;
      this.rowBins[2 * py + 1] = b1;
    }

    // STFT column range per pixel column. Column c is centered at time
    // (c * hop + window / 2) / fs; a pixel narrower than a hop takes the
    // nearest column.
    const pw = this.pw;
    const dt = this.duration / pw; // seconds per pixel
    const half = o.window / 2;
    const maxCols = Math.max(1, Math.floor((this.duration * o.fs - o.window) / o.hop) + 1);
    this.pixCols = new Int32Array(2 * pw);
    for (let px = 0; px < pw; px++) {
      let c0 = Math.ceil((px * dt * o.fs - half) / o.hop);
      let c1 = Math.ceil(((px + 1) * dt * o.fs - half) / o.hop);
      if (c1 <= c0) {
        c0 = Math.round(((px + 0.5) * dt * o.fs - half) / o.hop);
        c1 = c0 + 1;
      }
      c0 = Math.max(0, Math.min(maxCols - 1, c0));
      c1 = Math.max(c0 + 1, c1);
      this.pixCols[2 * px] = c0;
      this.pixCols[2 * px + 1] = c1;
    }

    this.image = this.ctx.createImageData(pw, ph);
    this.renderAll();
  }

  private renderAll(): void {
    if (!this.image) return;
    this.renderedPeak = this.peakDb;
    for (let px = 0; px < this.pw; px++) this.renderPixelColumn(px);
  }

  /** Renders the pixel columns touched by STFT columns [firstNew, columns.length). */
  private renderNew(firstNew: number): void {
    if (!this.image) return;
    if (this.peakDb > this.renderedPeak + PEAK_TOLERANCE_DB || !Number.isFinite(this.renderedPeak)) {
      this.renderAll();
      return;
    }
    const last = this.columns.length;
    for (let px = 0; px < this.pw; px++) {
      const c0 = this.pixCols[2 * px];
      const c1 = this.pixCols[2 * px + 1];
      if (c1 > firstNew && c0 < last) this.renderPixelColumn(px);
    }
  }

  private renderPixelColumn(px: number): void {
    const img = this.image;
    if (!img) return;
    const data = img.data;
    const { pw, ph } = this;
    const cols = this.columns;
    const c0 = this.pixCols[2 * px];
    const c1 = Math.min(cols.length, this.pixCols[2 * px + 1]);
    if (c0 >= c1) {
      for (let py = 0; py < ph; py++) {
        const o = (py * pw + px) * 4;
        data[o] = EMPTY_RGB[0];
        data[o + 1] = EMPTY_RGB[1];
        data[o + 2] = EMPTY_RGB[2];
        data[o + 3] = 255;
      }
      return;
    }
    const ref = this.renderedPeak;
    const range = this.opts.dbRange;
    const lut = colormapLut();
    for (let py = 0; py < ph; py++) {
      const b0 = this.rowBins[2 * py];
      const b1 = this.rowBins[2 * py + 1];
      let v = -Infinity;
      for (let c = c0; c < c1; c++) {
        const col = cols[c];
        for (let b = b0; b < b1; b++) if (col[b] > v) v = col[b];
      }
      let t = (v - ref + range) / range;
      if (!(t > 0)) t = 0;
      else if (t > 1) t = 1;
      const i = Math.round(t * 255) * 3;
      const o = (py * pw + px) * 4;
      data[o] = lut[i];
      data[o + 1] = lut[i + 1];
      data[o + 2] = lut[i + 2];
      data[o + 3] = 255;
    }
  }

  // ---- Axes ----

  private drawAxes(): void {
    const ctx = this.ctx;
    const { dpr, left, top, pw, ph, fmin, fmax } = this;
    ctx.fillStyle = AXIS_COLOR;
    ctx.font = `${Math.round(10 * dpr)}px system-ui, -apple-system, "Segoe UI", Helvetica, Arial, sans-serif`;
    ctx.textAlign = "right";
    ctx.textBaseline = "middle";
    const fticks = this.opts.logFreq
      ? LOG_TICKS.filter((f) => f >= fmin && f <= fmax)
      : niceTicks(0, fmax, Math.max(2, Math.floor(ph / (36 * dpr))));
    const thin = Math.max(1, Math.round(dpr));
    let lastY = -Infinity;
    for (const f of fticks) {
      const y = Math.round(top + (1 - this.freqFrac(f)) * ph);
      if (Math.abs(y - lastY) < 12 * dpr) continue;
      lastY = y;
      ctx.fillRect(left - 4 * dpr, y, 4 * dpr, thin);
      ctx.fillText(formatHz(f), left - 6 * dpr, y);
    }
    ctx.textAlign = "center";
    ctx.textBaseline = "top";
    const tticks = niceTicks(0, this.duration, Math.max(2, Math.floor(pw / (64 * dpr))));
    const yText = top + ph + 4 * dpr;
    for (let k = 0; k < tticks.length; k++) {
      const t = tticks[k];
      const x = Math.round(left + (t / this.duration) * pw);
      ctx.fillRect(x, top + ph, thin, 4 * dpr);
      let label = formatNumber(t);
      if (k === tticks.length - 1) label += " s";
      ctx.fillText(label, x, yText);
    }
  }

  private scheduleDraw(): void {
    if (this.destroyed || this.raf) return;
    if (typeof requestAnimationFrame === "undefined") {
      this.draw();
      return;
    }
    this.raf = requestAnimationFrame(() => {
      this.raf = 0;
      this.draw();
    });
  }
}

const LOG_TICKS = [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000];

function validate(o: SpectrogramOptions): void {
  if (!(o.window >= 2) || (o.window & (o.window - 1)) !== 0) throw new Error(`Spectrogram: window ${o.window} is not a power of two`);
  if (!(o.hop >= 1)) throw new Error(`Spectrogram: hop must be positive`);
  if (!(o.fs > 0)) throw new Error(`Spectrogram: fs must be positive`);
  if (!(o.dbRange > 0)) throw new Error(`Spectrogram: dbRange must be positive`);
  o.window = Math.round(o.window);
  o.hop = Math.round(o.hop);
}

function formatHz(f: number): string {
  if (f >= 1000) return `${formatNumber(f / 1000)} kHz`;
  return `${formatNumber(f)} Hz`;
}

function formatNumber(x: number): string {
  const s = x.toFixed(3);
  return s.replace(/\.?0+$/, "");
}

/** Round tick positions covering [lo, hi] with about `count` intervals. */
function niceTicks(lo: number, hi: number, count: number): number[] {
  const span = hi - lo;
  if (!(span > 0)) return [lo];
  const raw = span / count;
  const mag = 10 ** Math.floor(Math.log10(raw));
  const norm = raw / mag;
  const step = (norm < 1.5 ? 1 : norm < 3.5 ? 2 : norm < 7.5 ? 5 : 10) * mag;
  const ticks: number[] = [];
  const first = Math.ceil(lo / step - 1e-9);
  const last = Math.floor(hi / step + 1e-9);
  for (let k = first; k <= last; k++) ticks.push(+(k * step).toPrecision(12));
  return ticks;
}

// Inferno-like colormap: anchor colors at t = 0, 0.1, ..., 1 interpolated
// linearly into a 256-entry RGB table.
const INFERNO_ANCHORS: [number, number, number][] = [
  [0, 0, 4],
  [22, 11, 57],
  [66, 10, 104],
  [106, 23, 110],
  [147, 38, 103],
  [188, 55, 84],
  [221, 81, 58],
  [243, 120, 25],
  [252, 165, 10],
  [246, 215, 70],
  [252, 255, 164],
];

let lut: Uint8ClampedArray | null = null;

function colormapLut(): Uint8ClampedArray {
  if (lut) return lut;
  const table = new Uint8ClampedArray(256 * 3);
  const segs = INFERNO_ANCHORS.length - 1;
  for (let i = 0; i < 256; i++) {
    const t = (i / 255) * segs;
    const s = Math.min(segs - 1, Math.floor(t));
    const u = t - s;
    const a = INFERNO_ANCHORS[s];
    const b = INFERNO_ANCHORS[s + 1];
    for (let k = 0; k < 3; k++) table[i * 3 + k] = Math.round(a[k] + (b[k] - a[k]) * u);
  }
  lut = table;
  return table;
}