//! Spectral convolution layers for Fourier Neural Operators.
//!
//! Spectral convolutions operate in the Fourier domain, enabling efficient
//! learning of long-range dependencies in spatial data. These layers are the
//! core building blocks of Fourier Neural Operators (FNO).
//!
//! ## Algorithm
//!
//! For SpectralConv2d:
//! 1. Apply 2D FFT to input: x_freq = FFT2d(x)
//! 2. Truncate to n_modes (keep only low frequencies)
//! 3. Multiply by learnable weights in frequency domain
//! 4. Apply inverse FFT: y = IFFT2d(x_freq * W)
//!
//! This is dramatically more efficient than spatial convolutions for learning
//! global patterns, as multiplication in Fourier space is O(n_modes) vs O(kernel_size).

use rtx_backend::Backend;
use rtx_nn::generic::{GenericModule, GenericModule4D};
use rtx_tensor::generic::GenericTensor;
use rtx_tensor::{ComplexTensor, Device, Tensor};
use std::fmt::Debug;

/// 1D Spectral Convolution layer.
///
/// Performs convolution in the Fourier domain for 1D signals.
#[derive(Debug)]
pub struct SpectralConv1d<B: Backend<FloatElem = f32>> {
    in_channels: usize,
    out_channels: usize,
    n_modes: usize,
    device: B::Device,
    // Weights are stored as [in_channels, out_channels, n_modes, 2]
    // where the last dimension holds [real, imag] components
    weights: GenericTensor<B, 4>,
}

impl<B: Backend<FloatElem = f32>> SpectralConv1d<B> {
    /// Create a new 1D spectral convolution layer.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `n_modes` - Number of Fourier modes to keep (frequency truncation)
    /// * `device` - Device to create the layer on
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        n_modes: usize,
        device: &B::Device,
    ) -> Self {
        // Initialize weights with Xavier uniform: scale * (2 * rand - 1)
        let scale = (2.0 / (in_channels + out_channels) as f32).sqrt();
        let weights = GenericTensor::rand([in_channels, out_channels, n_modes, 2], device)
            .mul_scalar(2.0 * scale)
            .add_scalar(-scale);

        Self {
            in_channels,
            out_channels,
            n_modes,
            device: device.clone(),
            weights,
        }
    }

    /// Get the number of Fourier modes.
    pub fn n_modes(&self) -> usize {
        self.n_modes
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule<B> for SpectralConv1d<B> {
    fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
        panic!("Use forward_3d for 1D spectral convolution");
    }

    fn device(&self) -> &B::Device {
        &self.device
    }
}

/// 2D Spectral Convolution layer.
///
/// Performs convolution in the Fourier domain for 2D spatial data.
/// This is the core component of the Fourier Neural Operator (FNO).
///
/// # Algorithm (neuraloperator v2.0 compatible)
///
/// Given input x of shape [batch, in_channels, height, width]:
/// 1. Apply 2D FFT to get frequency representation
/// 2. Apply weights1 to upper portion of frequency domain (low frequencies)
/// 3. Apply weights2 to lower portion of frequency domain (high frequencies in h)
/// 4. Apply inverse 2D FFT to return to spatial domain
///
/// The learnable parameters are two sets of complex weights in frequency space,
/// matching neuraloperator v2.0's dual-weight architecture.
#[derive(Debug)]
pub struct SpectralConv2d<B: Backend<FloatElem = f32>> {
    in_channels: usize,
    out_channels: usize,
    n_modes_h: usize,
    n_modes_w: usize,
    device: B::Device,
    /// Complex weights for upper frequency portion: [in_ch, out_ch, n_modes_h, n_modes_w, 2]
    weights1: GenericTensor<B, 5>,
    /// Complex weights for lower frequency portion: [in_ch, out_ch, n_modes_h, n_modes_w, 2]
    weights2: GenericTensor<B, 5>,
}

impl<B: Backend<FloatElem = f32>> SpectralConv2d<B> {
    /// Create a new 2D spectral convolution layer.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `n_modes_h` - Number of Fourier modes along height dimension
    /// * `n_modes_w` - Number of Fourier modes along width dimension
    /// * `device` - Device to create the layer on
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        n_modes_h: usize,
        n_modes_w: usize,
        device: &B::Device,
    ) -> Self {
        // Xavier initialization for complex weights: scale * (2 * rand - 1)
        let scale = (2.0 / (in_channels + out_channels) as f32).sqrt();

        // weights1: for upper portion of frequency domain
        let weights1 =
            GenericTensor::rand([in_channels, out_channels, n_modes_h, n_modes_w, 2], device)
                .mul_scalar(2.0 * scale)
                .add_scalar(-scale);

        // weights2: for lower portion of frequency domain
        let weights2 =
            GenericTensor::rand([in_channels, out_channels, n_modes_h, n_modes_w, 2], device)
                .mul_scalar(2.0 * scale)
                .add_scalar(-scale);

        Self {
            in_channels,
            out_channels,
            n_modes_h,
            n_modes_w,
            device: device.clone(),
            weights1,
            weights2,
        }
    }

    /// Get the number of Fourier modes in height dimension.
    pub fn n_modes_h(&self) -> usize {
        self.n_modes_h
    }

    /// Get the number of Fourier modes in width dimension.
    pub fn n_modes_w(&self) -> usize {
        self.n_modes_w
    }

    /// Get a reference to weights1 tensor.
    pub fn weights1(&self) -> &GenericTensor<B, 5> {
        &self.weights1
    }

    /// Get a reference to weights2 tensor.
    pub fn weights2(&self) -> &GenericTensor<B, 5> {
        &self.weights2
    }

    /// Get a reference to the weights tensor (returns weights1 for compatibility).
    #[deprecated(note = "Use weights1() or weights2() instead")]
    pub fn weights(&self) -> &GenericTensor<B, 5> {
        &self.weights1
    }

    /// Extract weights1 as separate real and imaginary parts.
    ///
    /// Returns `(real, imag)` where each is a `Vec<f32>` of shape
    /// `[in_channels * out_channels * n_modes_h * n_modes_w]`.
    pub fn weights1_real_imag(&self) -> (Vec<f32>, Vec<f32>) {
        let data = self.weights1.to_vec();
        let n_complex = self.in_channels * self.out_channels * self.n_modes_h * self.n_modes_w;
        let mut real = Vec::with_capacity(n_complex);
        let mut imag = Vec::with_capacity(n_complex);

        for i in 0..n_complex {
            real.push(data[i * 2]);
            imag.push(data[i * 2 + 1]);
        }

        (real, imag)
    }

    /// Extract weights2 as separate real and imaginary parts.
    ///
    /// Returns `(real, imag)` where each is a `Vec<f32>` of shape
    /// `[in_channels * out_channels * n_modes_h * n_modes_w]`.
    pub fn weights2_real_imag(&self) -> (Vec<f32>, Vec<f32>) {
        let data = self.weights2.to_vec();
        let n_complex = self.in_channels * self.out_channels * self.n_modes_h * self.n_modes_w;
        let mut real = Vec::with_capacity(n_complex);
        let mut imag = Vec::with_capacity(n_complex);

        for i in 0..n_complex {
            real.push(data[i * 2]);
            imag.push(data[i * 2 + 1]);
        }

        (real, imag)
    }

    /// Get the input channels.
    pub fn in_channels(&self) -> usize {
        self.in_channels
    }

    /// Get the output channels.
    pub fn out_channels(&self) -> usize {
        self.out_channels
    }

    /// Create a spectral convolution layer with pre-defined weights.
    ///
    /// # Arguments
    /// * `weights1_real` - Real part of weights1 [in_ch, out_ch, n_modes_h, n_modes_w]
    /// * `weights1_imag` - Imaginary part of weights1
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `n_modes_h` - Number of Fourier modes along height
    /// * `n_modes_w` - Number of Fourier modes along width
    /// * `device` - Device to create the layer on
    ///
    /// Note: This creates weights2 as a copy of weights1 for backwards compatibility.
    /// Use `from_weights_dual()` for full neuraloperator v2.0 compatibility.
    pub fn from_weights(
        weights1_real: &[f32],
        weights1_imag: &[f32],
        in_channels: usize,
        out_channels: usize,
        n_modes_h: usize,
        n_modes_w: usize,
        device: &B::Device,
    ) -> Self {
        // For backwards compatibility, use the same weights for both
        Self::from_weights_dual(
            weights1_real,
            weights1_imag,
            weights1_real,
            weights1_imag,
            in_channels,
            out_channels,
            n_modes_h,
            n_modes_w,
            device,
        )
    }

    /// Create a spectral convolution layer with dual pre-defined weights.
    ///
    /// This matches neuraloperator v2.0 architecture with two weight matrices.
    ///
    /// # Arguments
    /// * `weights1_real` - Real part of weights1 [in_ch, out_ch, n_modes_h, n_modes_w]
    /// * `weights1_imag` - Imaginary part of weights1
    /// * `weights2_real` - Real part of weights2 [in_ch, out_ch, n_modes_h, n_modes_w]
    /// * `weights2_imag` - Imaginary part of weights2
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `n_modes_h` - Number of Fourier modes along height
    /// * `n_modes_w` - Number of Fourier modes along width
    /// * `device` - Device to create the layer on
    #[allow(clippy::too_many_arguments)]
    pub fn from_weights_dual(
        weights1_real: &[f32],
        weights1_imag: &[f32],
        weights2_real: &[f32],
        weights2_imag: &[f32],
        in_channels: usize,
        out_channels: usize,
        n_modes_h: usize,
        n_modes_w: usize,
        device: &B::Device,
    ) -> Self {
        let total_complex = in_channels * out_channels * n_modes_h * n_modes_w;

        // Interleave weights1
        let mut interleaved1 = Vec::with_capacity(total_complex * 2);
        for i in 0..total_complex {
            interleaved1.push(weights1_real[i]);
            interleaved1.push(weights1_imag[i]);
        }
        let weights1 = GenericTensor::from_slice(
            &interleaved1,
            [in_channels, out_channels, n_modes_h, n_modes_w, 2],
            device,
        );

        // Interleave weights2
        let mut interleaved2 = Vec::with_capacity(total_complex * 2);
        for i in 0..total_complex {
            interleaved2.push(weights2_real[i]);
            interleaved2.push(weights2_imag[i]);
        }
        let weights2 = GenericTensor::from_slice(
            &interleaved2,
            [in_channels, out_channels, n_modes_h, n_modes_w, 2],
            device,
        );

        Self {
            in_channels,
            out_channels,
            n_modes_h,
            n_modes_w,
            device: device.clone(),
            weights1,
            weights2,
        }
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule<B> for SpectralConv2d<B> {
    fn forward(&self, _input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
        panic!("Use forward_4d for 2D spectral convolution");
    }

    fn device(&self) -> &B::Device {
        &self.device
    }
}

impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for SpectralConv2d<B> {
    fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        let shape = input.shape();
        let batch = shape[0];
        let in_ch = shape[1];
        let height = shape[2];
        let width = shape[3];

        // Validate input channels match
        assert_eq!(
            in_ch, self.in_channels,
            "Input channels {} do not match layer in_channels {}",
            in_ch, self.in_channels
        );

        // Convert GenericTensor to Vec<f32> for processing
        let input_data = input.to_vec();

        // Allocate output buffer [batch, out_channels, height, width]
        let output_size = batch * self.out_channels * height * width;
        let mut output_data = vec![0.0f32; output_size];

        // Use CPU device for ComplexTensor operations
        let cpu_device = Device::cpu();

        // Get weight data for both weight matrices
        let weight1_data = self.weights1.to_vec();
        let weight2_data = self.weights2.to_vec();

        // Process each batch
        for b in 0..batch {
            // Process each output channel
            for out_c in 0..self.out_channels {
                // Accumulate contributions from all input channels
                let mut channel_sum_real = Tensor::zeros(vec![height, width], &cpu_device)
                    .expect("Failed to create zeros");
                let mut channel_sum_imag = Tensor::zeros(vec![height, width], &cpu_device)
                    .expect("Failed to create zeros");

                for in_c in 0..self.in_channels {
                    // Extract 2D slice [height, width] for this batch and input channel
                    let slice_start = (b * in_ch + in_c) * height * width;
                    let slice_end = slice_start + (height * width);
                    let slice_data = input_data[slice_start..slice_end].to_vec();

                    // Create real tensor and convert to ComplexTensor
                    let real_tensor =
                        Tensor::from_data(slice_data, vec![height, width], &cpu_device)
                            .expect("Failed to create tensor from data");

                    let complex_input = ComplexTensor::<f32>::from_real(real_tensor)
                        .expect("Failed to create complex tensor");

                    // Apply 2D FFT
                    let fft_result = complex_input.fft2d().expect("Failed to perform FFT2D");

                    // Truncate to n_modes
                    let modes_h = self.n_modes_h.min(height);
                    let modes_w = self.n_modes_w.min(width);

                    // Extract FFT data
                    let fft_real_data = fft_result
                        .real()
                        .to_cpu()
                        .expect("Failed to get FFT real data");
                    let fft_imag_data = fft_result
                        .imag()
                        .to_cpu()
                        .expect("Failed to get FFT imag data");

                    // Weight offset for [in_c, out_c, :, :, :]
                    let weight_slice_start =
                        (in_c * self.out_channels + out_c) * self.n_modes_h * self.n_modes_w * 2;

                    // Perform complex multiplication in frequency domain
                    let mut freq_result_real = vec![0.0f32; height * width];
                    let mut freq_result_imag = vec![0.0f32; height * width];

                    // Apply weights1 to upper portion (low frequencies in h)
                    // Positions [0..modes_h, 0..modes_w]
                    for h in 0..modes_h {
                        for w in 0..modes_w {
                            let fft_idx = h * width + w;

                            // Get complex FFT value at this position
                            let fft_r = fft_real_data[fft_idx];
                            let fft_i = fft_imag_data[fft_idx];

                            // Get complex weight from weights1
                            let weight_idx = weight_slice_start + (h * self.n_modes_w + w) * 2;
                            let weight_r = weight1_data[weight_idx];
                            let weight_i = weight1_data[weight_idx + 1];

                            // Complex multiplication: (a+bi) * (c+di) = (ac-bd) + (ad+bc)i
                            let result_r = fft_r * weight_r - fft_i * weight_i;
                            let result_i = fft_r * weight_i + fft_i * weight_r;

                            freq_result_real[fft_idx] = result_r;
                            freq_result_imag[fft_idx] = result_i;
                        }
                    }

                    // Apply weights2 to lower portion (high frequencies in h)
                    // Positions [height-modes_h..height, 0..modes_w]
                    // This corresponds to negative frequencies in the FFT output
                    if height > modes_h {
                        for h in 0..modes_h {
                            let fft_h = height - modes_h + h;
                            for w in 0..modes_w {
                                let fft_idx = fft_h * width + w;

                                // Get complex FFT value at this position
                                let fft_r = fft_real_data[fft_idx];
                                let fft_i = fft_imag_data[fft_idx];

                                // Get complex weight from weights2
                                let weight_idx = weight_slice_start + (h * self.n_modes_w + w) * 2;
                                let weight_r = weight2_data[weight_idx];
                                let weight_i = weight2_data[weight_idx + 1];

                                // Complex multiplication
                                let result_r = fft_r * weight_r - fft_i * weight_i;
                                let result_i = fft_r * weight_i + fft_i * weight_r;

                                freq_result_real[fft_idx] = result_r;
                                freq_result_imag[fft_idx] = result_i;
                            }
                        }
                    }

                    // Create ComplexTensor from frequency domain result
                    let freq_real_tensor =
                        Tensor::from_data(freq_result_real, vec![height, width], &cpu_device)
                            .expect("Failed to create frequency real tensor");

                    let freq_imag_tensor =
                        Tensor::from_data(freq_result_imag, vec![height, width], &cpu_device)
                            .expect("Failed to create frequency imag tensor");

                    let freq_complex =
                        ComplexTensor::<f32>::from_real_imag(freq_real_tensor, freq_imag_tensor)
                            .expect("Failed to create frequency complex tensor");

                    // Apply inverse 2D FFT
                    let ifft_result = freq_complex.ifft2d().expect("Failed to perform IFFT2D");

                    // Accumulate to channel sum
                    let ifft_real = ifft_result.real();
                    let ifft_imag = ifft_result.imag();

                    channel_sum_real = channel_sum_real
                        .add(ifft_real)
                        .expect("Failed to add real part");
                    channel_sum_imag = channel_sum_imag
                        .add(ifft_imag)
                        .expect("Failed to add imag part");
                }

                // Extract real part and copy to output
                let final_real_data = channel_sum_real
                    .to_cpu()
                    .expect("Failed to get final real data");

                let output_offset = (b * self.out_channels + out_c) * height * width;
                for i in 0..final_real_data.len() {
                    output_data[output_offset + i] = final_real_data[i];
                }
            }
        }

        // Convert back to GenericTensor
        GenericTensor::from_slice(
            &output_data,
            [batch, self.out_channels, height, width],
            &self.device,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rtx_backend_cpu::{CpuBackend, CpuDevice};

    // ==================== SpectralConv1d Tests ====================

    #[test]
    fn test_spectral_conv1d_creation() {
        let device = CpuDevice::new();
        let conv = SpectralConv1d::<CpuBackend>::new(4, 8, 12, &device);

        assert_eq!(conv.in_channels, 4);
        assert_eq!(conv.out_channels, 8);
        assert_eq!(conv.n_modes(), 12);
    }

    #[test]
    fn test_spectral_conv1d_weight_shape() {
        let device = CpuDevice::new();
        let conv = SpectralConv1d::<CpuBackend>::new(3, 6, 8, &device);

        // Weights should be [in_ch, out_ch, n_modes, 2]
        assert_eq!(conv.weights.shape(), [3, 6, 8, 2]);
    }

    #[test]
    fn test_spectral_conv1d_weight_initialization() {
        let device = CpuDevice::new();
        let conv = SpectralConv1d::<CpuBackend>::new(4, 8, 10, &device);

        // Check Xavier initialization bounds
        let scale = (2.0_f32 / (4.0_f32 + 8.0_f32)).sqrt();
        let data = conv.weights.to_vec();

        // All values should be in range [-scale, scale]
        for &val in &data {
            assert!(
                val >= -scale && val <= scale,
                "Value {} outside Xavier range [{}, {}]",
                val,
                -scale,
                scale
            );
        }
    }

    // ==================== SpectralConv2d Tests ====================

    #[test]
    fn test_spectral_conv2d_creation() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(3, 16, 12, 12, &device);

        assert_eq!(conv.in_channels, 3);
        assert_eq!(conv.out_channels, 16);
        assert_eq!(conv.n_modes_h(), 12);
        assert_eq!(conv.n_modes_w(), 12);
    }

    #[test]
    fn test_spectral_conv2d_weight_shape() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(4, 8, 6, 10, &device);

        // Both weight matrices should be [in_ch, out_ch, n_modes_h, n_modes_w, 2]
        assert_eq!(conv.weights1().shape(), [4, 8, 6, 10, 2]);
        assert_eq!(conv.weights2().shape(), [4, 8, 6, 10, 2]);
    }

    #[test]
    fn test_spectral_conv2d_weight_initialization() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(3, 12, 8, 8, &device);

        // Check Xavier initialization bounds for both weight matrices
        let scale = (2.0_f32 / (3.0_f32 + 12.0_f32)).sqrt();

        for weights in [conv.weights1(), conv.weights2()] {
            let data = weights.to_vec();
            // All values should be in range [-scale, scale]
            for &val in &data {
                assert!(
                    val >= -scale && val <= scale,
                    "Value {} outside Xavier range [{}, {}]",
                    val,
                    -scale,
                    scale
                );
            }
        }
    }

    #[test]
    fn test_spectral_conv2d_forward_shape() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(4, 8, 12, 12, &device);

        let input = GenericTensor::randn([2, 4, 64, 64], &device);
        let output = conv.forward_4d(&input);

        // Output should have shape [batch, out_channels, height, width]
        assert_eq!(output.shape(), [2, 8, 64, 64]);
    }

    #[test]
    fn test_spectral_conv2d_preserves_spatial_dims() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(1, 32, 8, 8, &device);

        let input = GenericTensor::randn([4, 1, 32, 48], &device);
        let output = conv.forward_4d(&input);

        // Spatial dimensions should be preserved
        assert_eq!(output.shape()[2], 32); // height
        assert_eq!(output.shape()[3], 48); // width
    }

    #[test]
    fn test_spectral_conv2d_different_modes() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 6, 10, &device);

        let input = GenericTensor::randn([1, 2, 64, 64], &device);
        let output = conv.forward_4d(&input);

        // Output channels should match conv.out_channels
        assert_eq!(output.shape(), [1, 4, 64, 64]);
    }

    #[test]
    fn test_spectral_conv2d_batch_independence() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(3, 6, 8, 8, &device);

        // Single sample
        let input1 = GenericTensor::randn([1, 3, 32, 32], &device);
        let output1 = conv.forward_4d(&input1);

        // Batch of 4
        let input4 = GenericTensor::randn([4, 3, 32, 32], &device);
        let output4 = conv.forward_4d(&input4);

        // Batch dimension should be the only difference
        assert_eq!(output1.shape()[1..], output4.shape()[1..]);
    }

    #[test]
    fn test_spectral_conv2d_mode_truncation() {
        let device = CpuDevice::new();

        // Small mode count (strong low-pass filtering)
        let conv_small = SpectralConv2d::<CpuBackend>::new(1, 1, 4, 4, &device);

        // Large mode count (weak filtering)
        let conv_large = SpectralConv2d::<CpuBackend>::new(1, 1, 30, 30, &device);

        let input = GenericTensor::randn([2, 1, 64, 64], &device);

        let output_small = conv_small.forward_4d(&input);
        let output_large = conv_large.forward_4d(&input);

        // Both should preserve spatial dimensions
        assert_eq!(output_small.shape(), [2, 1, 64, 64]);
        assert_eq!(output_large.shape(), [2, 1, 64, 64]);
    }

    #[test]
    #[should_panic(expected = "Use forward_4d for 2D spectral convolution")]
    fn test_spectral_conv2d_forward_panics() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 8, 8, &device);

        let input = GenericTensor::randn([4, 2], &device);
        let _ = conv.forward(&input);
    }

    // ==================== forward_4d Implementation Tests ====================

    #[test]
    fn test_forward_4d_output_channels() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(3, 8, 12, 12, &device);

        let input = GenericTensor::randn([2, 3, 64, 64], &device);
        let output = conv.forward_4d(&input);

        // Output should transform from in_channels to out_channels
        assert_eq!(output.shape(), [2, 8, 64, 64]);
    }

    #[test]
    fn test_forward_4d_non_zero_output() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(2, 4, 8, 8, &device);

        // Create non-zero input
        let input = GenericTensor::ones([1, 2, 32, 32], &device);
        let output = conv.forward_4d(&input);

        // FFT-based transformation should produce non-zero output
        let output_data = output.to_vec();
        let non_zero_count = output_data.iter().filter(|&&x| x.abs() > 1e-6).count();

        assert!(
            non_zero_count > 0,
            "FFT transformation should produce non-zero outputs"
        );
    }

    #[test]
    fn test_forward_4d_complex_multiplication() {
        let device = CpuDevice::new();

        // Use small modes for easier verification
        let conv = SpectralConv2d::<CpuBackend>::new(1, 1, 4, 4, &device);

        // Create a simple input pattern (constant)
        let input = GenericTensor::ones([1, 1, 16, 16], &device);
        let output = conv.forward_4d(&input);

        // Output shape should be preserved
        assert_eq!(output.shape(), [1, 1, 16, 16]);

        // The output should be different from input (due to complex multiplication)
        let input_data = input.to_vec();
        let output_data = output.to_vec();

        let mut differences = 0;
        for (inp, out) in input_data.iter().zip(output_data.iter()) {
            if (inp - out).abs() > 1e-5 {
                differences += 1;
            }
        }

        assert!(
            differences > 0,
            "Complex multiplication in frequency domain should modify values"
        );
    }

    #[test]
    fn test_forward_4d_multiple_batches() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(2, 3, 8, 8, &device);

        let input = GenericTensor::randn([5, 2, 32, 32], &device);
        let output = conv.forward_4d(&input);

        // Should handle multiple batches correctly
        assert_eq!(output.shape(), [5, 3, 32, 32]);
    }

    #[test]
    fn test_forward_4d_different_spatial_sizes() {
        let device = CpuDevice::new();
        let conv = SpectralConv2d::<CpuBackend>::new(1, 2, 4, 6, &device);

        // Non-square spatial dimensions
        let input = GenericTensor::randn([1, 1, 24, 36], &device);
        let output = conv.forward_4d(&input);

        // Spatial dimensions should be preserved
        assert_eq!(output.shape(), [1, 2, 24, 36]);
    }
}