//! Safe Mathematical Operations for ML
//!
//! This module provides mathematically safe operations that handle
//! NaN, Infinity, and edge cases gracefully for all ML computations.

use super::{MLSafetyConfig, MLSafetyError, SafetyResult};

/// Safe mathematical operations with comprehensive error handling
#[derive(Debug, Clone)]
pub struct SafeMathOps {
    config: MLSafetyConfig,
}

impl SafeMathOps {
    /// Create new safe math operations
    pub fn new(config: &MLSafetyConfig) -> Self {
        Self {
            config: config.clone(),
        }
    }

    /// Safely divide two numbers with fallback and validation
    pub fn safe_divide(&self, numerator: f64, denominator: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !numerator.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Division numerator: {}", numerator),
                });
            }
            if !denominator.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Division denominator: {}", denominator),
                });
            }
        }

        if denominator.abs() < f64::EPSILON {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Division by zero: {} / {}", numerator, denominator),
            });
        }

        let result = numerator / denominator;

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!(
                    "Division result: {} / {} = {}",
                    numerator, denominator, result
                ),
            });
        }

        Ok(result)
    }

    /// Safely compute square root with validation
    pub fn safe_sqrt(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Square root input: {}", value),
            });
        }

        if value < 0.0 {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Square root of negative number: {}", value),
            });
        }

        let result = value.sqrt();

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Square root result: sqrt({}) = {}", value, result),
            });
        }

        Ok(result)
    }

    /// Safely compute natural logarithm with validation
    pub fn safe_ln(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Natural log input: {}", value),
            });
        }

        if value <= 0.0 {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Natural log of non-positive number: {}", value),
            });
        }

        let result = value.ln();

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Natural log result: ln({}) = {}", value, result),
            });
        }

        Ok(result)
    }

    /// Safely compute exponential with overflow protection
    pub fn safe_exp(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Exponential input: {}", value),
            });
        }

        // Prevent overflow
        const MAX_EXP: f64 = 700.0;
        if value > MAX_EXP {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Exponential input too large (would overflow): {}", value),
            });
        }

        let result = if value < -MAX_EXP {
            0.0 // Underflow to zero
        } else {
            value.exp()
        };

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Exponential result: exp({}) = {}", value, result),
            });
        }

        Ok(result)
    }

    /// Safely compute power with overflow protection
    pub fn safe_pow(&self, base: f64, exponent: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !base.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Power base: {}", base),
                });
            }
            if !exponent.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Power exponent: {}", exponent),
                });
            }
        }

        // Check for potentially problematic cases
        if base == 0.0 && exponent <= 0.0 {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Zero to non-positive power: 0^{}", exponent),
            });
        }

        if base < 0.0 && exponent.fract() != 0.0 {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Negative base to fractional power: {}^{}", base, exponent),
            });
        }

        let result = base.powf(exponent);

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Power result: {}^{} = {}", base, exponent, result),
            });
        }

        Ok(result)
    }

    /// Safely normalize a vector with validation
    pub fn safe_normalize(&self, values: &[f64]) -> SafetyResult<Vec<f64>> {
        if values.is_empty() {
            return Err(MLSafetyError::MathSafety {
                reason: "Cannot normalize empty vector".to_owned(),
            });
        }

        // Check for potential overflow in length conversion
        if values.len() > (f64::MAX as usize) {
            return Err(MLSafetyError::MathSafety {
                reason: format!(
                    "Vector too large for safe processing: {} elements",
                    values.len()
                ),
            });
        }

        // Check for NaN/Infinity in input
        if self.config.nan_infinity_checks {
            for (i, &value) in values.into_iter().enumerate() {
                if !value.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("Normalize input at index {}: {}", i, value),
                    });
                }
            }
        }

        // Calculate sum with overflow protection
        let sum = self.safe_sum(values)?;

        if sum.abs() < f64::EPSILON {
            return Err(MLSafetyError::MathSafety {
                reason: "Cannot normalize vector with zero sum".to_owned(),
            });
        }

        // Normalize
        let normalized: SafetyResult<Vec<f64>> = values
            .iter()
            .map(|&value| self.safe_divide(value, sum))
            .collect();

        normalized
    }

    /// Safely clamp value between bounds
    pub fn safe_clamp(&self, value: f64, min: f64, max: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !value.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Clamp value: {}", value),
                });
            }
            if !min.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Clamp min: {}", min),
                });
            }
            if !max.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Clamp max: {}", max),
                });
            }
        }

        if min > max {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Invalid clamp range: min {} > max {}", min, max),
            });
        }

        Ok(value.max(min).min(max))
    }

    /// Safely compute softmax with numerical stability
    pub fn safe_softmax(&self, values: &[f64]) -> SafetyResult<Vec<f64>> {
        if values.is_empty() {
            return Err(MLSafetyError::MathSafety {
                reason: "Cannot compute softmax of empty vector".to_owned(),
            });
        }

        // Check for NaN/Infinity in input
        if self.config.nan_infinity_checks {
            for (i, &value) in values.into_iter().enumerate() {
                if !value.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("Softmax input at index {}: {}", i, value),
                    });
                }
            }
        }

        // Find maximum for numerical stability
        let max_val = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));

        if !max_val.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Softmax max value: {}", max_val),
            });
        }

        // Compute shifted exponentials
        let shifted_exp: SafetyResult<Vec<f64>> =
            values.iter().map(|&x| self.safe_exp(x - max_val)).collect();
        let shifted_exp = shifted_exp?;

        // Compute sum
        let sum = shifted_exp.iter().sum::<f64>();

        if sum <= f64::EPSILON {
            return Err(MLSafetyError::MathSafety {
                reason: "Softmax sum too small (numerical instability)".to_owned(),
            });
        }

        // Normalize
        let result: SafetyResult<Vec<f64>> = shifted_exp
            .iter()
            .map(|&x| self.safe_divide(x, sum))
            .collect();

        result
    }

    /// Safely compute sigmoid activation
    pub fn safe_sigmoid(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Sigmoid input: {}", value),
            });
        }

        // Use numerically stable sigmoid computation
        let result = if value >= 0.0 {
            let exp_neg = self.safe_exp(-value)?;
            self.safe_divide(1.0, 1.0 + exp_neg)?
        } else {
            let exp_pos = self.safe_exp(value)?;
            self.safe_divide(exp_pos, 1.0 + exp_pos)?
        };

        Ok(result)
    }

    /// Safely compute tanh activation
    pub fn safe_tanh(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Tanh input: {}", value),
            });
        }

        // Use numerically stable tanh computation
        if value.abs() > 20.0 {
            // Avoid overflow for large values
            Ok(if value > 0.0 { 1.0 } else { -1.0 })
        } else {
            let result = value.tanh();
            if self.config.nan_infinity_checks && !result.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Tanh result: tanh({}) = {}", value, result),
                });
            }
            Ok(result)
        }
    }

    /// Safely compute ReLU activation
    pub fn safe_relu(&self, value: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks && !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("ReLU input: {}", value),
            });
        }

        Ok(value.max(0.0))
    }

    /// Safely cast usize to f64 with overflow protection
    pub fn safe_cast_usize_to_f64(&self, value: usize) -> SafetyResult<f64> {
        if value > (f64::MAX as usize) {
            return Err(MLSafetyError::MathSafety {
                reason: format!("usize value {} too large for f64 conversion", value),
            });
        }
        Ok(value as f64)
    }

    /// Safely sum a slice of f64 values with overflow protection
    pub fn safe_sum(&self, values: &[f64]) -> SafetyResult<f64> {
        if values.is_empty() {
            return Ok(0.0);
        }

        // Check for NaN/Infinity in input if configured
        if self.config.nan_infinity_checks {
            for (i, &value) in values.into_iter().enumerate() {
                if !value.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("Sum input at index {}: {}", i, value),
                    });
                }
            }
        }

        // Use Kahan summation algorithm for better numerical stability
        let mut sum = 0.0;
        let mut compensation = 0.0;

        for &value in values {
            let compensated_value = value - compensation;
            let temp_sum = sum + compensated_value;
            compensation = (temp_sum - sum) - compensated_value;
            sum = temp_sum;

            // Check for overflow during summation
            if !sum.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Sum overflow detected: {}", sum),
                });
            }
        }

        Ok(sum)
    }

    /// Safely add two f64 values with overflow protection
    pub fn safe_add(&self, a: f64, b: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !a.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Addition operand a: {}", a),
                });
            }
            if !b.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Addition operand b: {}", b),
                });
            }
        }

        let result = a + b;

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Addition result: {} + {} = {}", a, b, result),
            });
        }

        Ok(result)
    }

    /// Safely multiply two f64 values with overflow protection
    pub fn safe_multiply(&self, a: f64, b: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !a.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Multiplication operand a: {}", a),
                });
            }
            if !b.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Multiplication operand b: {}", b),
                });
            }
        }

        let result = a * b;

        if self.config.nan_infinity_checks && !result.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Multiplication result: {} * {} = {}", a, b, result),
            });
        }

        Ok(result)
    }

    /// Safely compute leaky ReLU activation
    pub fn safe_leaky_relu(&self, value: f64, alpha: f64) -> SafetyResult<f64> {
        if self.config.nan_infinity_checks {
            if !value.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Leaky ReLU input: {}", value),
                });
            }
            if !alpha.is_finite() {
                return Err(MLSafetyError::InvalidFloat {
                    operation: format!("Leaky ReLU alpha: {}", alpha),
                });
            }
        }

        if alpha < 0.0 || alpha >= 1.0 {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Invalid leaky ReLU alpha: {} (must be in [0, 1))", alpha),
            });
        }

        Ok(if value >= 0.0 { value } else { alpha * value })
    }

    /// Safely compute mean squared error
    pub fn safe_mse(&self, predicted: &[f64], actual: &[f64]) -> SafetyResult<f64> {
        if predicted.len() != actual.len() {
            return Err(MLSafetyError::MathSafety {
                reason: format!(
                    "Length mismatch: predicted {} vs actual {}",
                    predicted.len(),
                    actual.len()
                ),
            });
        }

        if predicted.is_empty() {
            return Err(MLSafetyError::MathSafety {
                reason: "Cannot compute MSE of empty arrays".to_owned(),
            });
        }

        // Check for potential overflow in length conversion
        if predicted.len() > (f64::MAX as usize) {
            return Err(MLSafetyError::MathSafety {
                reason: format!(
                    "Arrays too large for safe processing: {} elements",
                    predicted.len()
                ),
            });
        }

        // Check for NaN/Infinity
        if self.config.nan_infinity_checks {
            for (i, (&pred, &act)) in predicted.into_iter().zip(actual.into_iter()).enumerate() {
                if !pred.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("MSE predicted at index {}: {}", i, pred),
                    });
                }
                if !act.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("MSE actual at index {}: {}", i, act),
                    });
                }
            }
        }

        // Compute squared errors safely
        let squared_errors: SafetyResult<Vec<f64>> = predicted
            .iter()
            .zip(actual.iter())
            .map(|(&pred, &act)| {
                let diff = pred - act;
                self.safe_pow(diff, 2.0)
            })
            .collect();

        let squared_errors = squared_errors?;
        let sum = self.safe_sum(&squared_errors)?;

        self.safe_divide(sum, self.safe_cast_usize_to_f64(predicted.len())?)
    }

    /// Safely compute correlation coefficient
    pub fn safe_correlation(&self, x: &[f64], y: &[f64]) -> SafetyResult<f64> {
        if x.len() != y.len() {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Length mismatch: x {} vs y {}", x.len(), y.len()),
            });
        }

        if x.len() < 2 {
            return Err(MLSafetyError::MathSafety {
                reason: "Need at least 2 points for correlation".to_owned(),
            });
        }

        // Check for potential overflow in length conversion
        if x.len() > (f64::MAX as usize) {
            return Err(MLSafetyError::MathSafety {
                reason: format!("Arrays too large for safe processing: {} elements", x.len()),
            });
        }

        // Check for NaN/Infinity
        if self.config.nan_infinity_checks {
            for (i, (&xi, &yi)) in x.into_iter().zip(y.into_iter()).enumerate() {
                if !xi.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("Correlation x at index {}: {}", i, xi),
                    });
                }
                if !yi.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!("Correlation y at index {}: {}", i, yi),
                    });
                }
            }
        }

        let n = self.safe_cast_usize_to_f64(x.len())?;
        let sum_x = self.safe_sum(x)?;
        let sum_y = self.safe_sum(y)?;
        let mean_x = self.safe_divide(sum_x, n)?;
        let mean_y = self.safe_divide(sum_y, n)?;

        let mut sum_xy = 0.0;
        let mut sum_x2 = 0.0;
        let mut sum_y2 = 0.0;

        for (&xi, &yi) in x.into_iter().zip(y.into_iter()) {
            let dx = xi - mean_x;
            let dy = yi - mean_y;

            // Check for overflow in intermediate calculations
            let xy_term = self.safe_multiply(dx, dy)?;
            let x2_term = self.safe_multiply(dx, dx)?;
            let y2_term = self.safe_multiply(dy, dy)?;

            sum_xy = self.safe_add(sum_xy, xy_term)?;
            sum_x2 = self.safe_add(sum_x2, x2_term)?;
            sum_y2 = self.safe_add(sum_y2, y2_term)?;
        }

        let sqrt_x2 = self.safe_sqrt(sum_x2)?;
        let sqrt_y2 = self.safe_sqrt(sum_y2)?;
        let denominator = self.safe_multiply(sqrt_x2, sqrt_y2)?;

        if denominator.abs() < f64::EPSILON {
            return Err(MLSafetyError::MathSafety {
                reason: "Zero variance in correlation computation".to_owned(),
            });
        }

        self.safe_divide(sum_xy, denominator)
    }
}

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

    fn create_test_ops() -> SafeMathOps {
        SafeMathOps::new(&MLSafetyConfig::default())
    }

    #[test]
    fn test_safe_divide() {
        let ops = create_test_ops();

        // Valid division
        let result = ops.safe_divide(10.0, 2.0);
        assert!(
            result.is_ok(),
            "Valid division should succeed: {:?}",
            result.err()
        );
        if let Ok(value) = result {
            assert!(
                (value - 5.0).abs() < f64::EPSILON,
                "Division result should be 5.0, got {}",
                value
            );
        }

        // Division by zero
        assert!(
            ops.safe_divide(10.0, 0.0).is_err(),
            "Division by zero should fail"
        );

        // NaN inputs
        assert!(
            ops.safe_divide(f64::NAN, 2.0).is_err(),
            "Division with NaN numerator should fail"
        );
        assert!(
            ops.safe_divide(10.0, f64::NAN).is_err(),
            "Division with NaN denominator should fail"
        );
    }

    #[test]
    fn test_safe_sqrt() {
        let ops = create_test_ops();

        // Valid sqrt
        let result = ops.safe_sqrt(4.0);
        assert!(
            result.is_ok(),
            "Valid square root should succeed: {:?}",
            result.err()
        );
        if let Ok(value) = result {
            assert!(
                (value - 2.0).abs() < f64::EPSILON,
                "Square root of 4.0 should be 2.0, got {}",
                value
            );
        }

        // Negative input
        assert!(
            ops.safe_sqrt(-1.0).is_err(),
            "Square root of negative number should fail"
        );

        // NaN input
        assert!(
            ops.safe_sqrt(f64::NAN).is_err(),
            "Square root of NaN should fail"
        );
    }

    #[test]
    fn test_safe_softmax() {
        let ops = create_test_ops();

        // Valid softmax
        let input = vec![1.0, 2.0, 3.0];
        let result = ops.safe_softmax(&input);
        assert!(
            result.is_ok(),
            "Valid softmax should succeed: {:?}",
            result.err()
        );
        if let Ok(softmax_result) = result {
            let sum = softmax_result.iter().sum::<f64>();
            assert!(
                (sum - 1.0).abs() < 1e-10,
                "Softmax should sum to 1.0, got {}",
                sum
            );
            assert!(
                softmax_result.len() == input.len(),
                "Softmax output length should match input length"
            );
            assert!(
                softmax_result.iter().all(|&x| x >= 0.0 && x <= 1.0),
                "All softmax values should be in [0,1]"
            );
        }

        // Empty input
        assert!(
            ops.safe_softmax(&[]).is_err(),
            "Softmax of empty vector should fail"
        );

        // NaN input
        assert!(
            ops.safe_softmax(&[1.0, f64::NAN, 3.0]).is_err(),
            "Softmax with NaN input should fail"
        );
    }

    #[test]
    fn test_safe_correlation() {
        let ops = create_test_ops();

        // Perfect correlation
        let x = vec![1.0, 2.0, 3.0, 4.0];
        let y = vec![2.0, 4.0, 6.0, 8.0];
        let result = ops.safe_correlation(&x, &y);
        assert!(
            result.is_ok(),
            "Valid correlation should succeed: {:?}",
            result.err()
        );
        if let Ok(corr) = result {
            assert!(
                (corr - 1.0).abs() < 1e-10,
                "Perfect positive correlation should be 1.0, got {}",
                corr
            );
            assert!(
                corr >= -1.0 && corr <= 1.0,
                "Correlation should be in [-1,1], got {}",
                corr
            );
        }

        // Length mismatch
        assert!(
            ops.safe_correlation(&[1.0, 2.0], &[1.0]).is_err(),
            "Correlation with mismatched lengths should fail"
        );

        // Too few points
        assert!(
            ops.safe_correlation(&[1.0], &[2.0]).is_err(),
            "Correlation with insufficient data should fail"
        );
    }
}
