package main

import (
	"math"
	"math/cmplx"
	"testing"
)

func TestFFTBasic(t *testing.T) {
	// Test with simple data
	data := []complex128{
		complex(1, 0),
		complex(2, 0),
		complex(3, 0),
		complex(4, 0),
	}

	result := FFT(data)

	// Check that result has same length
	if len(result) != len(data) {
		t.Errorf("FFT result length %d, expected %d", len(result), len(data))
	}

	// Check that result is not all zeros
	allZero := true
	for _, val := range result {
		if cmplx.Abs(val) > 1e-10 {
			allZero = false
			break
		}
	}
	if allZero {
		t.Error("FFT result is all zeros")
	}
}

func TestFFTPowerOfTwo(t *testing.T) {
	// Test with non-power-of-2 length
	data := []complex128{
		complex(1, 0),
		complex(2, 0),
		complex(3, 0),
		complex(4, 0),
		complex(5, 0),
	}

	result := FFT(data)

	// Should be padded to next power of 2 (8)
	expectedLen := 8
	if len(result) != expectedLen {
		t.Errorf("FFT result length %d, expected %d", len(result), expectedLen)
	}
}

func TestIFFT(t *testing.T) {
	// Test that IFFT(FFT(data)) ≈ data
	data := []complex128{
		complex(1, 0),
		complex(2, 0),
		complex(3, 0),
		complex(4, 0),
	}

	fftResult := FFT(data)
	ifftResult := IFFT(fftResult)

	// Check that IFFT recovers original data (within numerical precision)
	tolerance := 1e-10
	for i, original := range data {
		recovered := ifftResult[i]
		diff := cmplx.Abs(original - recovered)
		if diff > tolerance {
			t.Errorf("IFFT recovery failed at index %d: original=%v, recovered=%v, diff=%v",
				i, original, recovered, diff)
		}
	}
}

func TestFFTComplexData(t *testing.T) {
	// Test with complex input data
	data := []complex128{
		complex(1, 1),
		complex(2, -1),
		complex(-3, 2),
		complex(4, 0),
	}

	result := FFT(data)

	// Check that result has same length
	if len(result) != len(data) {
		t.Errorf("FFT result length %d, expected %d", len(result), len(data))
	}

	// Check that result is not all zeros
	allZero := true
	for _, val := range result {
		if cmplx.Abs(val) > 1e-10 {
			allZero = false
			break
		}
	}
	if allZero {
		t.Error("FFT result is all zeros")
	}
}

func TestFFTEmpty(t *testing.T) {
	// Test with empty slice
	var data []complex128
	result := FFT(data)

	if len(result) != 0 {
		t.Errorf("FFT of empty slice should return empty slice, got length %d", len(result))
	}
}

func TestFFTSingle(t *testing.T) {
	// Test with single element
	data := []complex128{complex(5, 3)}
	result := FFT(data)

	if len(result) != 1 {
		t.Errorf("FFT of single element should return single element, got length %d", len(result))
	}

	// Single element FFT should return the same value
	if cmplx.Abs(result[0]-data[0]) > 1e-10 {
		t.Errorf("FFT of single element should return same value, got %v, expected %v",
			result[0], data[0])
	}
}

func TestFFTMathematical(t *testing.T) {
	// Test with mathematical properties of FFT
	// FFT of [1, 0, 0, 0] should be [1, 1, 1, 1]
	data := []complex128{
		complex(1, 0),
		complex(0, 0),
		complex(0, 0),
		complex(0, 0),
	}

	result := FFT(data)

	// All elements should be approximately 1
	tolerance := 1e-10
	for i, val := range result {
		if cmplx.Abs(val-complex(1, 0)) > tolerance {
			t.Errorf("FFT of impulse should be all ones, got %v at index %d", val, i)
		}
	}
}

func BenchmarkFFT(b *testing.B) {
	// Benchmark with power of 2 size
	size := 1024
	data := make([]complex128, size)
	for i := range data {
		data[i] = complex(float64(i), float64(i%10))
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		FFT(data)
	}
}

func BenchmarkFFTLarge(b *testing.B) {
	// Benchmark with larger size
	size := 4096
	data := make([]complex128, size)
	for i := range data {
		data[i] = complex(float64(i), float64(i%10))
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		FFT(data)
	}
}

func BenchmarkIFFT(b *testing.B) {
	// Benchmark IFFT
	size := 1024
	data := make([]complex128, size)
	for i := range data {
		data[i] = complex(float64(i), float64(i%10))
	}

	fftResult := FFT(data)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		IFFT(fftResult)
	}
}