package dsp

import (
	"math"
	"math/cmplx"
)

// FFT computes the discrete Fourier transform of x (in-place, radix-2).
// len(x) must be a power of 2.
func FFT(x []complex128) {
	n := len(x)
	if n <= 1 {
		return
	}
	// Bit-reversal permutation
	j := 0
	for i := 1; i < n; i++ {
		bit := n >> 1
		for j&bit != 0 {
			j ^= bit
			bit >>= 1
		}
		j ^= bit
		if i < j {
			x[i], x[j] = x[j], x[i]
		}
	}
	// Cooley-Tukey butterfly
	for size := 2; size <= n; size <<= 1 {
		half := size >> 1
		wn := cmplx.Exp(complex(0, -2*math.Pi/float64(size)))
		for start := 0; start < n; start += size {
			w := complex(1, 0)
			for k := 0; k < half; k++ {
				u := x[start+k]
				v := x[start+k+half] * w
				x[start+k] = u + v
				x[start+k+half] = u - v
				w *= wn
			}
		}
	}
}

// IFFT computes the inverse DFT of x (in-place).
func IFFT(x []complex128) {
	n := len(x)
	// Conjugate, FFT, conjugate, scale
	for i := range x {
		x[i] = cmplx.Conj(x[i])
	}
	FFT(x)
	scale := 1.0 / float64(n)
	for i := range x {
		x[i] = cmplx.Conj(x[i]) * complex(scale, 0)
	}
}

// HannWindow fills w with a Hann window of length n.
func HannWindow(w []float64) {
	n := len(w)
	for i := range w {
		w[i] = 0.5 * (1.0 - math.Cos(2*math.Pi*float64(i)/float64(n)))
	}
}