// The modal views of a run. The spectrogram is time against mode frequency with colour the modal
// energy in dB below the loudest value seen so far: each mode is one horizontal trace, so energy
// arriving in a mode is visible as its trace lighting up, resolved per mode rather than smeared
// by an STFT window. The bar view is the energy of every mode at the cursor time. Both share the
// frequency axis (linear or logarithmic, with a selectable top) and the dB range, and clicking
// near a trace or a bar selects that mode.
import type { Stft } from "../analysis/spectrogram.ts";
import type { ModalTracker } from "../modal/tracker.ts";

export type FrequencyScale = "linear" | "log";

/** Perceptual-ish dark-to-bright colour ramp for t in [0, 1]. */
function rampColor(t: number): [number, number, number] {
  const stops: [number, number, number][] = [
    [12, 13, 18], [48, 26, 74], [120, 35, 107], [193, 62, 88], [240, 122, 52], [252, 200, 96], [252, 252, 200],
  ];
  const x = Math.min(1, Math.max(0, t)) * (stops.length - 1), i = Math.min(stops.length - 2, Math.floor(x)), f = x - i;
  return [0, 1, 2].map((k) => Math.round(stops[i][k] + f * (stops[i + 1][k] - stops[i][k]))) as [number, number, number];
}

class FrequencyAxis {
  scale: FrequencyScale = "linear";
  fmax = 20000;
  fmin = 30;
  /** Pixel y (0 at the top) of frequency f within a plot of height H. */
  y(f: number, H: number): number {
    if (this.scale === "linear") return H * (1 - f / this.fmax);
    const l0 = Math.log(this.fmin), l1 = Math.log(this.fmax);
    return H * (1 - (Math.log(Math.max(f, this.fmin)) - l0) / (l1 - l0));
  }
  /** Frequency at pixel y, the inverse of ``y``. */
  f(y: number, H: number): number {
    if (this.scale === "linear") return this.fmax * (1 - y / H);
    const l0 = Math.log(this.fmin), l1 = Math.log(this.fmax);
    return Math.exp(l0 + (1 - y / H) * (l1 - l0));
  }
  ticks(): number[] {
    if (this.scale === "linear") {
      const step = this.fmax > 10000 ? 2000 : this.fmax > 4000 ? 1000 : this.fmax > 1500 ? 500 : 200;
      const out = [];
      for (let f = step; f <= this.fmax; f += step) out.push(f);
      return out;
    }
    return [50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000].filter((f) => f >= this.fmin && f <= this.fmax);
  }
}

const LEFT = 44, BOTTOM = 16;

export class ModalSpectrogram {
  private readonly canvas: HTMLCanvasElement;
  private readonly image = document.createElement("canvas");
  readonly axis = new FrequencyAxis();
  dbRange = 80;
  tracker: ModalTracker | null = null;
  duration = 1;                 // x extent [s]
  cursorTime: number | null = null;
  selectedMode = -1;
  onSelect: ((mode: number, time: number) => void) | null = null;
  onHover: ((text: string) => void) | null = null;
  private dirty = true;

  constructor(canvas: HTMLCanvasElement) {
    this.canvas = canvas;
    canvas.addEventListener("pointerdown", (e) => {
      const pick = this.pick(e);
      if (pick) this.onSelect?.(pick.mode, pick.time);
    });
    canvas.addEventListener("pointermove", (e) => {
      const pick = this.pick(e);
      if (!pick || !this.tracker) { this.onHover?.(""); return; }
      const s = this.tracker.indexAtTime(pick.time);
      const E = s >= 0 ? this.tracker.energies[s][pick.mode] : 0;
      this.onHover?.(`mode ${pick.mode} at ${this.freqOf(pick.mode).toFixed(1)} Hz, ${this.db(E).toFixed(0)} dB at t = ${pick.time.toFixed(2)} s`);
    });
    canvas.addEventListener("pointerleave", () => this.onHover?.(""));
  }

  private freqOf(m: number) { return this.tracker ? this.tracker.omega[m] / (2 * Math.PI) : 0; }
  private db(E: number) { return 10 * Math.log10(Math.max(E, 1e-300) / Math.max(this.tracker?.maxEnergy ?? 1, 1e-300)); }

  private pick(e: PointerEvent): { mode: number; time: number } | null {
    const t = this.tracker;
    if (!t || !t.length) return null;
    const r = this.canvas.getBoundingClientRect();
    const W = r.width - LEFT, H = r.height - BOTTOM;
    const px = e.clientX - r.left - LEFT, py = e.clientY - r.top;
    if (px < 0 || py > H) return null;
    const time = Math.max(0, Math.min(this.duration, (px / W) * this.duration));
    let best = -1, bestD = 8;   // within 8 px of a trace
    for (let m = 0; m < t.nModes; m++) {
      const d = Math.abs(this.axis.y(this.freqOf(m), H) - py);
      if (d < bestD) { bestD = d; best = m; }
    }
    return best >= 0 ? { mode: best, time } : null;
  }

  invalidate() { this.dirty = true; }

  /** Rebuild the cached image: max-aggregate the snapshots into pixel columns. */
  private renderImage(W: number, H: number) {
    const img = this.image;
    if (img.width !== W || img.height !== H) { img.width = W; img.height = H; }
    const g = img.getContext("2d")!;
    g.fillStyle = "#0c0d12"; g.fillRect(0, 0, W, H);
    const t = this.tracker;
    if (!t || !t.length) return;
    const data = g.getImageData(0, 0, W, H), px = data.data;
    const cols = new Float32Array(W * H).fill(-1e30);
    const maxE = Math.max(t.maxEnergy, 1e-300);
    for (let s = 0; s < t.length; s++) {
      const x = Math.min(W - 1, Math.floor((t.times[s] / this.duration) * W));
      const E = t.energies[s];
      for (let m = 0; m < t.nModes; m++) {
        const f = t.omega[m] / (2 * Math.PI);
        if (f > this.axis.fmax) continue;
        const y = Math.round(this.axis.y(f, H));
        if (y < 0 || y >= H) continue;
        const dB = 10 * Math.log10(Math.max(E[m], 1e-300) / maxE);
        const i = y * W + x;
        if (dB > cols[i]) cols[i] = dB;
      }
    }
    for (let i = 0; i < W * H; i++) {
      if (cols[i] < -1e29) continue;
      const [r, gg, b] = rampColor(1 + cols[i] / this.dbRange);
      px[4 * i] = r; px[4 * i + 1] = gg; px[4 * i + 2] = b; px[4 * i + 3] = 255;
    }
    g.putImageData(data, 0, 0);
    this.dirty = false;
  }

  draw() {
    const cv = this.canvas, dpr = window.devicePixelRatio || 1;
    const CW = Math.max(1, Math.round(cv.clientWidth * dpr)), CH = Math.max(1, Math.round(cv.clientHeight * dpr));
    if (cv.width !== CW || cv.height !== CH) { cv.width = CW; cv.height = CH; this.dirty = true; }
    const ctx = cv.getContext("2d")!;
    ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
    const W = cv.clientWidth - LEFT, H = cv.clientHeight - BOTTOM;
    if (W < 10 || H < 10) return;
    if (this.dirty || this.image.width !== W || this.image.height !== H) this.renderImage(W, H);
    ctx.fillStyle = "#14161a"; ctx.fillRect(0, 0, cv.clientWidth, cv.clientHeight);
    ctx.drawImage(this.image, LEFT, 0);
    // axes
    ctx.fillStyle = "#9aa0a8"; ctx.font = "10px system-ui"; ctx.textAlign = "right"; ctx.textBaseline = "middle";
    for (const f of this.axis.ticks()) {
      const y = this.axis.y(f, H);
      ctx.fillText(f >= 1000 ? `${f / 1000}k` : String(f), LEFT - 4, y);
      ctx.strokeStyle = "rgba(154,160,168,0.12)"; ctx.beginPath(); ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + W, y); ctx.stroke();
    }
    ctx.textAlign = "center"; ctx.textBaseline = "top";
    const tStep = this.duration > 8 ? 2 : this.duration > 3 ? 1 : this.duration > 1.2 ? 0.5 : 0.2;
    for (let tt = 0; tt <= this.duration + 1e-9; tt += tStep) {
      const x = LEFT + (tt / this.duration) * W;
      ctx.fillText(`${tt.toFixed(tt < 1 && tStep < 1 ? 1 : 0)}s`, x, H + 3);
    }
    // selected mode trace and cursor
    if (this.tracker && this.selectedMode >= 0 && this.selectedMode < this.tracker.nModes) {
      const y = this.axis.y(this.freqOf(this.selectedMode), H);
      ctx.strokeStyle = "rgba(224,162,58,0.8)"; ctx.setLineDash([3, 4]);
      ctx.beginPath(); ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + W, y); ctx.stroke(); ctx.setLineDash([]);
    }
    if (this.cursorTime !== null) {
      const x = LEFT + (this.cursorTime / this.duration) * W;
      ctx.strokeStyle = "rgba(230,230,230,0.7)";
      ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, H); ctx.stroke();
    }
  }
}

export class EnergyBars {
  private readonly canvas: HTMLCanvasElement;
  readonly axis: FrequencyAxis;    // shared with the spectrogram
  dbRange = 80;
  tracker: ModalTracker | null = null;
  snapshot = -1;
  selectedMode = -1;
  onSelect: ((mode: number) => void) | null = null;
  onHover: ((text: string) => void) | null = null;

  constructor(canvas: HTMLCanvasElement, axis: FrequencyAxis) {
    this.canvas = canvas; this.axis = axis;
    canvas.addEventListener("pointerdown", (e) => { const m = this.pick(e); if (m >= 0) this.onSelect?.(m); });
    canvas.addEventListener("pointermove", (e) => {
      const m = this.pick(e);
      if (m < 0 || !this.tracker || this.snapshot < 0) { this.onHover?.(""); return; }
      const E = this.tracker.energies[this.snapshot][m];
      this.onHover?.(`mode ${m} at ${(this.tracker.omega[m] / (2 * Math.PI)).toFixed(1)} Hz, ${this.db(E).toFixed(0)} dB`);
    });
    canvas.addEventListener("pointerleave", () => this.onHover?.(""));
  }

  private db(E: number) { return 10 * Math.log10(Math.max(E, 1e-300) / Math.max(this.tracker?.maxEnergy ?? 1, 1e-300)); }

  /** Frequency is on x here (the transpose of the spectrogram's y). */
  private x(f: number, W: number): number { return W - this.axis.y(f, W); }

  private pick(e: PointerEvent): number {
    const t = this.tracker;
    if (!t) return -1;
    const r = this.canvas.getBoundingClientRect();
    const W = r.width - LEFT, px = e.clientX - r.left - LEFT;
    let best = -1, bestD = 5;
    for (let m = 0; m < t.nModes; m++) {
      const d = Math.abs(this.x(t.omega[m] / (2 * Math.PI), W) - px);
      if (d < bestD) { bestD = d; best = m; }
    }
    return best;
  }

  draw() {
    const cv = this.canvas, dpr = window.devicePixelRatio || 1;
    const CW = Math.max(1, Math.round(cv.clientWidth * dpr)), CH = Math.max(1, Math.round(cv.clientHeight * dpr));
    if (cv.width !== CW || cv.height !== CH) { cv.width = CW; cv.height = CH; }
    const ctx = cv.getContext("2d")!;
    ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
    const W = cv.clientWidth - LEFT, H = cv.clientHeight - BOTTOM;
    ctx.fillStyle = "#0c0d12"; ctx.fillRect(0, 0, cv.clientWidth, cv.clientHeight);
    if (W < 10 || H < 10) return;
    ctx.fillStyle = "#9aa0a8"; ctx.font = "10px system-ui"; ctx.textAlign = "center"; ctx.textBaseline = "top";
    for (const f of this.axis.ticks()) ctx.fillText(f >= 1000 ? `${f / 1000}k` : String(f), LEFT + this.x(f, W), H + 3);
    ctx.textAlign = "right"; ctx.textBaseline = "middle";
    for (let dB = 0; dB >= -this.dbRange; dB -= 20) {
      const y = (-dB / this.dbRange) * H;
      ctx.fillText(String(dB), LEFT - 4, Math.min(H - 6, Math.max(6, y)));
      ctx.strokeStyle = "rgba(154,160,168,0.12)"; ctx.beginPath(); ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + W, y); ctx.stroke();
    }
    const t = this.tracker;
    if (!t || this.snapshot < 0 || this.snapshot >= t.length) return;
    const E = t.energies[this.snapshot];
    for (let m = 0; m < t.nModes; m++) {
      const f = t.omega[m] / (2 * Math.PI);
      if (f > this.axis.fmax) continue;
      const dB = this.db(E[m]);
      if (dB < -this.dbRange) continue;
      const x = LEFT + this.x(f, W), h = (1 + dB / this.dbRange) * H;
      const driven = m === t.drivenMode, sel = m === this.selectedMode;
      ctx.strokeStyle = sel ? "#e0a23a" : driven ? "#6fbf73" : "#7f9fd0";
      ctx.lineWidth = sel || driven ? 2 : 1;
      ctx.beginPath(); ctx.moveTo(x, H); ctx.lineTo(x, H - h); ctx.stroke();
    }
    ctx.lineWidth = 1;
  }
}

/** A true spectrogram (of the listener signal or of one modal coordinate) on the same frequency
 * axis as the modal view, with the linear mode frequencies overlaid as ticks at the left edge.
 * This is the complement of the modal-energy view: the modal view says which shapes hold the
 * energy, pinned at their linear frequencies; this view says at which frequencies the signal
 * actually oscillates, so combination tones between the modes, amplitude-shifted peaks and the
 * broadband turbulent floor are visible against the modal grid. */
export class StftView {
  private readonly canvas: HTMLCanvasElement;
  private readonly image = document.createElement("canvas");
  readonly axis: FrequencyAxis;
  dbRange = 80;
  stft: Stft | null = null;
  duration = 1;
  cursorTime: number | null = null;
  modeFreqs: Float64Array | null = null;
  selectedMode = -1;
  drivenMode = -1;
  onSelect: ((time: number) => void) | null = null;
  onHover: ((text: string) => void) | null = null;
  private dirty = true;

  constructor(canvas: HTMLCanvasElement, axis: FrequencyAxis) {
    this.canvas = canvas; this.axis = axis;
    canvas.addEventListener("pointerdown", (e) => {
      const p = this.pos(e);
      if (p) this.onSelect?.(p.time);
    });
    canvas.addEventListener("pointermove", (e) => {
      const p = this.pos(e);
      if (!p) { this.onHover?.(""); return; }
      let text = `t = ${p.time.toFixed(2)} s, f = ${p.f.toFixed(0)} Hz`;
      if (this.modeFreqs && this.modeFreqs.length) {
        let best = 0;
        for (let m = 1; m < this.modeFreqs.length; m++) if (Math.abs(this.modeFreqs[m] - p.f) < Math.abs(this.modeFreqs[best] - p.f)) best = m;
        const df = p.f - this.modeFreqs[best];
        text += `; nearest mode ${best} at ${this.modeFreqs[best].toFixed(0)} Hz (${df >= 0 ? "+" : ""}${df.toFixed(0)} Hz)`;
      }
      this.onHover?.(text);
    });
    canvas.addEventListener("pointerleave", () => this.onHover?.(""));
  }

  private pos(e: PointerEvent): { time: number; f: number } | null {
    const r = this.canvas.getBoundingClientRect();
    const W = r.width - LEFT, H = r.height - BOTTOM;
    const px = e.clientX - r.left - LEFT, py = e.clientY - r.top;
    if (px < 0 || py > H || W < 10) return null;
    return { time: Math.max(0, Math.min(this.duration, (px / W) * this.duration)), f: this.axis.f(py, H) };
  }

  invalidate() { this.dirty = true; }

  private renderImage(W: number, H: number) {
    const img = this.image;
    if (img.width !== W || img.height !== H) { img.width = W; img.height = H; }
    const g = img.getContext("2d")!;
    g.fillStyle = "#0c0d12"; g.fillRect(0, 0, W, H);
    const st = this.stft;
    if (!st || !st.columns.length) { this.dirty = false; return; }
    let maxP = 0;
    for (const col of st.columns) for (const p of col) maxP = Math.max(maxP, p);
    if (!(maxP > 0)) { this.dirty = false; return; }
    const data = g.getImageData(0, 0, W, H), px = data.data;
    const cells = new Float32Array(W * H).fill(-1e30);
    for (let c = 0; c < st.columns.length; c++) {
      const x = Math.min(W - 1, Math.floor((st.columnTime(c) / this.duration) * W));
      if (x < 0) continue;
      const col = st.columns[c];
      for (let b = 1; b < st.nBins; b++) {
        const f = st.binFrequency(b);
        if (f > this.axis.fmax) break;
        const y = Math.round(this.axis.y(f, H));
        if (y < 0 || y >= H) continue;
        const dB = 10 * Math.log10(col[b] / maxP);
        const i = y * W + x;
        if (dB > cells[i]) cells[i] = dB;
      }
    }
    for (let i = 0; i < W * H; i++) {
      if (cells[i] < -1e29) continue;
      const [r, gg, b] = rampColor(1 + cells[i] / this.dbRange);
      px[4 * i] = r; px[4 * i + 1] = gg; px[4 * i + 2] = b; px[4 * i + 3] = 255;
    }
    g.putImageData(data, 0, 0);
    this.dirty = false;
  }

  draw() {
    const cv = this.canvas, dpr = window.devicePixelRatio || 1;
    const CW = Math.max(1, Math.round(cv.clientWidth * dpr)), CH = Math.max(1, Math.round(cv.clientHeight * dpr));
    if (cv.width !== CW || cv.height !== CH) { cv.width = CW; cv.height = CH; this.dirty = true; }
    const ctx = cv.getContext("2d")!;
    ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
    const W = cv.clientWidth - LEFT, H = cv.clientHeight - BOTTOM;
    if (W < 10 || H < 10) return;
    if (this.dirty || this.image.width !== W || this.image.height !== H) this.renderImage(W, H);
    ctx.fillStyle = "#14161a"; ctx.fillRect(0, 0, cv.clientWidth, cv.clientHeight);
    ctx.drawImage(this.image, LEFT, 0);
    // axes
    ctx.fillStyle = "#9aa0a8"; ctx.font = "10px system-ui"; ctx.textAlign = "right"; ctx.textBaseline = "middle";
    for (const f of this.axis.ticks()) {
      const y = this.axis.y(f, H);
      ctx.fillText(f >= 1000 ? `${f / 1000}k` : String(f), LEFT - 4, y);
      ctx.strokeStyle = "rgba(154,160,168,0.12)"; ctx.beginPath(); ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + W, y); ctx.stroke();
    }
    ctx.textAlign = "center"; ctx.textBaseline = "top";
    const tStep = this.duration > 8 ? 2 : this.duration > 3 ? 1 : this.duration > 1.2 ? 0.5 : 0.2;
    for (let tt = 0; tt <= this.duration + 1e-9; tt += tStep) {
      const x = LEFT + (tt / this.duration) * W;
      ctx.fillText(`${tt.toFixed(tt < 1 && tStep < 1 ? 1 : 0)}s`, x, H + 3);
    }
    // the linear modal grid: a tick at the left edge for every mode frequency
    if (this.modeFreqs) {
      ctx.strokeStyle = "rgba(127,159,208,0.55)";
      ctx.beginPath();
      for (let m = 0; m < this.modeFreqs.length; m++) {
        const f = this.modeFreqs[m];
        if (f > this.axis.fmax) continue;
        const y = this.axis.y(f, H);
        ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + 6, y);
      }
      ctx.stroke();
      const line = (m: number, color: string) => {
        if (m < 0 || m >= this.modeFreqs!.length) return;
        const y = this.axis.y(this.modeFreqs![m], H);
        ctx.strokeStyle = color; ctx.setLineDash([3, 4]);
        ctx.beginPath(); ctx.moveTo(LEFT, y); ctx.lineTo(LEFT + W, y); ctx.stroke(); ctx.setLineDash([]);
      };
      line(this.drivenMode, "rgba(111,191,115,0.7)");
      if (this.selectedMode !== this.drivenMode) line(this.selectedMode, "rgba(224,162,58,0.8)");
    }
    if (this.cursorTime !== null) {
      const x = LEFT + (this.cursorTime / this.duration) * W;
      ctx.strokeStyle = "rgba(230,230,230,0.7)";
      ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, H); ctx.stroke();
    }
  }
}