// 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;
  }
}