// Fast Fourier Transform (adapted from spectral_editor/dct.js)
// Radix-2 Cooley-Tukey algorithm

// Bit-reversal permutation (in-place)
function bitReversePermute(real, imag, N) {
  let temp_bits = N;
  let num_bits = 0;
  while (temp_bits > 1) {
    temp_bits >>= 1;
    num_bits++;
  }

  for (let i = 0; i < N; ++i) {
    let j = 0;
    let temp = i;
    for (let b = 0; b < num_bits; ++b) {
      j = (j << 1) | (temp & 1);
      temp >>= 1;
    }

    if (j > i) {
      const tmp_real = real[i];
      const tmp_imag = imag[i];
      real[i] = real[j];
      imag[i] = imag[j];
      real[j] = tmp_real;
      imag[j] = tmp_imag;
    }
  }
}

// In-place radix-2 FFT
// direction: +1 for forward, -1 for inverse
function fftRadix2(real, imag, N, direction) {
  const PI = Math.PI;

  for (let stage_size = 2; stage_size <= N; stage_size *= 2) {
    const half_stage = stage_size / 2;
    const angle = direction * 2.0 * PI / stage_size;

    let wr = 1.0;
    let wi = 0.0;
    const wr_delta = Math.cos(angle);
    const wi_delta = Math.sin(angle);

    for (let k = 0; k < half_stage; ++k) {
      for (let group_start = k; group_start < N; group_start += stage_size) {
        const i = group_start;
        const j = group_start + half_stage;

        const temp_real = real[j] * wr - imag[j] * wi;
        const temp_imag = real[j] * wi + imag[j] * wr;

        real[j] = real[i] - temp_real;
        imag[j] = imag[i] - temp_imag;
        real[i] = real[i] + temp_real;
        imag[i] = imag[i] + temp_imag;
      }

      const wr_old = wr;
      wr = wr_old * wr_delta - wi * wi_delta;
      wi = wr_old * wi_delta + wi * wr_delta;
    }
  }
}

// Forward FFT: Time domain → Frequency domain
function fftForward(real, imag, N) {
  bitReversePermute(real, imag, N);
  fftRadix2(real, imag, N, +1);
}

// Real FFT wrapper for MQ extraction
// Input: Float32Array (time-domain signal)
// Output: Float32Array (interleaved [re0, im0, re1, im1, ...])
function realFFT(signal) {
  const N = signal.length;

  // Must be power of 2
  if ((N & (N - 1)) !== 0) {
    throw new Error('FFT size must be power of 2');
  }

  const real = new Float32Array(N);
  const imag = new Float32Array(N);

  // Copy input to real part
  for (let i = 0; i < N; ++i) {
    real[i] = signal[i];
  }

  // Compute FFT
  fftForward(real, imag, N);

  // Interleave output
  const spectrum = new Float32Array(N * 2);
  for (let i = 0; i < N; ++i) {
    spectrum[i * 2] = real[i];
    spectrum[i * 2 + 1] = imag[i];
  }

  return spectrum;
}

// STFT Cache - Pre-computes and caches windowed FFT frames
class STFTCache {
  constructor(signal, sampleRate, fftSize, hopSize) {
    this.signal = signal;
    this.sampleRate = sampleRate;
    this.fftSize = fftSize;
    this.hopSize = hopSize;
    this.frames = []; // Array of {time, offset, squaredAmplitude}
    this.compute();
  }

  compute() {
    this.frames = [];
    const numFrames = Math.floor((this.signal.length - this.fftSize) / this.hopSize);

    for (let frameIdx = 0; frameIdx < numFrames; ++frameIdx) {
      const offset = frameIdx * this.hopSize;
      const time = offset / this.sampleRate;

      // Extract frame
      const frame = this.signal.slice(offset, offset + this.fftSize);

      // Apply Hann window
      const windowed = new Float32Array(this.fftSize);
      for (let i = 0; i < this.fftSize; ++i) {
        const w = 0.5 - 0.5 * Math.cos(2 * Math.PI * i / this.fftSize);
        windowed[i] = frame[i] * w;
      }

      // Compute FFT, store only squared amplitudes (re*re + im*im, no sqrt)
      const fftOut = realFFT(windowed);
      const squaredAmplitude = new Float32Array(this.fftSize / 2);
      for (let i = 0; i < this.fftSize / 2; ++i) {
        const re = fftOut[i * 2];
        const im = fftOut[i * 2 + 1];
        squaredAmplitude[i] = re * re + im * im;
      }

      this.frames.push({time, offset, squaredAmplitude});
    }
  }

  setHopSize(hopSize) {
    if (hopSize === this.hopSize) return;
    this.hopSize = hopSize;
    this.compute();
  }

  setFFTSize(fftSize) {
    if (fftSize === this.fftSize) return;
    this.fftSize = fftSize;
    this.compute();
  }

  getNumFrames() {
    return this.frames.length;
  }

  getFrameAtIndex(frameIdx) {
    if (frameIdx < 0 || frameIdx >= this.frames.length) return null;
    return this.frames[frameIdx];
  }

  getFrameAtTime(t) {
    if (this.frames.length === 0) return null;

    // Find closest frame
    const frameIdx = Math.floor(t * this.sampleRate / this.hopSize);
    return this.getFrameAtIndex(frameIdx);
  }

  getSquaredAmplitude(t) {
    const frame = this.getFrameAtTime(t);
    return frame ? frame.squaredAmplitude : null;
  }

  // Get magnitude in dB at specific time and frequency
  getMagnitudeDB(t, freq) {
    const sq = this.getSquaredAmplitude(t);
    if (!sq) return -80;

    const bin = Math.round(freq * this.fftSize / this.sampleRate);
    if (bin < 0 || bin >= this.fftSize / 2) return -80;

    return 10 * Math.log10(Math.max(sq[bin], 1e-20));
  }
}