package audiofeatures

import "math"

// fft computes the in-place iterative radix-2 Cooley-Tukey FFT of a,
// whose length must be a power of two. It's the only DSP primitive this
// package needs (spectral centroid/rolloff and the chroma vector all
// derive from the resulting magnitude spectrum), so a small
// hand-written implementation avoids pulling in an external numerics
// dependency for one function.
func fft(a []complex128) {
	n := len(a)
	if n <= 1 {
		return
	}

	// Bit-reversal permutation.
	for i, j := 1, 0; i < n; i++ {
		bit := n >> 1
		for ; j&bit != 0; bit >>= 1 {
			j ^= bit
		}
		j ^= bit
		if i < j {
			a[i], a[j] = a[j], a[i]
		}
	}

	for length := 2; length <= n; length <<= 1 {
		ang := -2 * math.Pi / float64(length)
		wlen := complex(math.Cos(ang), math.Sin(ang))
		for i := 0; i < n; i += length {
			w := complex(1, 0)
			half := length / 2
			for k := 0; k < half; k++ {
				u := a[i+k]
				v := a[i+k+half] * w
				a[i+k] = u + v
				a[i+k+half] = u - v
				w *= wlen
			}
		}
	}
}

// magnitudeSpectrum returns |FFT(frame)| for bins [0, n/2], given a
// real-valued frame whose length is a power of two.
func magnitudeSpectrum(frame []float64) []float64 {
	n := len(frame)
	buf := make([]complex128, n)
	for i, v := range frame {
		buf[i] = complex(v, 0)
	}
	fft(buf)
	mags := make([]float64, n/2+1)
	for i := range mags {
		mags[i] = cmplxAbs(buf[i])
	}
	return mags
}

func cmplxAbs(c complex128) float64 {
	re, im := real(c), imag(c)
	return math.Sqrt(re*re + im*im)
}