my_ai_agent_celue_/src/indicators/volatility.rs

//! Volatility indicators: ATR, Bollinger Bands.

use crate::core::error::RaptorError;
use crate::core::Result;

/// Average True Range (ATR).
///
/// # Arguments
/// * `high` - High prices
/// * `low` - Low prices
/// * `close` - Close prices
/// * `period` - Lookback period (default: 14)
///
/// # Returns
/// Vector of ATR values (NaN for warmup period)
pub fn atr(high: &[f64], low: &[f64], close: &[f64], period: usize) -> Result<Vec<f64>> {
    let n = close.len();
    if n != high.len() || n != low.len() {
        return Err(RaptorError::length_mismatch(n, high.len()));
    }
    if period == 0 {
        return Err(RaptorError::invalid_parameter("ATR period must be > 0"));
    }
    Ok(ferro_ta_core::volatility::atr(high, low, close, period))
}

/// True Range calculation (single bar).
#[inline]
pub fn true_range(high: f64, low: f64, prev_close: f64) -> f64 {
    let hl = high - low;
    let hc = (high - prev_close).abs();
    let lc = (low - prev_close).abs();
    hl.max(hc).max(lc)
}

/// Bollinger Bands result.
#[derive(Debug, Clone)]
pub struct BollingerBandsResult {
    /// Middle band (SMA).
    pub middle: Vec<f64>,
    /// Upper band (SMA + std_dev * multiplier).
    pub upper: Vec<f64>,
    /// Lower band (SMA - std_dev * multiplier).
    pub lower: Vec<f64>,
    /// Bandwidth: (upper - lower) / middle.
    pub bandwidth: Vec<f64>,
    /// %B: (price - lower) / (upper - lower).
    pub percent_b: Vec<f64>,
}

/// Bollinger Bands.
///
/// # Arguments
/// * `data` - Price data (typically close prices)
/// * `period` - Lookback period (default: 20)
/// * `std_dev` - Standard deviation multiplier (default: 2.0)
///
/// # Returns
/// BollingerBandsResult with middle, upper, lower bands, bandwidth, and %B
pub fn bollinger_bands(data: &[f64], period: usize, std_dev: f64) -> Result<BollingerBandsResult> {
    if period == 0 {
        return Err(RaptorError::invalid_parameter("Bollinger Bands period must be > 0"));
    }
    if std_dev <= 0.0 {
        return Err(RaptorError::invalid_parameter("Bollinger Bands std_dev must be > 0"));
    }

    let n = data.len();
    let (upper, middle, lower) = ferro_ta_core::overlap::bbands(data, period, std_dev, std_dev);

    let mut bandwidth = vec![f64::NAN; n];
    let mut percent_b = vec![f64::NAN; n];

    for i in 0..n {
        if !middle[i].is_nan() && middle[i].abs() > f64::EPSILON {
            bandwidth[i] = (upper[i] - lower[i]) / middle[i].abs();
        }
        let band_width = upper[i] - lower[i];
        if band_width > f64::EPSILON {
            percent_b[i] = (data[i] - lower[i]) / band_width;
        }
    }

    Ok(BollingerBandsResult { middle, upper, lower, bandwidth, percent_b })
}

/// Keltner Channels (ATR-based bands).
///
/// # Arguments
/// * `high` - High prices
/// * `low` - Low prices
/// * `close` - Close prices
/// * `ema_period` - EMA period for middle band
/// * `atr_period` - ATR period
/// * `multiplier` - ATR multiplier
///
/// # Returns
/// Tuple of (middle, upper, lower) bands
pub fn keltner_channels(
    high: &[f64],
    low: &[f64],
    close: &[f64],
    ema_period: usize,
    atr_period: usize,
    multiplier: f64,
) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>)> {
    let n = close.len();
    if n != high.len() || n != low.len() {
        return Err(RaptorError::length_mismatch(n, high.len()));
    }

    // Calculate EMA for middle band
    let middle = super::trend::ema(close, ema_period)?;

    // Calculate ATR
    let atr_values = atr(high, low, close, atr_period)?;

    // Calculate bands
    let mut upper = vec![f64::NAN; n];
    let mut lower = vec![f64::NAN; n];

    for i in 0..n {
        if !middle[i].is_nan() && !atr_values[i].is_nan() {
            upper[i] = middle[i] + multiplier * atr_values[i];
            lower[i] = middle[i] - multiplier * atr_values[i];
        }
    }

    Ok((middle, upper, lower))
}

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

    #[test]
    fn test_atr() {
        let high = vec![50.0, 51.0, 52.0, 51.5, 50.5, 51.0, 52.0, 53.0, 52.5, 51.5];
        let low = vec![48.0, 49.0, 50.0, 49.5, 48.5, 49.0, 50.0, 51.0, 50.5, 49.5];
        let close = vec![49.0, 50.0, 51.0, 50.0, 49.0, 50.0, 51.0, 52.0, 51.0, 50.0];

        let result = atr(&high, &low, &close, 5).unwrap();

        // ATR should be valid from index 4
        assert!(result[3].is_nan());
        assert!(!result[4].is_nan());
        assert!(result[4] > 0.0);
    }

    #[test]
    fn test_bollinger_bands() {
        let data: Vec<f64> = (1..=30).map(|x| x as f64 + (x as f64 * 0.1).sin()).collect();

        let result = bollinger_bands(&data, 20, 2.0).unwrap();

        // Bands should be valid from index 19
        assert!(result.middle[18].is_nan());
        assert!(!result.middle[19].is_nan());

        // Upper > Middle > Lower
        assert!(result.upper[19] > result.middle[19]);
        assert!(result.middle[19] > result.lower[19]);

        // %B should be between 0 and 1 for data within bands
        assert!(result.percent_b[19] >= -0.5 && result.percent_b[19] <= 1.5);
    }

    #[test]
    fn test_true_range() {
        // Simple case
        assert!((true_range(52.0, 48.0, 50.0) - 4.0).abs() < 1e-10);

        // Gap up case
        assert!((true_range(55.0, 53.0, 50.0) - 5.0).abs() < 1e-10);

        // Gap down case
        assert!((true_range(48.0, 45.0, 50.0) - 5.0).abs() < 1e-10);
    }
}