/**
 * In-place iterative radix-2 complex FFT on separate real and imaginary
 * arrays. Sizes must be powers of two. The inverse transform is scaled by 1/n.
 */
export class FFT {
  readonly n: number
  private rev: Uint32Array
  private cos: Float64Array
  private sin: Float64Array

  constructor(n: number) {
    if (n < 2 || (n & (n - 1)) !== 0) throw new Error(`FFT size must be a power of two, got ${n}`)
    this.n = n
    const bits = Math.round(Math.log2(n))
    this.rev = new Uint32Array(n)
    for (let i = 0; i < n; i++) {
      let r = 0
      let x = i
      for (let b = 0; b < bits; b++) {
        r = (r << 1) | (x & 1)
        x >>= 1
      }
      this.rev[i] = r >>> 0
    }
    this.cos = new Float64Array(n / 2)
    this.sin = new Float64Array(n / 2)
    for (let i = 0; i < n / 2; i++) {
      const a = (-2 * Math.PI * i) / n
      this.cos[i] = Math.cos(a)
      this.sin[i] = Math.sin(a)
    }
  }

  transform(re: Float64Array, im: Float64Array, inverse = false): void {
    const n = this.n
    const rev = this.rev
    for (let i = 0; i < n; i++) {
      const j = rev[i]
      if (j > i) {
        const tr = re[i]
        re[i] = re[j]
        re[j] = tr
        const ti = im[i]
        im[i] = im[j]
        im[j] = ti
      }
    }
    const sgn = inverse ? -1 : 1
    for (let size = 2; size <= n; size <<= 1) {
      const half = size >> 1
      const step = n / size
      for (let start = 0; start < n; start += size) {
        for (let k = 0, tw = 0; k < half; k++, tw += step) {
          const wr = this.cos[tw]
          const wi = sgn * this.sin[tw]
          const i = start + k
          const j = i + half
          const tr = re[j] * wr - im[j] * wi
          const ti = re[j] * wi + im[j] * wr
          re[j] = re[i] - tr
          im[j] = im[i] - ti
          re[i] += tr
          im[i] += ti
        }
      }
    }
    if (inverse) {
      for (let i = 0; i < n; i++) {
        re[i] /= n
        im[i] /= n
      }
    }
  }
}

export function nextPow2(n: number): number {
  let p = 1
  while (p < n) p <<= 1
  return p
}

const cache = new Map<number, FFT>()

/** Shared FFT plans, one per size. */
export function fftPlan(n: number): FFT {
  let p = cache.get(n)
  if (!p) {
    p = new FFT(n)
    cache.set(n, p)
  }
  return p
}