//! Fourier Neural Operator (FNO) architectures.
//!
//! FNO is a neural operator architecture that learns mappings between function
//! spaces using spectral convolutions in the Fourier domain.
//!
//! ## Architecture
//!
//! FNO consists of:
//! 1. Lifting layer: P (input → latent space)
//! 2. Multiple Fourier layers: spectral conv + residual + activation
//! 3. Projection layer: Q (latent space → output)
//!
//! ## References
//!
//! Li, Z., et al. (2020). "Fourier Neural Operator for Parametric Partial
//! Differential Equations." arXiv:2010.08895

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

use crate::Result;
use crate::layers::{GridPositionalEncoding, LiftingMLP, Projection};
use crate::spectral::SpectralConv2d;
use crate::weights::FNO2dWeights;

/// 1D Fourier Neural Operator.
#[derive(Debug)]
pub struct FNO1d<B: Backend<FloatElem = f32>> {
    _device: B::Device,
}

impl<B: Backend<FloatElem = f32>> FNO1d<B> {
    /// Create a new 1D FNO.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `width` - Width of the latent space
    /// * `n_modes` - Number of Fourier modes
    /// * `device` - Device to create the model on
    pub fn new(
        _in_channels: usize,
        _out_channels: usize,
        _width: usize,
        _n_modes: usize,
        device: &B::Device,
    ) -> Result<Self> {
        Ok(Self {
            _device: device.clone(),
        })
    }
}

/// 2D Fourier Neural Operator.
///
/// FNO2d learns mappings between 2D function spaces by alternating between
/// spectral convolutions (in Fourier space) and pointwise operations (in physical space).
///
/// ## Architecture (neuraloperator v2.0 compatible)
///
/// 1. Positional encoding: Adds 2D grid coordinates to input
/// 2. Lifting MLP (P): Two-layer MLP projects (in_channels + 2) to width
/// 3. N Fourier blocks: Each has spectral_conv + 1x1 conv + residual + GELU
/// 4. Projection layers (Q): Projects width -> 128 -> output channels
#[derive(Debug)]
pub struct FNO2d<B: Backend<FloatElem = f32>> {
    /// Grid positional encoding layer
    positional_encoding: GridPositionalEncoding<B>,
    /// Two-layer MLP lifting (neuraloperator v2.0 compatible)
    lifting: LiftingMLP<B>,
    /// Spectral convolution layers
    spectral_convs: Vec<SpectralConv2d<B>>,
    /// 1x1 convolution layers (skip connections)
    convs: Vec<GenericLinear<B>>,
    /// First projection layer (width -> hidden)
    projection1: Projection<B>,
    /// Second projection layer (hidden -> out_channels)
    projection2: Projection<B>,
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> FNO2d<B> {
    /// Create a new 2D FNO with default number of layers (4).
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `width` - Width of the latent space
    /// * `n_modes` - Number of Fourier modes (same for both dimensions)
    /// * `device` - Device to create the model on
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        width: usize,
        n_modes: usize,
        device: &B::Device,
    ) -> Result<Self> {
        Self::new_with_layers(in_channels, out_channels, width, n_modes, 4, device)
    }

    /// Create a new 2D FNO with custom number of Fourier blocks.
    ///
    /// # Arguments
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `width` - Width of the latent space
    /// * `n_modes` - Number of Fourier modes (same for both dimensions)
    /// * `n_layers` - Number of Fourier blocks
    /// * `device` - Device to create the model on
    /// Hidden dimension for projection layers (matches neuraloperator)
    const PROJECTION_HIDDEN: usize = 128;

    /// Creates a new FNO2d model with a specified number of Fourier layers.
    ///
    /// # Arguments
    ///
    /// * `in_channels` - Number of input channels
    /// * `out_channels` - Number of output channels
    /// * `width` - Width of hidden channels in Fourier blocks
    /// * `n_modes` - Number of Fourier modes to keep
    /// * `n_layers` - Number of Fourier blocks
    /// * `device` - Device to create the model on
    pub fn new_with_layers(
        in_channels: usize,
        out_channels: usize,
        width: usize,
        n_modes: usize,
        n_layers: usize,
        device: &B::Device,
    ) -> Result<Self> {
        // 1. Create positional encoding layer
        let positional_encoding = GridPositionalEncoding::new(device);

        // 2. Create lifting MLP: (in_channels + 2) -> width
        // The +2 accounts for x and y positional encoding channels
        let lifting = LiftingMLP::new(in_channels + 2, width, device);

        // 3. Create n_layers spectral convolution layers: width -> width
        let mut spectral_convs = Vec::with_capacity(n_layers);
        for _ in 0..n_layers {
            spectral_convs.push(SpectralConv2d::new(width, width, n_modes, n_modes, device));
        }

        // 4. Create n_layers 1x1 convolution layers (skip connections): width -> width
        // Using GenericLinear as 1x1 conv (pointwise linear transformation)
        let mut convs = Vec::with_capacity(n_layers);
        for _ in 0..n_layers {
            convs.push(GenericLinear::new(width, width, true, device));
        }

        // 5. Create two-stage projection: width -> 128 -> out_channels
        let projection1 = Projection::new(width, Self::PROJECTION_HIDDEN, device);
        let projection2 = Projection::new(Self::PROJECTION_HIDDEN, out_channels, device);

        Ok(Self {
            positional_encoding,
            lifting,
            spectral_convs,
            convs,
            projection1,
            projection2,
            device: device.clone(),
        })
    }

    /// Create a new FNO2d model from pre-trained weights.
    ///
    /// # Arguments
    /// * `weights` - Pre-trained weights loaded from SafeTensors file
    /// * `device` - Device to create the model on
    ///
    /// # Example
    /// ```rust,ignore
    /// use rtx_neural_operator::{FNO2d, load_fno2d_weights};
    ///
    /// let weights = load_fno2d_weights("./weights/fno2d_darcy/model.safetensors")?;
    /// let fno = FNO2d::<CpuBackend>::from_weights(&weights, &device)?;
    /// ```
    pub fn from_weights(weights: &FNO2dWeights, device: &B::Device) -> Result<Self> {
        let config = &weights.config;

        // Create positional encoding layer
        let positional_encoding = GridPositionalEncoding::new(device);

        // Create lifting MLP from weights
        let lifting = LiftingMLP::from_weights(
            &weights.lifting_fc1_weight,
            &weights.lifting_fc1_bias,
            &weights.lifting_fc2_weight,
            &weights.lifting_fc2_bias,
            config.in_channels + 2, // +2 for positional encoding
            config.width,
            device,
        );

        // Create spectral convolution layers from weights
        let mut spectral_convs = Vec::with_capacity(config.n_layers);
        for sw in &weights.spectral_weights {
            spectral_convs.push(SpectralConv2d::from_weights_dual(
                &sw.weights1_real,
                &sw.weights1_imag,
                &sw.weights2_real,
                &sw.weights2_imag,
                config.width,
                config.width,
                config.n_modes.0,
                config.n_modes.1,
                device,
            ));
        }

        // Create 1x1 convolution layers (skip connections) from weights
        let mut convs = Vec::with_capacity(config.n_layers);
        for (weight, bias) in &weights.conv_weights {
            convs.push(GenericLinear::from_weights(
                weight,
                Some(bias),
                config.width,
                config.width,
                device,
            ));
        }

        // Create two-stage projection layers: width -> 128 -> out_channels
        let (proj0_weight, proj0_bias) = &weights.projection_weights[0];
        let projection1 = Projection::from_weights(
            proj0_weight,
            proj0_bias,
            config.width,
            Self::PROJECTION_HIDDEN,
            device,
        );

        let (proj1_weight, proj1_bias) = &weights.projection_weights[1];
        let projection2 = Projection::from_weights(
            proj1_weight,
            proj1_bias,
            Self::PROJECTION_HIDDEN,
            config.out_channels,
            device,
        );

        Ok(Self {
            positional_encoding,
            lifting,
            spectral_convs,
            convs,
            projection1,
            projection2,
            device: device.clone(),
        })
    }
}

impl<B: Backend<FloatElem = f32>> FNO2d<B> {
    /// Extract all model weights into a serializable format.
    ///
    /// This is the inverse of `from_weights()` and can be used to save
    /// trained models to SafeTensors format.
    pub fn to_weights(&self, pde_type: &str) -> FNO2dWeights {
        // Extract lifting MLP weights
        let (lifting_fc1_weight, lifting_fc1_bias) = self.lifting.fc1_weights();
        let (lifting_fc2_weight, lifting_fc2_bias) = self.lifting.fc2_weights();

        // Extract spectral convolution weights
        let spectral_weights: Vec<crate::weights::SpectralConvWeights> = self
            .spectral_convs
            .iter()
            .map(|sc| {
                let (w1r, w1i) = sc.weights1_real_imag();
                let (w2r, w2i) = sc.weights2_real_imag();
                crate::weights::SpectralConvWeights {
                    weights1_real: w1r,
                    weights1_imag: w1i,
                    weights2_real: w2r,
                    weights2_imag: w2i,
                }
            })
            .collect();

        // Extract 1x1 convolution (skip connection) weights
        let conv_weights: Vec<(Vec<f32>, Vec<f32>)> = self
            .convs
            .iter()
            .map(|linear| {
                let weight = linear.weight().to_vec();
                let bias = linear
                    .bias()
                    .map(rtx_tensor::GenericTensor::to_vec)
                    .unwrap_or_default();
                (weight, bias)
            })
            .collect();

        // Extract projection weights
        let proj1 = self.projection1.weights();
        let proj2 = self.projection2.weights();
        let projection_weights = vec![proj1, proj2];

        // Determine config from model structure
        let width = self.lifting.out_features();
        let n_layers = self.spectral_convs.len();
        let n_modes = if !self.spectral_convs.is_empty() {
            (
                self.spectral_convs[0].n_modes_h(),
                self.spectral_convs[0].n_modes_w(),
            )
        } else {
            (12, 12)
        };
        let in_channels = self.lifting.in_features().saturating_sub(2); // Remove positional encoding
        let out_channels = self.projection2.out_features();

        let config = crate::weights::FNO2dConfig {
            in_channels,
            out_channels,
            width,
            n_modes,
            n_layers,
            pde_type: pde_type.to_string(),
        };

        FNO2dWeights {
            lifting_fc1_weight,
            lifting_fc1_bias,
            lifting_fc2_weight,
            lifting_fc2_bias,
            spectral_weights,
            conv_weights,
            projection_weights,
            config,
        }
    }
}

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

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

impl<B: Backend<FloatElem = f32>> GenericModule4D<B> for FNO2d<B> {
    fn forward_4d(&self, input: &GenericTensor<B, 4>) -> GenericTensor<B, 4> {
        // Step 1: Add positional encoding
        // input: [batch, in_channels, height, width]
        // with_pos: [batch, in_channels + 2, height, width]
        let with_pos = self.positional_encoding.forward_4d(input);

        // Step 2: Lift input to latent space using MLP
        // with_pos: [batch, in_channels + 2, height, width]
        // x: [batch, width, height, width_spatial]
        let mut x = self.lifting.forward_4d(&with_pos);

        // Step 3: Apply Fourier blocks
        let n_layers = self.spectral_convs.len();
        for i in 0..n_layers {
            // Apply spectral convolution in Fourier domain
            let spectral_out = self.spectral_convs[i].forward_4d(&x);

            // Apply 1x1 convolution (skip connection) in physical space
            // For 1x1 conv, we need to apply GenericLinear pointwise across spatial dimensions
            let conv_out = self.apply_pointwise_linear(&x, &self.convs[i]);

            // Residual connection: combine spectral and conv paths
            let combined = spectral_out.add(&conv_out);

            // Apply GELU activation (except on last layer, matching neuraloperator v2.0)
            if i < n_layers - 1 {
                x = combined.gelu();
            } else {
                x = combined;
            }
        }

        // Step 4: Project back to output space
        // Two-stage projection: width -> 128 -> out_channels with GELU
        let x = self.projection1.forward_4d(&x);
        let x = x.gelu();
        self.projection2.forward_4d(&x)
    }
}

impl<B: Backend<FloatElem = f32>> FNO2d<B> {
    /// Apply pointwise linear transformation (1x1 convolution) to 4D tensor.
    ///
    /// This mimics a 1x1 convolution by applying the linear layer independently
    /// at each spatial location.
    fn apply_pointwise_linear(
        &self,
        input: &GenericTensor<B, 4>,
        linear: &GenericLinear<B>,
    ) -> GenericTensor<B, 4> {
        let shape = input.shape();
        let batch = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        // Reshape: [batch, channels, height, width] -> [batch, height, width, channels]
        // First swap channels and height: [B, C, H, W] -> [B, H, C, W]
        let step1 = input.swap_dims(1, 2);
        // Then swap channels and width: [B, H, C, W] -> [B, H, W, C]
        let permuted = step1.swap_dims(2, 3);

        // Reshape to [batch * height * width, channels]
        let reshaped = permuted.reshape([batch * height * width, channels]);

        // Apply linear transformation
        let transformed = linear.forward(&reshaped);

        // Get output channels
        let out_channels = transformed.shape()[1];

        // Reshape back to [batch, height, width, out_channels]
        let output_nhwc = transformed.reshape([batch, height, width, out_channels]);

        // Permute back to [batch, out_channels, height, width]
        // First swap width and channels: [B, H, W, C] -> [B, H, C, W]
        let step1 = output_nhwc.swap_dims(2, 3);
        // Then swap height and channels: [B, H, C, W] -> [B, C, H, W]
        step1.swap_dims(1, 2)
    }
}

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

    #[test]
    fn test_fno2d_creation() {
        let device = CpuDevice::new();
        let result = FNO2d::<CpuBackend>::new(1, 1, 32, 12, &device);
        assert!(result.is_ok());
    }

    #[test]
    fn test_fno2d_forward_shape() {
        let device = CpuDevice::new();
        let fno = FNO2d::<CpuBackend>::new(2, 3, 64, 16, &device).unwrap();

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

        // For now, output has same shape as input (placeholder)
        assert_eq!(output.shape()[0], 4); // batch
        assert_eq!(output.shape()[2], 32); // height
        assert_eq!(output.shape()[3], 32); // width
    }

    // ==================== TDD Tests for Full FNO2d Implementation ====================

    #[test]
    fn test_fno2d_output_channels() {
        let device = CpuDevice::new();
        let fno = FNO2d::<CpuBackend>::new(3, 8, 64, 12, &device).unwrap();

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

        // Output should transform from in_channels to out_channels
        assert_eq!(output.shape()[0], 2); // batch
        assert_eq!(output.shape()[1], 8); // out_channels
        assert_eq!(output.shape()[2], 32); // height preserved
        assert_eq!(output.shape()[3], 32); // width preserved
    }

    #[test]
    fn test_fno2d_non_zero_output() {
        let device = CpuDevice::new();
        let fno = FNO2d::<CpuBackend>::new(2, 4, 32, 8, &device).unwrap();

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

        // FNO should produce meaningful 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,
            "FNO should produce non-zero outputs for non-zero inputs"
        );

        // Output should not be all zeros
        let all_zeros = output_data.iter().all(|&x| x.abs() < 1e-6);
        assert!(!all_zeros, "FNO output should not be all zeros");
    }

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

        // Test with different latent widths
        let fno_small = FNO2d::<CpuBackend>::new(1, 1, 16, 4, &device).unwrap();
        let fno_large = FNO2d::<CpuBackend>::new(1, 1, 128, 4, &device).unwrap();

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

        let output_small = fno_small.forward_4d(&input);
        let output_large = fno_large.forward_4d(&input);

        // Both should produce valid outputs with correct shape
        assert_eq!(output_small.shape(), [1, 1, 32, 32]);
        assert_eq!(output_large.shape(), [1, 1, 32, 32]);

        // Outputs should be non-zero
        let data_small = output_small.to_vec();
        let data_large = output_large.to_vec();

        assert!(data_small.iter().any(|&x| x.abs() > 1e-6));
        assert!(data_large.iter().any(|&x| x.abs() > 1e-6));
    }

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

        // Test with different number of Fourier blocks
        let fno_2layers = FNO2d::<CpuBackend>::new_with_layers(2, 3, 32, 8, 2, &device).unwrap();
        let fno_6layers = FNO2d::<CpuBackend>::new_with_layers(2, 3, 32, 8, 6, &device).unwrap();

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

        let output_2 = fno_2layers.forward_4d(&input);
        let output_6 = fno_6layers.forward_4d(&input);

        // Both should produce valid outputs
        assert_eq!(output_2.shape(), [1, 3, 24, 24]);
        assert_eq!(output_6.shape(), [1, 3, 24, 24]);

        // Outputs should be different (different network depths)
        let data_2 = output_2.to_vec();
        let data_6 = output_6.to_vec();

        assert!(data_2.iter().any(|&x| x.abs() > 1e-6));
        assert!(data_6.iter().any(|&x| x.abs() > 1e-6));
    }
}