//! Log-space arithmetic utilities for numerically stable computations.
//!
//! These functions perform arithmetic operations in log-space to avoid
//! underflow/overflow when dealing with very small or very large probabilities.

/// Compute log(exp(a) + exp(b)) in a numerically stable way.
///
/// Uses the log-sum-exp trick: log(exp(a) + exp(b)) = max(a,b) + log(exp(a-max) + exp(b-max))
/// This avoids overflow/underflow by keeping intermediate values bounded.
///
/// # Arguments
/// * `a` - First log-space value
/// * `b` - Second log-space value
///
/// # Returns
/// The log of the sum of exponentials
///
/// # Examples
/// ```
/// use fastbaps_core::logspace::log_sum_exp;
/// let result = log_sum_exp(-1.0, -2.0);
/// assert!((result - (-0.6867)).abs() < 1e-4);
/// ```
pub fn log_sum_exp(a: f64, b: f64) -> f64 {
    // Handle case where both are -inf
    if a.is_infinite() && a.is_sign_negative() && b.is_infinite() && b.is_sign_negative() {
        return f64::NEG_INFINITY;
    }

    let max_val = a.max(b);
    max_val + ((a - max_val).exp() + (b - max_val).exp()).ln()
}

/// Compute log(exp(a) - exp(b)) in a numerically stable way.
///
/// This is the log-space equivalent of subtraction. Note that this requires a >= b
/// for the result to be meaningful (you cannot take the log of a negative number).
///
/// # Arguments
/// * `a` - First log-space value (must be >= b in log-space)
/// * `b` - Second log-space value to subtract
///
/// # Returns
/// The log of the difference of exponentials
///
/// # Special cases
/// * If b is -infinity, returns a unchanged (subtracting zero in linear space)
///
/// # Examples
/// ```
/// use fastbaps_core::logspace::log_subtract;
/// let result = log_subtract(0.0, f64::NEG_INFINITY);
/// assert_eq!(result, 0.0);
/// ```
pub fn log_subtract(a: f64, b: f64) -> f64 {
    // If b is -inf, exp(b) = 0, so exp(a) - exp(b) = exp(a)
    if b.is_infinite() && b.is_sign_negative() {
        return a;
    }

    // Use log1p for numerical stability: log(exp(a) - exp(b)) = a + log(1 - exp(b-a))
    a + (-((b - a).exp())).ln_1p()
}

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

    #[test]
    fn test_log_sum_exp_basic() {
        // log(exp(-1) + exp(-2)) ≈ log(0.368 + 0.135) ≈ log(0.503) ≈ -0.687
        let result = log_sum_exp(-1.0, -2.0);
        assert!((result - (-0.6867)).abs() < 1e-4);
    }

    #[test]
    fn test_log_sum_exp_both_neg_inf() {
        let result = log_sum_exp(f64::NEG_INFINITY, f64::NEG_INFINITY);
        assert_eq!(result, f64::NEG_INFINITY);
    }

    #[test]
    fn test_log_sum_exp_one_neg_inf() {
        // log(exp(-1) + exp(-inf)) = log(exp(-1) + 0) = -1
        let result = log_sum_exp(-1.0, f64::NEG_INFINITY);
        assert!((result - (-1.0)).abs() < 1e-10);
    }

    #[test]
    fn test_log_sum_exp_symmetric() {
        // Should be commutative
        let r1 = log_sum_exp(-1.0, -2.0);
        let r2 = log_sum_exp(-2.0, -1.0);
        assert!((r1 - r2).abs() < 1e-10);
    }

    #[test]
    fn test_log_subtract_with_neg_inf() {
        // log(exp(0) - exp(-inf)) = log(1 - 0) = 0
        let result = log_subtract(0.0, f64::NEG_INFINITY);
        assert_eq!(result, 0.0);
    }

    #[test]
    fn test_log_subtract_basic() {
        // log(exp(0) - exp(-1)) = log(1 - 0.368) ≈ log(0.632) ≈ -0.459
        let result = log_subtract(0.0, -1.0);
        assert!((result - (-0.459)).abs() < 1e-3);
    }

    #[test]
    fn test_log_subtract_close_values() {
        // log(exp(-1) - exp(-1.1)) should be negative
        let result = log_subtract(-1.0, -1.1);
        assert!(result < -1.0);
    }
}
