// Incremental STFT spectrogram drawn on a 2D canvas.

import { FFT } from "./fft.ts";

export interface SpectrogramOptions {
  windowSize: number;
  hop: number;
  fmax: number; // Hz
  dbFloor: number; // dB below the peak
  logFrequency: boolean;
}

export class Spectrogram {
  private fft: FFT;
  private window: Float32Array;
  private re: Float32Array;
  private im: Float32Array;
  /** Magnitude (dB) columns, each of length windowSize / 2. */
  private columns: Float32Array[] = [];
  private consumed = 0;
  private peakDb = -Infinity;
  private opts: SpectrogramOptions;

  constructor(opts: SpectrogramOptions) {
    this.opts = opts;
    this.fft = new FFT(opts.windowSize);
    this.window = new Float32Array(opts.windowSize);
    for (let k = 0; k < opts.windowSize; k++) this.window[k] = 0.5 - 0.5 * Math.cos((2 * Math.PI * k) / opts.windowSize);
    this.re = new Float32Array(opts.windowSize);
    this.im = new Float32Array(opts.windowSize);
  }

  get options(): SpectrogramOptions {
    return this.opts;
  }

  /** Changes options; recomputes from scratch when the STFT parameters change. */
  configure(opts: Partial<SpectrogramOptions>, audio: Float32Array, nSamples: number): void {
    const next = { ...this.opts, ...opts };
    const recompute = next.windowSize !== this.opts.windowSize || next.hop !== this.opts.hop;
    this.opts = next;
    if (recompute) {
      this.fft = new FFT(next.windowSize);
      this.window = new Float32Array(next.windowSize);
      for (let k = 0; k < next.windowSize; k++) this.window[k] = 0.5 - 0.5 * Math.cos((2 * Math.PI * k) / next.windowSize);
      this.re = new Float32Array(next.windowSize);
      this.im = new Float32Array(next.windowSize);
      this.reset();
      this.update(audio, nSamples);
    }
  }

  reset(): void {
    this.columns = [];
    this.consumed = 0;
    this.peakDb = -Infinity;
  }

  /** Adds columns for all complete windows in audio[0..nSamples). */
  update(audio: Float32Array, nSamples: number): void {
    const { windowSize, hop } = this.opts;
    while (this.consumed + windowSize <= nSamples) {
      for (let k = 0; k < windowSize; k++) {
        this.re[k] = audio[this.consumed + k] * this.window[k];
        this.im[k] = 0;
      }
      this.fft.transform(this.re, this.im);
      const col = new Float32Array(windowSize / 2);
      for (let k = 0; k < windowSize / 2; k++) {
        const mag2 = this.re[k] * this.re[k] + this.im[k] * this.im[k];
        const db = 10 * Math.log10(mag2 + 1e-30);
        col[k] = db;
        if (db > this.peakDb) this.peakDb = db;
      }
      this.columns.push(col);
      this.consumed += hop;
    }
  }

  get columnCount(): number {
    return this.columns.length;
  }

  private imageKey = "";
  private image: ImageData | null = null;
  private drawKey = "";

  /** Forces the next draw to recompute everything. */
  invalidate(): void {
    this.imageKey = "";
    this.drawKey = "";
  }

  /**
   * Draws the spectrogram; `duration` [s] fixes the time axis, `playhead` [s] draws a cursor.
   * The plot image is cached and only recomputed when the data, options, or size change,
   * and nothing is drawn at all when the previous frame is still current.
   */
  draw(canvas: HTMLCanvasElement, fs: number, duration: number, playhead: number | null): void {
    const dpr = window.devicePixelRatio || 1;
    const W = Math.max(1, Math.floor(canvas.clientWidth * dpr));
    const H = Math.max(1, Math.floor(canvas.clientHeight * dpr));
    const o = this.opts;
    const imageKey = `${W}|${H}|${this.columns.length}|${duration}|${fs}|${o.windowSize}|${o.hop}|${o.fmax}|${o.dbFloor}|${o.logFrequency}`;
    const drawKey = `${imageKey}|${playhead === null ? "-" : playhead.toFixed(4)}`;
    if (drawKey === this.drawKey) return;
    this.drawKey = drawKey;
    if (canvas.width !== W || canvas.height !== H) {
      canvas.width = W;
      canvas.height = H;
    }
    const ctx = canvas.getContext("2d")!;
    // Whole pixels throughout: a fractional device pixel ratio would otherwise
    // make the ImageData and the pixel buffer disagree in size.
    const left = Math.round(52 * dpr), bottom = Math.round(22 * dpr), top = Math.round(6 * dpr), right = Math.round(8 * dpr);
    const pw = Math.max(1, Math.floor(W - left - right)), ph = Math.max(1, Math.floor(H - top - bottom));
    ctx.fillStyle = "#1c1e22";
    ctx.fillRect(0, 0, W, H);
    const { fmax, logFrequency } = this.opts;
    const fmin = logFrequency ? 20 : 0;
    const logMin = Math.log(fmin || 1), logMax = Math.log(fmax);
    if (imageKey !== this.imageKey || !this.image) {
      this.image = this.renderImage(ctx, pw, ph, fs, duration);
      this.imageKey = imageKey;
    }
    ctx.putImageData(this.image, left, top);
    this.drawAxes(ctx, dpr, left, top, pw, ph, fmin, fmax, logMin, logMax, logFrequency, duration);
    if (playhead !== null && playhead >= 0 && playhead <= duration) {
      const x = left + (playhead / duration) * pw;
      ctx.fillStyle = "rgba(255,255,255,0.8)";
      ctx.fillRect(x, top, Math.max(1, dpr), ph);
    }
  }

  private renderImage(ctx: CanvasRenderingContext2D, pw: number, ph: number, fs: number, duration: number): ImageData {
    const img = ctx.createImageData(pw, ph);
    img.data.set(spectrogramPixels(this.columns, this.peakDb, this.opts, fs, duration, pw, ph));
    return img;
  }

  /** STFT columns (dB) and their peak, for renderers outside the browser. */
  get data(): { columns: Float32Array[]; peakDb: number } {
    return { columns: this.columns, peakDb: this.peakDb };
  }

  private drawAxes(ctx: CanvasRenderingContext2D, dpr: number, left: number, top: number, pw: number, ph: number, fmin: number, fmax: number, logMin: number, logMax: number, logFrequency: boolean, duration: number): void {
    ctx.fillStyle = "#c8cbd0";
    ctx.font = `${11 * dpr}px system-ui, sans-serif`;
    ctx.textAlign = "right";
    ctx.textBaseline = "middle";
    const fticks = logFrequency ? [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000].filter((f) => f >= fmin && f <= fmax) : niceTicks(0, fmax, 6);
    for (const f of fticks) {
      const frac = logFrequency ? (Math.log(f) - logMin) / (logMax - logMin) : (f - fmin) / (fmax - fmin);
      const y = top + (1 - frac) * ph;
      ctx.fillText(f >= 1000 ? `${f / 1000}k` : `${f}`, left - 4 * dpr, y);
      ctx.fillRect(left - 3 * dpr, y, 3 * dpr, 1);
    }
    ctx.textAlign = "center";
    ctx.textBaseline = "top";
    for (const t of niceTicks(0, duration, 8)) {
      const x = left + (t / duration) * pw;
      ctx.fillText(`${+t.toFixed(2)}`, x, top + ph + 4 * dpr);
      ctx.fillRect(x, top + ph, 1, 3 * dpr);
    }
    ctx.save();
    ctx.translate(12 * dpr, top + ph / 2);
    ctx.rotate(-Math.PI / 2);
    ctx.textBaseline = "middle";
    ctx.fillText("Hz", 0, 0);
    ctx.restore();
  }
}

/** RGBA pixels (row-major, top row first) of the spectrogram plot area. */
export function spectrogramPixels(columns: Float32Array[], peakDb: number, opts: SpectrogramOptions, fs: number, duration: number, pw: number, ph: number, peakRef?: number): Uint8ClampedArray {
  const { windowSize, hop, fmax, dbFloor, logFrequency } = opts;
  const binHz = fs / windowSize;
  const fmin = logFrequency ? 20 : 0;
  const data = new Uint8ClampedArray(pw * ph * 4);
  const cols = columns.length;
  const colPerPx = (duration * fs) / hop / pw;
  const logMin = Math.log(fmin || 1), logMax = Math.log(fmax);
  const ref = peakRef ?? peakDb;
  for (let py = 0; py < ph; py++) {
    const frac = 1 - py / ph;
    const f = logFrequency ? Math.exp(logMin + frac * (logMax - logMin)) : fmin + frac * (fmax - fmin);
    const bin = Math.min(windowSize / 2 - 1, Math.max(0, Math.round(f / binHz)));
    for (let px = 0; px < pw; px++) {
      const c0 = Math.floor(px * colPerPx);
      const c1 = Math.max(c0 + 1, Math.floor((px + 1) * colPerPx));
      let v = -Infinity;
      for (let c = c0; c < c1 && c < cols; c++) v = Math.max(v, columns[c][bin]);
      const o = (py * pw + px) * 4;
      if (c0 >= cols) {
        data[o] = 28; data[o + 1] = 30; data[o + 2] = 34; data[o + 3] = 255;
        continue;
      }
      const t = ref > -200 ? Math.max(0, Math.min(1, (v - ref + dbFloor) / dbFloor)) : 0;
      const [r, g, b] = inferno(t);
      data[o] = r; data[o + 1] = g; data[o + 2] = b; data[o + 3] = 255;
    }
  }
  return data;
}

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[] = [];
  for (let t = Math.ceil(lo / step) * step; t <= hi + 1e-9; t += step) ticks.push(t);
  return ticks;
}

// Compact polynomial approximation of the inferno colormap.
function inferno(t: number): [number, number, number] {
  const c0 = [0.0002189403691192265, 0.001651004631001012, -0.01948089843709184];
  const c1 = [0.1065134194856116, 0.5639564367884091, 3.932712388889277];
  const c2 = [11.60249308247187, -3.972853965665698, -15.9423940062914];
  const c3 = [-41.70399613139459, 17.43639888205313, 44.35414519872813];
  const c4 = [77.162935699427, -33.40235894210092, -81.80730925738993];
  const c5 = [-71.31942824499214, 32.62606426397723, 73.20951985803202];
  const c6 = [25.13112622477341, -12.24266895238567, -23.07032500287172];
  const out: number[] = [];
  for (let k = 0; k < 3; k++) {
    const v = c0[k] + t * (c1[k] + t * (c2[k] + t * (c3[k] + t * (c4[k] + t * (c5[k] + t * c6[k])))));
    out.push(Math.max(0, Math.min(255, Math.round(v * 255))));
  }
  return [out[0], out[1], out[2]];
}