//! Fourier Neural Operator implementation.
//!
//! FNO learns operators by performing convolutions in Fourier space,
//! enabling resolution-independent learning.

use crate::{NeuralOpError, NeuralOperatorTrainer};
use neuralop_studio_shared::{
    EvaluationMetrics, OperatorConfig, OperatorType, PDEDefinition, TrainingConfig,
    TrainingProgress, TrainingResult,
};

/// Fourier Neural Operator.
#[derive(Debug)]
pub struct FourierNeuralOperator {
    /// Configuration.
    config: OperatorConfig,
    /// Trained weights (simplified representation).
    weights: Vec<f64>,
    /// Is trained.
    is_trained: bool,
    /// RNG state.
    rng_state: u64,
}

impl FourierNeuralOperator {
    /// Create a new FNO.
    pub fn new(config: OperatorConfig) -> Self {
        let num_params = Self::calculate_params(&config);
        Self {
            config,
            weights: vec![0.0; num_params],
            is_trained: false,
            rng_state: 42,
        }
    }

    /// Calculate number of parameters.
    fn calculate_params(config: &OperatorConfig) -> usize {
        let hidden = config.hidden_dim;
        let layers = config.num_layers;
        let modes = config.fourier_modes.unwrap_or(12);

        // Lifting + spectral layers + projection
        let lifting = hidden * 2; // Input to hidden
        let spectral = layers * (modes * modes * hidden * 2 + hidden * hidden); // Complex weights
        let projection = hidden * 2; // Hidden to output

        lifting + spectral + projection
    }

    /// Initialize weights with Xavier initialization.
    fn initialize_weights(&mut self) {
        let num_weights = self.weights.len();
        let random_values: Vec<f64> = (0..num_weights)
            .map(|_| self.random_normal() * 0.01)
            .collect();
        for (weight, value) in self.weights.iter_mut().zip(random_values) {
            *weight = value;
        }
    }

    /// Simple 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
    }

    /// Box-Muller for normal distribution.
    fn random_normal(&mut self) -> f64 {
        let u1 = self.random() + 1e-10;
        let u2 = self.random();
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    }

    /// Forward pass (simplified).
    fn forward(&self, input: &[f64]) -> Vec<f64> {
        // Simplified FNO forward pass
        // In reality, this would involve FFT operations
        let output_size = input.len();
        let mut output = vec![0.0; output_size];

        // Apply lifting
        let hidden_dim = self.config.hidden_dim;
        let scale = (hidden_dim as f64).sqrt().recip();

        for (i, out) in output.iter_mut().enumerate() {
            let mut sum = 0.0;
            for (j, &inp) in input.iter().enumerate().take(hidden_dim.min(input.len())) {
                let weight_idx = (i + j) % self.weights.len();
                sum += inp * self.weights[weight_idx] * scale;
            }
            *out = sum.tanh(); // Activation
        }

        output
    }

    /// Compute loss.
    fn compute_loss(&self, predictions: &[Vec<f64>], targets: &[Vec<f64>]) -> f64 {
        if predictions.is_empty() {
            return 0.0;
        }

        let mut total_loss = 0.0;
        for (pred, target) in predictions.iter().zip(targets.iter()) {
            let mse: f64 = pred
                .iter()
                .zip(target.iter())
                .map(|(p, t)| (p - t).powi(2))
                .sum::<f64>()
                / pred.len() as f64;
            total_loss += mse;
        }
        total_loss / predictions.len() as f64
    }

    /// Gradient descent step.
    fn gradient_step(&mut self, learning_rate: f64) {
        // Simplified gradient update (random perturbation for demo)
        let num_weights = self.weights.len();
        let gradients: Vec<f64> = (0..num_weights)
            .map(|_| self.random_normal() * 0.01)
            .collect();
        for (weight, gradient) in self.weights.iter_mut().zip(gradients) {
            *weight -= learning_rate * gradient;
        }
    }
}

impl NeuralOperatorTrainer for FourierNeuralOperator {
    fn train(
        &mut self,
        pde: &PDEDefinition,
        config: &TrainingConfig,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<TrainingResult, NeuralOpError> {
        self.rng_state = config.seed.unwrap_or(42);
        self.initialize_weights();

        let grid_size: usize = pde.domain.resolution.iter().product();
        let start_time = std::time::Instant::now();

        let mut train_loss_history = Vec::new();
        let mut val_loss_history = Vec::new();
        let mut best_val_loss = f64::MAX;
        let mut best_epoch = 0;

        // Generate synthetic training data
        let train_inputs: Vec<Vec<f64>> = (0..config.num_train_samples)
            .map(|i| {
                self.rng_state = config.seed.unwrap_or(42) + i as u64;
                (0..grid_size)
                    .map(|j| {
                        let x =
                            (j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
                        let y = if pde.domain.dimensions > 1 {
                            (j / pde.domain.resolution[0]) as f64
                                / pde.domain.resolution.get(1).copied().unwrap_or(1) as f64
                        } else {
                            0.0
                        };
                        (std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin()
                            + self.random_normal() * 0.1
                    })
                    .collect()
            })
            .collect();

        let train_targets: Vec<Vec<f64>> = train_inputs
            .iter()
            .map(|input| {
                input
                    .iter()
                    .map(|&v| v * 0.5 + self.random_normal() * 0.01)
                    .collect()
            })
            .collect();

        // Validation data
        let val_inputs: Vec<Vec<f64>> = (0..config.num_val_samples)
            .map(|i| {
                self.rng_state = config.seed.unwrap_or(42) + 10000 + i as u64;
                (0..grid_size)
                    .map(|j| {
                        let x =
                            (j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
                        (std::f64::consts::PI * x).sin() + self.random_normal() * 0.1
                    })
                    .collect()
            })
            .collect();

        let val_targets: Vec<Vec<f64>> = val_inputs
            .iter()
            .map(|input| input.iter().map(|&v| v * 0.5).collect())
            .collect();

        // Training loop
        let num_batches = config.num_train_samples.div_ceil(config.batch_size);

        for epoch in 0..config.epochs {
            let mut epoch_loss = 0.0;

            for batch in 0..num_batches {
                let batch_start = batch * config.batch_size;
                let batch_end = (batch_start + config.batch_size).min(config.num_train_samples);

                let batch_inputs: Vec<_> = train_inputs[batch_start..batch_end].to_vec();
                let batch_targets: Vec<_> = train_targets[batch_start..batch_end].to_vec();

                // Forward pass
                let predictions: Vec<Vec<f64>> = batch_inputs
                    .iter()
                    .map(|input| self.forward(input))
                    .collect();

                // Compute loss
                let batch_loss = self.compute_loss(&predictions, &batch_targets);
                epoch_loss += batch_loss;

                // Backward pass (simplified)
                let lr =
                    config.learning_rate * (1.0 - epoch as f64 / config.epochs as f64).max(0.1); // Learning rate decay
                self.gradient_step(lr);

                // Progress callback
                if let Some(ref callback) = progress_callback {
                    callback(TrainingProgress {
                        epoch: epoch + 1,
                        total_epochs: config.epochs,
                        batch: batch + 1,
                        total_batches: num_batches,
                        train_loss: batch_loss,
                        val_loss: None,
                        physics_loss: None,
                        relative_error: None,
                        learning_rate: lr,
                        elapsed_seconds: start_time.elapsed().as_secs_f64(),
                    });
                }
            }

            let avg_train_loss = epoch_loss / num_batches as f64;
            train_loss_history.push(avg_train_loss);

            // Validation
            let val_predictions: Vec<Vec<f64>> =
                val_inputs.iter().map(|input| self.forward(input)).collect();
            let val_loss = self.compute_loss(&val_predictions, &val_targets);
            val_loss_history.push(val_loss);

            if val_loss < best_val_loss {
                best_val_loss = val_loss;
                best_epoch = epoch + 1;
            }
        }

        self.is_trained = true;

        // Generate test data and evaluate
        let test_inputs: Vec<Vec<f64>> = (0..config.num_test_samples)
            .map(|i| {
                self.rng_state = config.seed.unwrap_or(42) + 20000 + i as u64;
                (0..grid_size)
                    .map(|j| {
                        let x =
                            (j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
                        (std::f64::consts::PI * x).sin() + self.random_normal() * 0.1
                    })
                    .collect()
            })
            .collect();

        let test_targets: Vec<Vec<f64>> = test_inputs
            .iter()
            .map(|input| input.iter().map(|&v| v * 0.5).collect())
            .collect();

        let test_metrics = self.evaluate(&test_inputs, &test_targets);

        Ok(TrainingResult {
            final_train_loss: *train_loss_history.last().unwrap_or(&0.0),
            final_val_loss: *val_loss_history.last().unwrap_or(&0.0),
            best_epoch,
            train_loss_history,
            val_loss_history,
            test_metrics,
            total_time_seconds: start_time.elapsed().as_secs_f64(),
            num_parameters: self.weights.len(),
        })
    }

    fn predict(&self, input: &[f64], query_points: &[Vec<f64>]) -> Result<Vec<f64>, NeuralOpError> {
        if !self.is_trained {
            return Err(NeuralOpError::PredictionFailed(
                "Model not trained".to_string(),
            ));
        }

        let full_output = self.forward(input);

        // Interpolate at query points
        let output: Vec<f64> = query_points
            .iter()
            .map(|point| {
                // Simple nearest-neighbor interpolation
                let idx = if point.is_empty() {
                    0
                } else {
                    let x = point[0].clamp(0.0, 1.0);
                    (x * (full_output.len() - 1) as f64).round() as usize
                };
                full_output.get(idx).copied().unwrap_or(0.0)
            })
            .collect();

        Ok(output)
    }

    fn evaluate(&self, test_inputs: &[Vec<f64>], test_outputs: &[Vec<f64>]) -> EvaluationMetrics {
        let start = std::time::Instant::now();

        let predictions: Vec<Vec<f64>> = test_inputs
            .iter()
            .map(|input| self.forward(input))
            .collect();

        let inference_time = start.elapsed().as_secs_f64() * 1000.0 / test_inputs.len() as f64;

        // Calculate metrics
        let mut total_mse = 0.0;
        let mut total_relative_l2 = 0.0;
        let mut max_error: f64 = 0.0;

        for (pred, target) in predictions.iter().zip(test_outputs.iter()) {
            let mse: f64 = pred
                .iter()
                .zip(target.iter())
                .map(|(p, t)| (p - t).powi(2))
                .sum::<f64>()
                / pred.len() as f64;

            let target_norm: f64 = target.iter().map(|t| t.powi(2)).sum::<f64>().sqrt();
            let error_norm: f64 = pred
                .iter()
                .zip(target.iter())
                .map(|(p, t)| (p - t).powi(2))
                .sum::<f64>()
                .sqrt();

            let relative_l2 = if target_norm > 1e-10 {
                error_norm / target_norm
            } else {
                error_norm
            };

            let local_max: f64 = pred
                .iter()
                .zip(target.iter())
                .map(|(p, t)| (p - t).abs())
                .fold(0.0, f64::max);

            total_mse += mse;
            total_relative_l2 += relative_l2;
            max_error = max_error.max(local_max);
        }

        let num_samples = test_inputs.len();
        EvaluationMetrics {
            mse: total_mse / num_samples as f64,
            relative_l2: total_relative_l2 / num_samples as f64,
            max_error,
            physics_residual: None,
            num_samples,
            avg_inference_time_ms: inference_time,
        }
    }

    fn num_parameters(&self) -> usize {
        self.weights.len()
    }

    fn operator_type(&self) -> OperatorType {
        OperatorType::FNO
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use neuralop_studio_shared::{
        sample_fno_config, sample_poisson_problem, sample_training_config,
    };

    #[test]
    fn test_fno_creation() {
        let config = sample_fno_config();
        let fno = FourierNeuralOperator::new(config);
        assert!(!fno.is_trained);
        assert!(!fno.weights.is_empty());
    }

    #[test]
    fn test_fno_forward() {
        let config = sample_fno_config();
        let fno = FourierNeuralOperator::new(config);

        let input = vec![1.0; 64];
        let output = fno.forward(&input);
        assert_eq!(output.len(), 64);
    }

    #[test]
    fn test_fno_training() {
        let config = sample_fno_config();
        let mut fno = FourierNeuralOperator::new(config);

        let mut training_config = sample_training_config();
        training_config.epochs = 3;
        training_config.num_train_samples = 20;
        training_config.num_val_samples = 5;
        training_config.num_test_samples = 5;

        let pde = sample_poisson_problem();
        let result = fno.train(&pde, &training_config, None);

        assert!(result.is_ok());
        assert!(fno.is_trained);
    }

    #[test]
    fn test_fno_predict_after_training() {
        let config = sample_fno_config();
        let mut fno = FourierNeuralOperator::new(config);

        let mut training_config = sample_training_config();
        training_config.epochs = 2;
        training_config.num_train_samples = 10;
        training_config.num_val_samples = 5;
        training_config.num_test_samples = 5;

        let pde = sample_poisson_problem();
        fno.train(&pde, &training_config, None).unwrap();

        let input = vec![1.0; 64];
        let query_points = vec![vec![0.5, 0.5], vec![0.25, 0.75]];
        let result = fno.predict(&input, &query_points);

        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 2);
    }

    #[test]
    fn test_fno_predict_without_training() {
        let config = sample_fno_config();
        let fno = FourierNeuralOperator::new(config);

        let input = vec![1.0; 64];
        let query_points = vec![vec![0.5, 0.5]];
        let result = fno.predict(&input, &query_points);

        assert!(result.is_err());
    }

    #[test]
    fn test_fno_num_parameters() {
        let config = sample_fno_config();
        let fno = FourierNeuralOperator::new(config);
        assert!(fno.num_parameters() > 0);
    }
}