"""Tests for FFT operations."""

import numpy as np
import pytest
from compute_core.fft import fft, fft2, fftfreq, ifft, ifft2, irfft, rfft


def test_fft_1d() -> None:
    """Test 1D FFT."""
    x = np.array([1, 2, 3, 4, 5])
    result = fft(x)
    expected = np.fft.fft(x)
    np.testing.assert_array_almost_equal(result, expected)


def test_ifft_1d() -> None:
    """Test 1D IFFT."""
    x = np.array([1, 2, 3, 4, 5])
    fft_result = fft(x)
    result = ifft(fft_result)
    np.testing.assert_array_almost_equal(result, x)


def test_fft_ifft_roundtrip() -> None:
    """Test FFT-IFFT roundtrip."""
    x = np.random.rand(100)
    result = ifft(fft(x))
    np.testing.assert_array_almost_equal(result, x)


def test_fft2_2d() -> None:
    """Test 2D FFT."""
    x = np.random.rand(10, 10)
    result = fft2(x)
    expected = np.fft.fft2(x)
    np.testing.assert_array_almost_equal(result, expected)


def test_ifft2_2d() -> None:
    """Test 2D IFFT."""
    x = np.random.rand(10, 10)
    fft_result = fft2(x)
    result = ifft2(fft_result)
    np.testing.assert_array_almost_equal(result, x)


def test_rfft_real_input() -> None:
    """Test real FFT with real input."""
    x = np.random.rand(100)
    result = rfft(x)
    expected = np.fft.rfft(x)
    np.testing.assert_array_almost_equal(result, expected)


def test_irfft_roundtrip() -> None:
    """Test RFFT-IRFFT roundtrip."""
    x = np.random.rand(100)
    result = irfft(rfft(x))
    np.testing.assert_array_almost_equal(result, x)


def test_fftfreq() -> None:
    """Test FFT frequency calculation."""
    n = 10
    d = 0.1
    result = fftfreq(n, d)
    expected = np.fft.fftfreq(n, d)
    np.testing.assert_array_almost_equal(result, expected)


@pytest.mark.gpu()
def test_fft_gpu() -> None:
    """Test FFT on GPU."""
    try:
        import cupy as cp
        from compute_core.arrays import to_gpu

        x = np.random.rand(100)
        x_gpu = to_gpu(x)
        result = fft(x_gpu)

        assert hasattr(result, "__cuda_array_interface__")

        # Verify correctness
        expected = np.fft.fft(x)
        np.testing.assert_array_almost_equal(cp.asnumpy(result), expected)
    except ImportError:
        pytest.skip("CuPy not available")


@pytest.mark.gpu()
def test_fft2_gpu() -> None:
    """Test 2D FFT on GPU."""
    try:
        import cupy as cp
        from compute_core.arrays import to_gpu

        x = np.random.rand(50, 50)
        x_gpu = to_gpu(x)
        result = fft2(x_gpu)

        assert hasattr(result, "__cuda_array_interface__")

        # Verify correctness
        expected = np.fft.fft2(x)
        np.testing.assert_array_almost_equal(cp.asnumpy(result), expected)
    except ImportError:
        pytest.skip("CuPy not available")
