// In-place radix-2 complex FFT (length must be a power of two).
export function fft(re: Float64Array, im: Float64Array): void {
  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;
    const wr = Math.cos(ang);
    const wi = Math.sin(ang);
    for (let i = 0; i < n; i += len) {
      let cr = 1;
      let ci = 0;
      for (let k = 0; k < len / 2; k++) {
        const a = i + k;
        const b = a + len / 2;
        const tr = re[b] * cr - im[b] * ci;
        const 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 interface Stft {
  mags: Float32Array[]; // dB per frame, nfft/2 + 1 bins
  hop: number;
  nfft: number;
  fs: number;
}

export function stft(x: Float32Array, fs: number, nfft = 2048, hop = 256): Stft {
  const win = new Float64Array(nfft);
  for (let i = 0; i < nfft; i++) win[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / nfft);
  const mags: Float32Array[] = [];
  const re = new Float64Array(nfft);
  const im = new Float64Array(nfft);
  for (let start = -nfft / 2; start < x.length - nfft / 2; start += hop) {
    for (let i = 0; i < nfft; i++) {
      const j = start + i;
      re[i] = j >= 0 && j < x.length ? x[j] * win[i] : 0;
      im[i] = 0;
    }
    fft(re, im);
    const m = new Float32Array(nfft / 2 + 1);
    for (let k = 0; k <= nfft / 2; k++) m[k] = 10 * Math.log10(re[k] * re[k] + im[k] * im[k] + 1e-30);
    mags.push(m);
  }
  return { mags, hop, nfft, fs };
}

// Power spectrum of the whole signal (Hann window, zero-padded), in dB.
export function spectrum(x: Float32Array, fs: number): { freq: Float64Array; db: Float64Array } {
  let n = 1;
  while (n < x.length) n <<= 1;
  const re = new Float64Array(n);
  const im = new Float64Array(n);
  for (let i = 0; i < x.length; i++) re[i] = x[i] * (0.5 - 0.5 * Math.cos((2 * Math.PI * i) / x.length));
  fft(re, im);
  const freq = new Float64Array(n / 2 + 1);
  const db = new Float64Array(n / 2 + 1);
  for (let k = 0; k <= n / 2; k++) {
    freq[k] = (k * fs) / n;
    db[k] = 10 * Math.log10(re[k] * re[k] + im[k] * im[k] + 1e-30);
  }
  return { freq, db };
}