// Spectrogram canvas: log-frequency rows, dB colour scale relative to the loudest column seen,
// a time axis spanning the whole run, the frontier of the generated signal and a playback cursor.
import type { Stft } from "../analysis/spectrogram.ts";

const ROWS = 256, DB_RANGE = 80, F_MIN = 40;
export type FrequencyScale = "linear" | "log";
const STOPS = [[0.0, 13, 8, 60], [0.25, 40, 40, 130], [0.5, 30, 120, 140], [0.75, 90, 200, 90], [1.0, 250, 240, 80]];

function color(t: number): [number, number, number] {
  const x = Math.min(Math.max(t, 0), 1);
  for (let i = 1; i < STOPS.length; i++) {
    if (x <= STOPS[i][0]) {
      const [t0, r0, g0, b0] = STOPS[i - 1], [t1, r1, g1, b1] = STOPS[i], u = (x - t0) / (t1 - t0);
      return [r0 + u * (r1 - r0), g0 + u * (g1 - g0), b0 + u * (b1 - b0)];
    }
  }
  return [250, 240, 80];
}

export class SpectrogramView {
  private readonly canvas: HTMLCanvasElement;
  private readonly off = document.createElement("canvas");
  private stft: Stft | null = null;
  private duration = 1; private drawn = 0; private maxPower = 0;
  private rowLo: Int32Array = new Int32Array(0); private rowHi: Int32Array = new Int32Array(0);
  cursorTime: number | null = null; frontierTime = 0;
  private scale: FrequencyScale = "linear";
  private fMaxShown = 0;   // 0 = up to Nyquist

  constructor(canvas: HTMLCanvasElement) { this.canvas = canvas; }

  /** Frequency axis; changing it redraws the stored columns. */
  setScale(scale: FrequencyScale) {
    if (scale === this.scale) return;
    this.scale = scale;
    if (this.stft) { this.buildRows(); this.drawn = 0; this.update(); } else this.draw();
  }
  get frequencyScale() { return this.scale; }

  /** Upper edge of the axis in Hz; 0 restores the Nyquist frequency. */
  setMaxFrequency(f: number) {
    this.fMaxShown = f;
    if (this.stft) { this.buildRows(); this.drawn = 0; this.update(); } else this.draw();
  }

  private get fMax() { return this.fMaxShown > 0 ? this.fMaxShown : this.stft ? this.stft.fs / 2 : 22050; }
  private get fMin() { return this.scale === "log" ? F_MIN : 0; }

  /** Frequency at the top edge of row r (r = 0 is the top of the image). */
  private rowFrequency(r: number): number {
    const u = 1 - r / ROWS;
    return this.scale === "log" ? this.fMin * Math.pow(this.fMax / this.fMin, u) : this.fMax * u;
  }

  private buildRows() {
    const st = this.stft!;
    this.rowLo = new Int32Array(ROWS); this.rowHi = new Int32Array(ROWS);
    const df = st.binFrequency(1);
    for (let r = 0; r < ROWS; r++) {
      const fHi = this.rowFrequency(r), fLo = this.rowFrequency(r + 1);
      this.rowLo[r] = Math.max(0, Math.floor(fLo / df));
      this.rowHi[r] = Math.max(this.rowLo[r] + 1, Math.ceil(fHi / df));
    }
  }

  reset(stft: Stft, duration: number) {
    this.stft = stft; this.duration = duration; this.drawn = 0; this.maxPower = 0; this.frontierTime = 0; this.cursorTime = null;
    this.off.width = Math.max(1, Math.ceil((duration * stft.fs) / stft.hop)); this.off.height = ROWS;
    this.off.getContext("2d")!.fillStyle = "#0d0838"; this.off.getContext("2d")!.fillRect(0, 0, this.off.width, ROWS);
    this.buildRows();
    this.draw();
  }

  /** Draw the columns computed since the last call; rescales everything when the maximum grows. */
  update() {
    const st = this.stft;
    if (!st) return;
    let mx = this.maxPower;
    for (let c = this.drawn; c < st.columns.length; c++) for (const p of st.columns[c]) mx = Math.max(mx, p);
    if (mx > this.maxPower * 2) { this.maxPower = mx; this.drawn = 0; } else this.maxPower = Math.max(this.maxPower, mx);
    const ctx = this.off.getContext("2d")!;
    const ref = this.maxPower || 1;
    for (let c = this.drawn; c < st.columns.length && c < this.off.width; c++) {
      const img = ctx.createImageData(1, ROWS), col = st.columns[c];
      for (let r = 0; r < ROWS; r++) {
        let p = 0;
        for (let b = this.rowLo[r]; b < Math.min(this.rowHi[r], col.length); b++) p = Math.max(p, col[b]);
        const db = 10 * Math.log10(p / ref + 1e-30);
        const [R, G, B] = color(1 + db / DB_RANGE);
        img.data[4 * r] = R; img.data[4 * r + 1] = G; img.data[4 * r + 2] = B; img.data[4 * r + 3] = 255;
      }
      ctx.putImageData(img, c, 0);
    }
    this.drawn = Math.min(st.columns.length, this.off.width);
    this.draw();
  }

  draw() {
    const cv = this.canvas, dpr = window.devicePixelRatio || 1;
    const w = cv.clientWidth, h = cv.clientHeight;
    if (cv.width !== Math.round(w * dpr) || cv.height !== Math.round(h * dpr)) { cv.width = Math.round(w * dpr); cv.height = Math.round(h * dpr); }
    const ctx = cv.getContext("2d")!;
    ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
    ctx.imageSmoothingEnabled = false;
    ctx.fillStyle = "#0d0838"; ctx.fillRect(0, 0, w, h);
    const left = 44, bottom = 18, pw = w - left - 6, ph = h - bottom - 4;
    if (this.stft) ctx.drawImage(this.off, 0, 0, this.off.width, ROWS, left, 4, pw, ph);
    ctx.fillStyle = "#c8c8c8"; ctx.font = "10px system-ui, sans-serif"; ctx.textAlign = "right";
    const fMax = this.fMax;
    const ticks = this.scale === "log"
      ? [50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000]
      : (() => {
          const step = [500, 1000, 2000, 2500, 5000, 10000].find((s) => fMax / s <= 11) ?? 10000;
          const out: number[] = [];
          for (let f = 0; f <= fMax + 1e-6; f += step) out.push(f);
          return out;
        })();
    for (const f of ticks) {
      if (f > fMax || (this.scale === "log" && f < F_MIN)) continue;
      const u = this.scale === "log" ? Math.log(f / this.fMin) / Math.log(fMax / this.fMin) : f / fMax;
      const y = 4 + ph * (1 - u);
      ctx.fillText(f === 0 ? "0" : f >= 1000 ? `${f / 1000}k` : String(f), left - 4, y + 3);
      if (f > 0) { ctx.fillStyle = "rgba(255,255,255,0.12)"; ctx.fillRect(left, y, pw, 1); ctx.fillStyle = "#c8c8c8"; }
    }
    ctx.textAlign = "center";
    const tick = this.duration > 4 ? 1 : this.duration > 1.5 ? 0.5 : 0.1;
    for (let t = 0; t <= this.duration + 1e-9; t += tick) {
      const x = left + (pw * t) / this.duration;
      ctx.fillText(`${t.toFixed(tick < 0.5 ? 1 : 1)}`, x, h - 5);
    }
    const xf = left + (pw * this.frontierTime) / this.duration;
    ctx.fillStyle = "rgba(255,255,255,0.5)"; ctx.fillRect(xf, 4, 1, ph);
    if (this.cursorTime !== null) { ctx.fillStyle = "#ff5a5a"; ctx.fillRect(left + (pw * this.cursorTime) / this.duration, 4, 2, ph); }
  }
}