// Incremental short-time Fourier transform of the listener signal: a Hann window of nfft samples
// advanced by hop, one power column per hop, plus the spectral centroid of each column.

/** In-place iterative radix-2 FFT of (re, im); n must be a power of two. */
export function fft(re: Float32Array, im: Float32Array) {
  const n = re.length;
  for (let i = 1, j = 0; i < n; i++) {
    let bit = n >> 1;
    for (; j & bit; bit >>= 1) j ^= bit;
    j ^= bit;
    if (i < j) { [re[i], re[j]] = [re[j], re[i]]; [im[i], im[j]] = [im[j], im[i]]; }
  }
  for (let len = 2; len <= n; len <<= 1) {
    const ang = (-2 * Math.PI) / len, wr = Math.cos(ang), wi = Math.sin(ang);
    for (let i = 0; i < n; i += len) {
      let cr = 1, ci = 0;
      for (let k = 0; k < len / 2; k++) {
        const a = i + k, b = a + len / 2;
        const tr = re[b] * cr - im[b] * ci, ti = re[b] * ci + im[b] * cr;
        re[b] = re[a] - tr; im[b] = im[a] - ti; re[a] += tr; im[a] += ti;
        const nr = cr * wr - ci * wi; ci = cr * wi + ci * wr; cr = nr;
      }
    }
  }
}

export class Stft {
  readonly nfft: number; readonly hop: number; readonly fs: number; readonly nBins: number;
  readonly columns: Float32Array[] = [];   // power per bin
  readonly centroid: number[] = [];        // Hz per column
  private readonly window: Float32Array;
  private buffer = new Float32Array(0);
  private consumed = 0;                    // samples already removed from the front of buffer
  private readonly re: Float32Array; private readonly im: Float32Array;

  constructor(fs: number, nfft = 1024, hop = 256) {
    this.fs = fs; this.nfft = nfft; this.hop = hop; this.nBins = nfft / 2 + 1;
    this.window = new Float32Array(nfft);
    for (let i = 0; i < nfft; i++) this.window[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / nfft);
    this.re = new Float32Array(nfft); this.im = new Float32Array(nfft);
  }

  /** Time of column c (its window centre) in seconds. */
  columnTime(c: number): number { return (c * this.hop + this.nfft / 2) / this.fs; }
  binFrequency(b: number): number { return (b * this.fs) / this.nfft; }

  /** Append samples; returns the number of new columns. */
  push(samples: Float32Array): number {
    const joined = new Float32Array(this.buffer.length + samples.length);
    joined.set(this.buffer); joined.set(samples, this.buffer.length);
    this.buffer = joined;
    let n = 0, start = 0;
    while (start + this.nfft <= this.buffer.length) {
      for (let i = 0; i < this.nfft; i++) { this.re[i] = this.buffer[start + i] * this.window[i]; this.im[i] = 0; }
      fft(this.re, this.im);
      const col = new Float32Array(this.nBins);
      let num = 0, den = 0;
      for (let b = 0; b < this.nBins; b++) {
        col[b] = this.re[b] * this.re[b] + this.im[b] * this.im[b];
        num += col[b] * this.binFrequency(b); den += col[b];
      }
      this.columns.push(col); this.centroid.push(den > 0 ? num / den : 0);
      start += this.hop; n++;
    }
    this.buffer = this.buffer.slice(start); this.consumed += start;
    return n;
  }
}

export const DEFAULT_BANDS = [60, 125, 250, 500, 1000, 2000, 4000, 8000, 16000];

export interface BandSummary { lo: number; hi: number; onsetMs: number; peakMs: number; levelDb: number }

/** Per-octave band envelopes: onset (first time at -20 dB of the band's own peak), peak time, level vs the loudest band. */
export function bandSummary(st: Stft, bands = DEFAULT_BANDS): { bands: BandSummary[]; centroidHz: { ms: number; hz: number }[] } {
  const nCol = st.columns.length;
  if (!nCol) return { bands: [], centroidHz: [] };
  const total = new Float64Array(nCol);
  for (let c = 0; c < nCol; c++) for (const p of st.columns[c]) total[c] += p;
  const totMax = Math.max(...total) || 1;
  const out: BandSummary[] = [];
  for (let b = 0; b + 1 < bands.length; b++) {
    const lo = bands[b], hi = bands[b + 1];
    const b0 = Math.ceil(lo / st.binFrequency(1)), b1 = Math.min(st.nBins, Math.ceil(hi / st.binFrequency(1)));
    if (b0 >= b1) continue;
    const env = new Float64Array(nCol);
    let mx = 0, peak = 0;
    for (let c = 0; c < nCol; c++) { let s = 0; for (let k = b0; k < b1; k++) s += st.columns[c][k]; env[c] = s; if (s > mx) { mx = s; peak = c; } }
    if (mx <= 0) continue;
    let onset = 0;
    for (let c = 0; c < nCol; c++) if (env[c] >= mx * 0.01) { onset = c; break; }
    out.push({ lo, hi, onsetMs: 1e3 * st.columnTime(onset), peakMs: 1e3 * st.columnTime(peak), levelDb: 10 * Math.log10(mx / totMax) });
  }
  const centroidHz = [5, 10, 20, 50, 100, 200, 500, 1000, 2000].filter((ms) => ms / 1e3 < st.columnTime(nCol - 1))
    .map((ms) => { let c = 0; while (c < nCol - 1 && st.columnTime(c) < ms / 1e3) c++; return { ms, hz: Math.round(st.centroid[c]) }; });
  return { bands: out, centroidHz };
}