package main

import (
	"fmt"
	"math"
	"math/cmplx"

	"github.com/klauspost/cpuid/v2"
)

// FFT performs Fast Fourier Transform on complex data
func FFT(data []complex128) []complex128 {
	if len(data) == 0 {
		return data
	}

	// Check if we can use AVX512
	if cpuid.CPU.AVX512F() && cpuid.CPU.AVX512DQ() {
		return fftAVX512(data)
	}

	// Fallback to standard Go implementation
	return fftGo(data)
}

// fftGo is the standard Go implementation of FFT
func fftGo(data []complex128) []complex128 {
	n := len(data)
	if n == 1 {
		return data
	}

	// Ensure n is a power of 2
	if n&(n-1) != 0 {
		// Pad with zeros to next power of 2
		nextPower := 1
		for nextPower < n {
			nextPower <<= 1
		}
		padded := make([]complex128, nextPower)
		copy(padded, data)
		data = padded
		n = nextPower
	}

	// Bit-reversal permutation
	rev := make([]int, n)
	for i := 0; i < n; i++ {
		rev[i] = rev[i>>1]>>1 | (i&1)<<int(math.Log2(float64(n))-1)
	}

	// Apply bit-reversal
	result := make([]complex128, n)
	for i := 0; i < n; i++ {
		result[i] = data[rev[i]]
	}

	// Cooley-Tukey FFT
	for size := 2; size <= n; size <<= 1 {
		half := size >> 1
		angle := -2 * math.Pi / float64(size)
		w := complex(math.Cos(angle), math.Sin(angle))

		for i := 0; i < n; i += size {
			wi := complex(1, 0)
			for j := 0; j < half; j++ {
				t := wi * result[i+j+half]
				result[i+j+half] = result[i+j] - t
				result[i+j] += t
				wi *= w
			}
		}
	}

	return result
}

// fftAVX512 calls the AVX512 assembly implementation
//go:noescape
func fftAVX512(data []complex128) []complex128

// Inverse FFT
func IFFT(data []complex128) []complex128 {
	n := len(data)
	if n == 0 {
		return data
	}

	// Conjugate input
	conj := make([]complex128, n)
	for i := 0; i < n; i++ {
		conj[i] = cmplx.Conj(data[i])
	}

	// Apply FFT
	fftResult := FFT(conj)

	// Conjugate output and scale
	result := make([]complex128, n)
	for i := 0; i < n; i++ {
		result[i] = cmplx.Conj(fftResult[i]) / complex(float64(n), 0)
	}

	return result
}

func main() {
	// Example usage
	fmt.Println("AVX512 Support:", cpuid.CPU.AVX512F() && cpuid.CPU.AVX512DQ())

	// Test data
	data := []complex128{
		complex(1, 0),
		complex(2, 0),
		complex(3, 0),
		complex(4, 0),
		complex(5, 0),
		complex(6, 0),
		complex(7, 0),
		complex(8, 0),
	}

	fmt.Println("Input:", data)

	// Forward FFT
	fftResult := FFT(data)
	fmt.Println("FFT Result:", fftResult)

	// Inverse FFT
	ifftResult := IFFT(fftResult)
	fmt.Println("IFFT Result:", ifftResult)
}