stock-bot/apps/stock/engine/src/indicators/bollinger_bands.rs

use super::{Indicator, IncrementalIndicator, IndicatorResult, IndicatorError, PriceData};
use super::sma::SMA;
use super::common::RollingWindow;

/// Bollinger Bands Indicator
/// 
/// Middle Band = SMA(n)
/// Upper Band = SMA(n) + (k × σ)
/// Lower Band = SMA(n) - (k × σ)
/// where σ is the standard deviation and k is typically 2
pub struct BollingerBands {
    period: usize,
    std_dev_multiplier: f64,
    sma: SMA,
    window: RollingWindow<f64>,
}

impl BollingerBands {
    pub fn new(period: usize, std_dev_multiplier: f64) -> Result<Self, IndicatorError> {
        if period == 0 {
            return Err(IndicatorError::InvalidParameter(
                "Period must be greater than 0".to_string()
            ));
        }
        
        if std_dev_multiplier <= 0.0 {
            return Err(IndicatorError::InvalidParameter(
                "Standard deviation multiplier must be positive".to_string()
            ));
        }
        
        Ok(Self {
            period,
            std_dev_multiplier,
            sma: SMA::new(period)?,
            window: RollingWindow::new(period),
        })
    }
    
    /// Standard Bollinger Bands with 20 period and 2 standard deviations
    pub fn standard() -> Result<Self, IndicatorError> {
        Self::new(20, 2.0)
    }
    
    /// Calculate standard deviation
    fn calculate_std_dev(values: &[f64], mean: f64) -> f64 {
        if values.is_empty() {
            return 0.0;
        }
        
        let variance = values.iter()
            .map(|x| (*x - mean).powi(2))
            .sum::<f64>() / values.len() as f64;
        
        variance.sqrt()
    }
    
    /// Calculate Bollinger Bands for a series of values
    pub fn calculate_series(
        values: &[f64], 
        period: usize, 
        std_dev_multiplier: f64
    ) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>), IndicatorError> {
        if period == 0 {
            return Err(IndicatorError::InvalidParameter(
                "Period must be greater than 0".to_string()
            ));
        }
        
        if values.len() < period {
            return Err(IndicatorError::InsufficientData {
                required: period,
                actual: values.len(),
            });
        }
        
        // Calculate SMA (middle band)
        let middle_band = SMA::calculate_series(values, period)?;
        let mut upper_band = Vec::with_capacity(middle_band.len());
        let mut lower_band = Vec::with_capacity(middle_band.len());
        
        // Calculate bands
        for i in 0..middle_band.len() {
            // Get the window of values for this position
            let start_idx = i;
            let end_idx = i + period;
            let window = &values[start_idx..end_idx];
            
            // Calculate standard deviation
            let std_dev = Self::calculate_std_dev(window, middle_band[i]);
            let band_width = std_dev * std_dev_multiplier;
            
            upper_band.push(middle_band[i] + band_width);
            lower_band.push(middle_band[i] - band_width);
        }
        
        Ok((middle_band, upper_band, lower_band))
    }
    
    /// Calculate the current bandwidth (distance between upper and lower bands)
    pub fn bandwidth(&self) -> Option<f64> {
        if let Some(values) = self.get_current_bands() {
            Some(values.upper - values.lower)
        } else {
            None
        }
    }
    
    /// Calculate %B (percent b) - position of price relative to bands
    /// %B = (Price - Lower Band) / (Upper Band - Lower Band)
    pub fn percent_b(&self, price: f64) -> Option<f64> {
        if let Some(values) = self.get_current_bands() {
            let width = values.upper - values.lower;
            if width > 0.0 {
                Some((price - values.lower) / width)
            } else {
                None
            }
        } else {
            None
        }
    }
    
    fn get_current_bands(&self) -> Option<BollingerBandsValues> {
        if let Some(middle) = self.sma.current() {
            if self.window.is_full() {
                let values = self.window.as_slice();
                let std_dev = Self::calculate_std_dev(&values, middle);
                let band_width = std_dev * self.std_dev_multiplier;
                
                Some(BollingerBandsValues {
                    middle,
                    upper: middle + band_width,
                    lower: middle - band_width,
                })
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl Indicator for BollingerBands {
    fn calculate(&mut self, data: &PriceData) -> Result<IndicatorResult, IndicatorError> {
        let values = &data.close;
        
        let (middle, upper, lower) = Self::calculate_series(
            values, 
            self.period, 
            self.std_dev_multiplier
        )?;
        
        Ok(IndicatorResult::BollingerBands {
            middle,
            upper,
            lower,
        })
    }
    
    fn reset(&mut self) {
        self.sma.reset();
        self.window.clear();
    }
    
    fn is_ready(&self) -> bool {
        self.sma.is_ready() && self.window.is_full()
    }
}

impl IncrementalIndicator for BollingerBands {
    fn update(&mut self, value: f64) -> Result<Option<f64>, IndicatorError> {
        // Update window
        self.window.push(value);
        
        // Update SMA
        let _sma_result = self.sma.update(value)?;
        
        // Return bandwidth if ready
        if self.is_ready() {
            Ok(self.bandwidth())
        } else {
            Ok(None)
        }
    }
    
    fn current(&self) -> Option<f64> {
        self.bandwidth()
    }
}

/// Structure to hold all Bollinger Bands values
pub struct BollingerBandsValues {
    pub middle: f64,
    pub upper: f64,
    pub lower: f64,
}

impl BollingerBands {
    /// Get all current band values
    pub fn current_values(&self) -> Option<BollingerBandsValues> {
        self.get_current_bands()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_bollinger_bands_calculation() {
        let values = vec![
            20.0, 21.0, 22.0, 23.0, 24.0,
            25.0, 24.0, 23.0, 22.0, 21.0,
            20.0, 21.0, 22.0, 23.0, 24.0,
            25.0, 24.0, 23.0, 22.0, 21.0
        ];
        
        let (middle, upper, lower) = BollingerBands::calculate_series(&values, 5, 2.0).unwrap();
        
        assert_eq!(middle.len(), 16);
        assert_eq!(upper.len(), 16);
        assert_eq!(lower.len(), 16);
        
        // Upper band should always be above middle
        for i in 0..middle.len() {
            assert!(upper[i] > middle[i]);
            assert!(lower[i] < middle[i]);
        }
    }
    
    #[test]
    fn test_percent_b() {
        let mut bb = BollingerBands::standard().unwrap();
        
        // Create some test data
        for i in 0..25 {
            let _ = bb.update(20.0 + (i as f64 % 5.0));
        }
        
        // Price at upper band should give %B ≈ 1.0
        if let Some(bands) = bb.current_values() {
            let percent_b = bb.percent_b(bands.upper).unwrap();
            assert!((percent_b - 1.0).abs() < 1e-10);
            
            // Price at lower band should give %B ≈ 0.0
            let percent_b = bb.percent_b(bands.lower).unwrap();
            assert!(percent_b.abs() < 1e-10);
            
            // Price at middle band should give %B ≈ 0.5
            let percent_b = bb.percent_b(bands.middle).unwrap();
            assert!((percent_b - 0.5).abs() < 0.1);
        }
    }
}