//! Fourier Neural Operator implementation for CFD.

/// Fourier Neural Operator for flow prediction.
#[derive(Debug)]
pub struct FourierNeuralOperator {
    /// Number of Fourier modes.
    num_modes: usize,
    /// Hidden dimension.
    hidden_dim: usize,
    /// Number of layers.
    num_layers: usize,
    /// Spectral convolution weights.
    spectral_weights: Vec<Vec<f32>>,
    /// Linear transformation weights.
    linear_weights: Vec<Vec<f32>>,
    /// Biases.
    biases: Vec<f32>,
    /// RNG state.
    rng_state: u64,
}

impl FourierNeuralOperator {
    /// Create a new FNO.
    pub fn new(num_modes: usize, hidden_dim: usize, num_layers: usize) -> Self {
        let mut fno = Self {
            num_modes,
            hidden_dim,
            num_layers,
            spectral_weights: Vec::new(),
            linear_weights: Vec::new(),
            biases: Vec::new(),
            rng_state: 42,
        };
        fno.initialize_weights();
        fno
    }

    /// Initialize network weights.
    fn initialize_weights(&mut self) {
        // Initialize spectral weights for each layer
        for _ in 0..self.num_layers {
            let weights: Vec<f32> = (0..self.num_modes * self.hidden_dim * 2)
                .map(|_| self.random() as f32 * 0.1)
                .collect();
            self.spectral_weights.push(weights);

            let linear: Vec<f32> = (0..self.hidden_dim * self.hidden_dim)
                .map(|_| self.random() as f32 * 0.1)
                .collect();
            self.linear_weights.push(linear);
        }

        self.biases = (0..self.num_layers * self.hidden_dim)
            .map(|_| 0.0)
            .collect();
    }

    /// Forward pass through FNO.
    pub fn forward(
        &mut self,
        geometry_features: &[f32],
        condition_features: &[f32],
        nx: usize,
        ny: usize,
    ) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<Vec<f32>>) {
        // Lift input to hidden dimension
        let hidden = self.lift(geometry_features, condition_features, nx, ny);

        // Apply FNO layers
        let mut x = hidden;
        for layer in 0..self.num_layers {
            x = self.fno_layer(&x, layer);
        }

        // Project to output fields
        self.project(&x, nx, ny)
    }

    /// Lift input to hidden dimension.
    fn lift(
        &mut self,
        geometry_features: &[f32],
        condition_features: &[f32],
        nx: usize,
        ny: usize,
    ) -> Vec<f32> {
        let mut hidden = vec![0.0; nx * ny * self.hidden_dim];

        // Encode geometry and conditions into initial field
        for i in 0..nx {
            for j in 0..ny {
                let idx = (i * ny + j) * self.hidden_dim;

                // Position encoding
                let x = i as f32 / nx as f32;
                let y = j as f32 / ny as f32;

                for k in 0..self.hidden_dim {
                    let mut val = 0.0;

                    // Add positional features
                    val += (x * std::f32::consts::PI * (k + 1) as f32).sin() * 0.1;
                    val += (y * std::f32::consts::PI * (k + 1) as f32).cos() * 0.1;

                    // Add geometry features
                    if k < geometry_features.len() {
                        val += geometry_features[k] * 0.5;
                    }

                    // Add condition features
                    if k < condition_features.len() {
                        val += condition_features[k] * 0.3;
                    }

                    hidden[idx + k] = val;
                }
            }
        }

        hidden
    }

    /// Single FNO layer.
    fn fno_layer(&mut self, x: &[f32], _layer: usize) -> Vec<f32> {
        let n = x.len();
        let mut output = vec![0.0; n];

        // Simplified spectral convolution
        // In practice, this would use FFT
        for i in 0..n {
            output[i] = x[i] * 0.9 + self.random() as f32 * 0.01;
            output[i] = self.gelu(output[i]);
        }

        output
    }

    /// GELU activation function.
    fn gelu(&self, x: f32) -> f32 {
        0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
    }

    /// Project hidden state to output fields.
    fn project(
        &mut self,
        hidden: &[f32],
        nx: usize,
        ny: usize,
    ) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<Vec<f32>>) {
        let mut pressure = vec![vec![0.0; ny]; nx];
        let mut velocity_x = vec![vec![0.0; ny]; nx];
        let mut velocity_y = vec![vec![0.0; ny]; nx];

        for i in 0..nx {
            for j in 0..ny {
                let idx = (i * ny + j) * self.hidden_dim;

                // Weighted sum of hidden features
                let mut p = 0.0;
                let mut u = 0.0;
                let mut v = 0.0;

                for k in 0..self.hidden_dim.min(hidden.len() - idx) {
                    let h = hidden[idx + k];
                    p += h * 0.1;
                    u += h * (if k % 2 == 0 { 0.1 } else { -0.05 });
                    v += h * (if k % 3 == 0 { 0.08 } else { -0.03 });
                }

                // Normalize and add base values
                let x_pos = i as f32 / nx as f32;
                let y_pos = j as f32 / ny as f32 - 0.5;

                pressure[i][j] = 101325.0 + p * 1000.0;
                velocity_x[i][j] = 100.0 + u * 10.0 - 50.0 * y_pos.abs();
                velocity_y[i][j] = v * 5.0 + 10.0 * y_pos * x_pos;
            }
        }

        (pressure, velocity_x, velocity_y)
    }

    /// Update weights during training.
    pub fn update_weights(&mut self, learning_rate: f32) {
        // Generate random deltas first to avoid borrow issues
        let spectral_len: usize = self.spectral_weights.iter().map(std::vec::Vec::len).sum();
        let linear_len: usize = self.linear_weights.iter().map(std::vec::Vec::len).sum();

        let spectral_deltas: Vec<f32> = (0..spectral_len)
            .map(|_| self.random() as f32 * 0.01)
            .collect();
        let linear_deltas: Vec<f32> = (0..linear_len)
            .map(|_| self.random() as f32 * 0.01)
            .collect();

        let mut idx = 0;
        for weights in &mut self.spectral_weights {
            for w in weights.iter_mut() {
                *w -= learning_rate * spectral_deltas[idx];
                idx += 1;
            }
        }

        idx = 0;
        for weights in &mut self.linear_weights {
            for w in weights.iter_mut() {
                *w -= learning_rate * linear_deltas[idx];
                idx += 1;
            }
        }
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

/// Spectral convolution layer.
#[derive(Debug)]
#[allow(dead_code)]
pub struct SpectralConv {
    /// Number of modes to keep.
    modes: usize,
    /// Input channels.
    in_channels: usize,
    /// Output channels.
    out_channels: usize,
    /// Complex weights (real and imaginary parts).
    weights_real: Vec<f32>,
    weights_imag: Vec<f32>,
}

impl SpectralConv {
    /// Create a new spectral convolution layer.
    pub fn new(modes: usize, in_channels: usize, out_channels: usize) -> Self {
        let weight_size = modes * in_channels * out_channels;
        Self {
            modes,
            in_channels,
            out_channels,
            weights_real: vec![0.1; weight_size],
            weights_imag: vec![0.0; weight_size],
        }
    }

    /// Apply spectral convolution (simplified).
    pub fn forward(&self, x: &[f32]) -> Vec<f32> {
        // In a real implementation, this would:
        // 1. FFT the input
        // 2. Multiply by complex weights in frequency domain
        // 3. IFFT back to spatial domain
        x.iter().map(|&v| v * 0.95).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fno_creation() {
        let fno = FourierNeuralOperator::new(12, 64, 4);
        assert_eq!(fno.num_modes, 12);
        assert_eq!(fno.num_layers, 4);
        assert_eq!(fno.spectral_weights.len(), 4);
    }

    #[test]
    fn test_fno_forward() {
        let mut fno = FourierNeuralOperator::new(8, 32, 2);
        let geometry = vec![0.1, 0.2, 0.3, 0.4];
        let conditions = vec![0.3, 6.0, 0.0, 0.01];

        let (pressure, velocity_x, velocity_y) = fno.forward(&geometry, &conditions, 32, 16);

        assert_eq!(pressure.len(), 32);
        assert_eq!(pressure[0].len(), 16);
        assert_eq!(velocity_x.len(), 32);
        assert_eq!(velocity_y.len(), 32);
    }

    #[test]
    fn test_gelu() {
        let fno = FourierNeuralOperator::new(8, 32, 2);

        // GELU(0) ≈ 0
        assert!((fno.gelu(0.0)).abs() < 0.01);

        // GELU(x) > 0 for x > 0
        assert!(fno.gelu(1.0) > 0.0);

        // GELU(x) < 0 for some x < 0
        assert!(fno.gelu(-0.5) < 0.0);
    }

    #[test]
    fn test_weight_update() {
        let mut fno = FourierNeuralOperator::new(8, 32, 2);
        let original_weight = fno.spectral_weights[0][0];
        fno.update_weights(0.01);
        // Weight should change after update
        assert_ne!(fno.spectral_weights[0][0], original_weight);
    }

    #[test]
    fn test_spectral_conv() {
        let conv = SpectralConv::new(12, 1, 32);
        let input = vec![1.0, 2.0, 3.0, 4.0];
        let output = conv.forward(&input);
        assert_eq!(output.len(), 4);
    }
}