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

use super::{Indicator, IncrementalIndicator, IndicatorResult, IndicatorError, PriceData};

/// Exponential Moving Average (EMA) Indicator
/// 
/// Calculates the exponentially weighted moving average
/// giving more weight to recent prices
pub struct EMA {
    period: usize,
    alpha: f64,
    value: Option<f64>,
    initialized: bool,
}

impl EMA {
    pub fn new(period: usize) -> Result<Self, IndicatorError> {
        if period == 0 {
            return Err(IndicatorError::InvalidParameter(
                "Period must be greater than 0".to_string()
            ));
        }
        
        // Calculate smoothing factor (alpha)
        // Common formula: 2 / (period + 1)
        let alpha = 2.0 / (period as f64 + 1.0);
        
        Ok(Self {
            period,
            alpha,
            value: None,
            initialized: false,
        })
    }
    
    /// Create EMA with custom smoothing factor
    pub fn with_alpha(alpha: f64) -> Result<Self, IndicatorError> {
        if alpha <= 0.0 || alpha > 1.0 {
            return Err(IndicatorError::InvalidParameter(
                "Alpha must be between 0 and 1".to_string()
            ));
        }
        
        // Calculate equivalent period for reference
        let period = ((2.0 / alpha) - 1.0) as usize;
        
        Ok(Self {
            period,
            alpha,
            value: None,
            initialized: false,
        })
    }
    
    /// Calculate EMA for a series of values
    pub fn calculate_series(values: &[f64], period: usize) -> Result<Vec<f64>, IndicatorError> {
        if period == 0 {
            return Err(IndicatorError::InvalidParameter(
                "Period must be greater than 0".to_string()
            ));
        }
        
        if values.is_empty() {
            return Ok(vec![]);
        }
        
        let alpha = 2.0 / (period as f64 + 1.0);
        let mut result = Vec::with_capacity(values.len());
        
        // Start with first value as initial EMA
        let mut ema = values[0];
        result.push(ema);
        
        // Calculate EMA for remaining values
        for i in 1..values.len() {
            ema = alpha * values[i] + (1.0 - alpha) * ema;
            result.push(ema);
        }
        
        Ok(result)
    }
    
    /// Alternative initialization using SMA of first N values
    pub fn calculate_series_sma_init(values: &[f64], period: usize) -> Result<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(),
            });
        }
        
        let alpha = 2.0 / (period as f64 + 1.0);
        let mut result = Vec::with_capacity(values.len() - period + 1);
        
        // Calculate initial SMA
        let initial_sma: f64 = values[0..period].iter().sum::<f64>() / period as f64;
        let mut ema = initial_sma;
        result.push(ema);
        
        // Calculate EMA for remaining values
        for i in period..values.len() {
            ema = alpha * values[i] + (1.0 - alpha) * ema;
            result.push(ema);
        }
        
        Ok(result)
    }
}

impl Indicator for EMA {
    fn calculate(&mut self, data: &PriceData) -> Result<IndicatorResult, IndicatorError> {
        let values = &data.close;
        
        if values.is_empty() {
            return Err(IndicatorError::InsufficientData {
                required: 1,
                actual: 0,
            });
        }
        
        // Reset and calculate from scratch
        self.reset();
        let ema_values = Self::calculate_series(values, self.period)?;
        
        // Update internal state with last value
        if let Some(&last) = ema_values.last() {
            self.value = Some(last);
            self.initialized = true;
        }
        
        Ok(IndicatorResult::Series(ema_values))
    }
    
    fn reset(&mut self) {
        self.value = None;
        self.initialized = false;
    }
    
    fn is_ready(&self) -> bool {
        self.initialized && self.value.is_some()
    }
}

impl IncrementalIndicator for EMA {
    fn update(&mut self, value: f64) -> Result<Option<f64>, IndicatorError> {
        match self.value {
            Some(prev_ema) => {
                // Update EMA: EMA = α × Price + (1 - α) × Previous EMA
                let new_ema = self.alpha * value + (1.0 - self.alpha) * prev_ema;
                self.value = Some(new_ema);
                self.initialized = true;
                Ok(Some(new_ema))
            }
            None => {
                // First value becomes the initial EMA
                self.value = Some(value);
                self.initialized = true;
                Ok(Some(value))
            }
        }
    }
    
    fn current(&self) -> Option<f64> {
        self.value
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_ema_calculation() {
        let values = vec![10.0, 11.0, 12.0, 13.0, 14.0, 15.0];
        let result = EMA::calculate_series(&values, 3).unwrap();
        
        assert_eq!(result.len(), 6);
        assert!((result[0] - 10.0).abs() < 1e-10); // First value
        
        // Verify EMA calculation
        let alpha = 2.0 / 4.0; // 0.5
        let expected_ema2 = alpha * 11.0 + (1.0 - alpha) * 10.0; // 10.5
        assert!((result[1] - expected_ema2).abs() < 1e-10);
    }
    
    #[test]
    fn test_incremental_ema() {
        let mut ema = EMA::new(3).unwrap();
        
        // First value
        assert_eq!(ema.update(10.0).unwrap(), Some(10.0));
        
        // Second value: EMA = 0.5 * 12 + 0.5 * 10 = 11
        assert_eq!(ema.update(12.0).unwrap(), Some(11.0));
        
        // Third value: EMA = 0.5 * 14 + 0.5 * 11 = 12.5
        assert_eq!(ema.update(14.0).unwrap(), Some(12.5));
    }
    
    #[test]
    fn test_ema_with_sma_init() {
        let values = vec![2.0, 4.0, 6.0, 8.0, 10.0, 12.0];
        let result = EMA::calculate_series_sma_init(&values, 3).unwrap();
        
        // Initial SMA = (2 + 4 + 6) / 3 = 4
        assert_eq!(result.len(), 4);
        assert!((result[0] - 4.0).abs() < 1e-10);
    }
}