//! Training module with AdamW optimizer

use crate::error::{PINNError, Result};

/// AdamW optimizer state
#[derive(Debug, Clone)]
pub struct AdamW {
    /// Learning rate
    learning_rate: f64,
    /// Beta1 (exponential decay for first moment)
    beta1: f64,
    /// Beta2 (exponential decay for second moment)
    beta2: f64,
    /// Weight decay coefficient
    weight_decay: f64,
    /// Epsilon for numerical stability
    epsilon: f64,
    /// Time step counter
    t: usize,
    /// First moment estimates for weights
    m_weights: Vec<Vec<Vec<f64>>>,
    /// First moment estimates for biases
    m_biases: Vec<Vec<f64>>,
    /// Second moment estimates for weights
    v_weights: Vec<Vec<Vec<f64>>>,
    /// Second moment estimates for biases
    v_biases: Vec<Vec<f64>>,
}

impl AdamW {
    /// Creates a new AdamW optimizer
    ///
    /// # Arguments
    ///
    /// * `learning_rate` - Learning rate (typically 0.001)
    /// * `beta1` - Exponential decay rate for first moment (typically 0.9)
    /// * `beta2` - Exponential decay rate for second moment (typically 0.999)
    /// * `weight_decay` - Weight decay coefficient (typically 0.01)
    /// * `weight_shapes` - Shapes of weight matrices
    /// * `bias_shapes` - Shapes of bias vectors
    ///
    /// # Errors
    ///
    /// Returns an error if hyperparameters are invalid
    pub fn new(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        weight_decay: f64,
        weight_shapes: &[Vec<Vec<f64>>],
        bias_shapes: &[Vec<f64>],
    ) -> Result<Self> {
        if learning_rate <= 0.0 {
            return Err(PINNError::training("learning rate must be positive"));
        }
        if !(0.0..1.0).contains(&beta1) {
            return Err(PINNError::training("beta1 must be in [0, 1)"));
        }
        if !(0.0..1.0).contains(&beta2) {
            return Err(PINNError::training("beta2 must be in [0, 1)"));
        }
        if weight_decay < 0.0 {
            return Err(PINNError::training("weight decay must be non-negative"));
        }

        // Initialize moment estimates with zeros
        let m_weights: Vec<Vec<Vec<f64>>> = weight_shapes
            .iter()
            .map(|layer| layer.iter().map(|row| vec![0.0; row.len()]).collect())
            .collect();

        let m_biases: Vec<Vec<f64>> = bias_shapes
            .iter()
            .map(|layer| vec![0.0; layer.len()])
            .collect();

        let v_weights = m_weights.clone();
        let v_biases = m_biases.clone();

        Ok(Self {
            learning_rate,
            beta1,
            beta2,
            weight_decay,
            epsilon: 1e-8,
            t: 0,
            m_weights,
            m_biases,
            v_weights,
            v_biases,
        })
    }

    /// Performs one optimization step
    ///
    /// # Arguments
    ///
    /// * `weight_grads` - Gradients for weights
    /// * `bias_grads` - Gradients for biases
    ///
    /// # Returns
    ///
    /// Weight and bias updates to apply to the network
    pub fn step(
        &mut self,
        weight_grads: &[Vec<Vec<f64>>],
        bias_grads: &[Vec<f64>],
    ) -> (Vec<Vec<Vec<f64>>>, Vec<Vec<f64>>) {
        self.t += 1;

        // Bias correction terms
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);

        let mut weight_updates = Vec::new();
        let mut bias_updates = Vec::new();

        // Update weights
        for (layer_idx, grads) in weight_grads.iter().enumerate() {
            let mut layer_updates = Vec::new();

            for (i, grad_row) in grads.iter().enumerate() {
                let mut row_updates = Vec::new();

                for (j, &grad) in grad_row.iter().enumerate() {
                    // Update biased first moment estimate
                    self.m_weights[layer_idx][i][j] =
                        self.beta1 * self.m_weights[layer_idx][i][j] + (1.0 - self.beta1) * grad;

                    // Update biased second moment estimate
                    self.v_weights[layer_idx][i][j] = self.beta2 * self.v_weights[layer_idx][i][j]
                        + (1.0 - self.beta2) * grad * grad;

                    // Compute bias-corrected moment estimates
                    let m_hat = self.m_weights[layer_idx][i][j] / bias_correction1;
                    let v_hat = self.v_weights[layer_idx][i][j] / bias_correction2;

                    // Compute update (negative because we want to minimize)
                    let update = -self.learning_rate * (m_hat / (v_hat.sqrt() + self.epsilon));

                    row_updates.push(update);
                }

                layer_updates.push(row_updates);
            }

            weight_updates.push(layer_updates);
        }

        // Update biases
        for (layer_idx, grads) in bias_grads.iter().enumerate() {
            let mut layer_updates = Vec::new();

            for (i, &grad) in grads.iter().enumerate() {
                // Update biased first moment estimate
                self.m_biases[layer_idx][i] =
                    self.beta1 * self.m_biases[layer_idx][i] + (1.0 - self.beta1) * grad;

                // Update biased second moment estimate
                self.v_biases[layer_idx][i] =
                    self.beta2 * self.v_biases[layer_idx][i] + (1.0 - self.beta2) * grad * grad;

                // Compute bias-corrected moment estimates
                let m_hat = self.m_biases[layer_idx][i] / bias_correction1;
                let v_hat = self.v_biases[layer_idx][i] / bias_correction2;

                // Compute update
                let update = -self.learning_rate * (m_hat / (v_hat.sqrt() + self.epsilon));

                layer_updates.push(update);
            }

            bias_updates.push(layer_updates);
        }

        (weight_updates, bias_updates)
    }

    /// Gets the current learning rate
    pub const fn learning_rate(&self) -> f64 {
        self.learning_rate
    }

    /// Sets a new learning rate
    pub fn set_learning_rate(&mut self, lr: f64) {
        self.learning_rate = lr;
    }

    /// Gets the current time step
    pub const fn time_step(&self) -> usize {
        self.t
    }
}

/// Training configuration
#[derive(Debug, Clone)]
pub struct TrainingConfig {
    /// Number of training epochs
    pub num_epochs: usize,
    /// Number of collocation points
    pub num_collocation_points: usize,
    /// Number of boundary points
    pub num_boundary_points: usize,
    /// Learning rate
    pub learning_rate: f64,
    /// Weight for physics loss
    pub physics_weight: f64,
    /// Weight for boundary loss
    pub boundary_weight: f64,
}

impl TrainingConfig {
    /// Creates a new training configuration with validation
    ///
    /// # Errors
    ///
    /// Returns an error if configuration is invalid
    pub fn new(
        num_epochs: usize,
        num_collocation_points: usize,
        num_boundary_points: usize,
        learning_rate: f64,
        physics_weight: f64,
        boundary_weight: f64,
    ) -> Result<Self> {
        if num_epochs == 0 {
            return Err(PINNError::invalid_config("num_epochs must be positive"));
        }
        if num_collocation_points == 0 {
            return Err(PINNError::invalid_config(
                "num_collocation_points must be positive",
            ));
        }
        if num_boundary_points == 0 {
            return Err(PINNError::invalid_config(
                "num_boundary_points must be positive",
            ));
        }
        if learning_rate <= 0.0 {
            return Err(PINNError::invalid_config("learning_rate must be positive"));
        }
        if physics_weight < 0.0 || boundary_weight < 0.0 {
            return Err(PINNError::invalid_config(
                "loss weights must be non-negative",
            ));
        }

        Ok(Self {
            num_epochs,
            num_collocation_points,
            num_boundary_points,
            learning_rate,
            physics_weight,
            boundary_weight,
        })
    }
}

/// Computes mean squared error loss
pub fn mse_loss(predictions: &[f64], targets: &[f64]) -> Result<f64> {
    if predictions.len() != targets.len() {
        return Err(PINNError::training(
            "predictions and targets length mismatch",
        ));
    }

    if predictions.is_empty() {
        return Err(PINNError::training("empty predictions"));
    }

    let sum_squared_error: f64 = predictions
        .iter()
        .zip(targets.iter())
        .map(|(pred, target)| (pred - target).powi(2))
        .sum();

    Ok(sum_squared_error / predictions.len() as f64)
}

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

    #[test]
    fn test_adamw_creation_valid() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3], vec![vec![0.0; 3]; 1]];
        let bias_shapes = vec![vec![0.0; 3], vec![0.0; 1]];

        let optimizer = AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_ok());

        let optimizer = optimizer.unwrap();
        assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);
        assert_eq!(optimizer.time_step(), 0);
    }

    #[test]
    fn test_adamw_creation_invalid_lr() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
        let bias_shapes = vec![vec![0.0; 3]];

        let optimizer = AdamW::new(0.0, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());

        let optimizer = AdamW::new(-0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());
    }

    #[test]
    fn test_adamw_creation_invalid_beta1() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
        let bias_shapes = vec![vec![0.0; 3]];

        let optimizer = AdamW::new(0.001, 1.0, 0.999, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());

        let optimizer = AdamW::new(0.001, -0.1, 0.999, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());
    }

    #[test]
    fn test_adamw_creation_invalid_beta2() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
        let bias_shapes = vec![vec![0.0; 3]];

        let optimizer = AdamW::new(0.001, 0.9, 1.0, 0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());
    }

    #[test]
    fn test_adamw_creation_invalid_weight_decay() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
        let bias_shapes = vec![vec![0.0; 3]];

        let optimizer = AdamW::new(0.001, 0.9, 0.999, -0.01, &weight_shapes, &bias_shapes);
        assert!(optimizer.is_err());
    }

    #[test]
    fn test_adamw_step_increments_time() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 3]];
        let bias_shapes = vec![vec![0.0; 3]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes).unwrap();

        assert_eq!(optimizer.time_step(), 0);

        let weight_grads = vec![vec![vec![0.1; 2]; 3]];
        let bias_grads = vec![vec![0.1; 3]];

        optimizer.step(&weight_grads, &bias_grads);
        assert_eq!(optimizer.time_step(), 1);

        optimizer.step(&weight_grads, &bias_grads);
        assert_eq!(optimizer.time_step(), 2);
    }

    #[test]
    fn test_adamw_step_returns_updates() {
        let weight_shapes = vec![vec![vec![0.0; 2]; 1]];
        let bias_shapes = vec![vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();

        let weight_grads = vec![vec![vec![1.0; 2]; 1]];
        let bias_grads = vec![vec![1.0; 1]];

        let (weight_updates, bias_updates) = optimizer.step(&weight_grads, &bias_grads);

        assert_eq!(weight_updates.len(), 1);
        assert_eq!(weight_updates[0].len(), 1);
        assert_eq!(weight_updates[0][0].len(), 2);

        assert_eq!(bias_updates.len(), 1);
        assert_eq!(bias_updates[0].len(), 1);
    }

    #[test]
    fn test_adamw_step_update_sign() {
        let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
        let bias_shapes = vec![vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();

        // Positive gradient should give negative update (gradient descent)
        let weight_grads = vec![vec![vec![1.0]; 1]];
        let bias_grads = vec![vec![1.0]];

        let (weight_updates, _) = optimizer.step(&weight_grads, &bias_grads);

        assert!(weight_updates[0][0][0] < 0.0);
    }

    #[test]
    fn test_adamw_bias_correction() {
        let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
        let bias_shapes = vec![vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();

        let weight_grads = vec![vec![vec![1.0]; 1]];
        let bias_grads = vec![vec![1.0]];

        let (updates1, _) = optimizer.step(&weight_grads, &bias_grads);
        let (updates2, _) = optimizer.step(&weight_grads, &bias_grads);

        // First step should have larger magnitude due to bias correction
        assert!(updates1[0][0][0].abs() > updates2[0][0][0].abs());
    }

    #[test]
    fn test_adamw_learning_rate_setter() {
        let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
        let bias_shapes = vec![vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();

        assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);

        optimizer.set_learning_rate(0.0005);
        assert_abs_diff_eq!(optimizer.learning_rate(), 0.0005);
    }

    #[test]
    fn test_training_config_creation_valid() {
        let config = TrainingConfig::new(1000, 10000, 100, 0.001, 1.0, 1.0);
        assert!(config.is_ok());

        let config = config.unwrap();
        assert_eq!(config.num_epochs, 1000);
        assert_eq!(config.num_collocation_points, 10000);
        assert_eq!(config.num_boundary_points, 100);
        assert_abs_diff_eq!(config.learning_rate, 0.001);
    }

    #[test]
    fn test_training_config_zero_epochs() {
        let config = TrainingConfig::new(0, 10000, 100, 0.001, 1.0, 1.0);
        assert!(config.is_err());
    }

    #[test]
    fn test_training_config_zero_collocation_points() {
        let config = TrainingConfig::new(1000, 0, 100, 0.001, 1.0, 1.0);
        assert!(config.is_err());
    }

    #[test]
    fn test_training_config_zero_boundary_points() {
        let config = TrainingConfig::new(1000, 10000, 0, 0.001, 1.0, 1.0);
        assert!(config.is_err());
    }

    #[test]
    fn test_training_config_invalid_learning_rate() {
        let config = TrainingConfig::new(1000, 10000, 100, 0.0, 1.0, 1.0);
        assert!(config.is_err());

        let config = TrainingConfig::new(1000, 10000, 100, -0.001, 1.0, 1.0);
        assert!(config.is_err());
    }

    #[test]
    fn test_training_config_negative_weights() {
        let config = TrainingConfig::new(1000, 10000, 100, 0.001, -1.0, 1.0);
        assert!(config.is_err());

        let config = TrainingConfig::new(1000, 10000, 100, 0.001, 1.0, -1.0);
        assert!(config.is_err());
    }

    #[test]
    fn test_mse_loss_perfect_predictions() {
        let predictions = vec![1.0, 2.0, 3.0, 4.0];
        let targets = vec![1.0, 2.0, 3.0, 4.0];

        let loss = mse_loss(&predictions, &targets).unwrap();
        assert_abs_diff_eq!(loss, 0.0, epsilon = 1e-10);
    }

    #[test]
    fn test_mse_loss_calculation() {
        let predictions = vec![1.0, 2.0, 3.0];
        let targets = vec![1.1, 2.1, 2.9];

        let loss = mse_loss(&predictions, &targets).unwrap();

        // MSE = ((0.1)^2 + (0.1)^2 + (0.1)^2) / 3 = 0.03 / 3 = 0.01
        assert_abs_diff_eq!(loss, 0.01, epsilon = 1e-10);
    }

    #[test]
    fn test_mse_loss_length_mismatch() {
        let predictions = vec![1.0, 2.0];
        let targets = vec![1.0, 2.0, 3.0];

        let result = mse_loss(&predictions, &targets);
        assert!(result.is_err());
    }

    #[test]
    fn test_mse_loss_empty_arrays() {
        let predictions: Vec<f64> = vec![];
        let targets: Vec<f64> = vec![];

        let result = mse_loss(&predictions, &targets);
        assert!(result.is_err());
    }

    #[test]
    fn test_mse_loss_single_value() {
        let predictions = vec![5.0];
        let targets = vec![3.0];

        let loss = mse_loss(&predictions, &targets).unwrap();

        // MSE = (5-3)^2 = 4
        assert_abs_diff_eq!(loss, 4.0, epsilon = 1e-10);
    }

    #[test]
    fn test_mse_loss_large_errors() {
        let predictions = vec![10.0, 20.0];
        let targets = vec![0.0, 0.0];

        let loss = mse_loss(&predictions, &targets).unwrap();

        // MSE = (100 + 400) / 2 = 250
        assert_abs_diff_eq!(loss, 250.0, epsilon = 1e-10);
    }

    #[test]
    fn test_adamw_multiple_layers() {
        let weight_shapes = vec![
            vec![vec![0.0; 2]; 5],
            vec![vec![0.0; 5]; 3],
            vec![vec![0.0; 3]; 1],
        ];
        let bias_shapes = vec![vec![0.0; 5], vec![0.0; 3], vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.01, &weight_shapes, &bias_shapes).unwrap();

        let weight_grads = vec![
            vec![vec![0.5; 2]; 5],
            vec![vec![0.3; 5]; 3],
            vec![vec![0.1; 3]; 1],
        ];
        let bias_grads = vec![vec![0.5; 5], vec![0.3; 3], vec![0.1; 1]];

        let (weight_updates, bias_updates) = optimizer.step(&weight_grads, &bias_grads);

        assert_eq!(weight_updates.len(), 3);
        assert_eq!(bias_updates.len(), 3);

        // Check shapes match
        for (i, updates) in weight_updates.iter().enumerate() {
            assert_eq!(updates.len(), weight_shapes[i].len());
            for (j, row) in updates.iter().enumerate() {
                assert_eq!(row.len(), weight_shapes[i][j].len());
            }
        }
    }

    #[test]
    fn test_adamw_momentum_accumulation() {
        let weight_shapes = vec![vec![vec![0.0; 1]; 1]];
        let bias_shapes = vec![vec![0.0; 1]];

        let mut optimizer =
            AdamW::new(0.001, 0.9, 0.999, 0.0, &weight_shapes, &bias_shapes).unwrap();

        let weight_grads = vec![vec![vec![1.0]; 1]];
        let bias_grads = vec![vec![1.0]];

        // Take multiple steps with same gradient
        let (updates1, _) = optimizer.step(&weight_grads, &bias_grads);
        let (updates2, _) = optimizer.step(&weight_grads, &bias_grads);
        let (updates3, _) = optimizer.step(&weight_grads, &bias_grads);

        // Update magnitudes should stabilize (not grow indefinitely)
        let mag1 = updates1[0][0][0].abs();
        let mag2 = updates2[0][0][0].abs();
        let mag3 = updates3[0][0][0].abs();

        assert!(mag2 < mag1); // Second step smaller due to bias correction
        assert!(mag3 < mag2 || (mag3 - mag2).abs() < 0.0001); // Should stabilize
    }
}