package odometry

import "math"

// fft is an in-place iterative radix-2 Cooley-Tukey FFT.
// len(a) must be a power of two. If inverse is true the result is normalized
// by 1/n (so ifft(fft(x)) == x).
func fft(a []complex128, inverse bool) {
	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]
		}
	}

	// Butterfly stages.
	for length := 2; length <= n; length <<= 1 {
		// Standard forward DFT uses exp(-i*2π/length); inverse negates.
		angle := -2 * math.Pi / float64(length)
		if inverse {
			angle = -angle
		}
		wlen := complex(math.Cos(angle), math.Sin(angle))
		half := length / 2
		for i := 0; i < n; i += length {
			var w complex128 = 1
			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
			}
		}
	}

	if inverse {
		invN := complex(1/float64(n), 0)
		for i := range a {
			a[i] *= invN
		}
	}
}

// fftRows applies a 1-D FFT to every row of m (in place).
func fftRows(m [][]complex128, inverse bool) {
	for i := range m {
		fft(m[i], inverse)
	}
}

// fftCols applies a 1-D FFT to every column of m (in place).
func fftCols(m [][]complex128, inverse bool) {
	rows := len(m)
	if rows == 0 {
		return
	}
	cols := len(m[0])
	col := make([]complex128, rows)
	for j := 0; j < cols; j++ {
		for i := 0; i < rows; i++ {
			col[i] = m[i][j]
		}
		fft(col, inverse)
		for i := 0; i < rows; i++ {
			m[i][j] = col[i]
		}
	}
}

// fft2 performs a forward 2-D FFT on the rows×cols complex matrix m.
func fft2(m [][]complex128, inverse bool) {
	fftRows(m, inverse)
	fftCols(m, inverse)
}