nyxs_owl 0.4.0

//! Moving average calculation implementations
//!
//! Contains implementations of various moving average types:
//! - Simple Moving Average (SMA)
//! - Exponential Moving Average (EMA)
//! - Volume-Weighted Moving Average (VWMA)

use super::{MathError, Result};
use std::collections::VecDeque;

/// Simple Moving Average (SMA) implementation with optimizations
#[derive(Debug, Clone)]
pub struct SimpleMovingAverage {
    period: usize,
    values: VecDeque<f64>,
    sum: f64,
    count: usize, // Track count separately for better performance
}

impl SimpleMovingAverage {
    /// Create a new Simple Moving Average with the specified period
    pub fn new(period: usize) -> Result<Self> {
        if period == 0 {
            return Err(MathError::InvalidInput(
                "Period must be greater than zero".to_string(),
            ));
        }

        Ok(Self {
            period,
            values: VecDeque::with_capacity(period), // Pre-allocate exact capacity
            sum: 0.0,
            count: 0,
        })
    }

    /// Update the SMA with a new value - optimized version
    #[inline]
    pub fn update(&mut self, value: f64) -> Result<()> {
        // Add new value
        self.values.push_back(value);
        self.sum += value;
        self.count += 1;

        // Remove oldest value if we have more than period values
        if self.count > self.period {
            if let Some(old_value) = self.values.pop_front() {
                self.sum -= old_value;
                self.count -= 1;
            }
        }

        Ok(())
    }

    /// Get the current SMA value - optimized version
    #[inline]
    pub fn value(&self) -> Result<f64> {
        if self.count < self.period {
            return Err(MathError::InsufficientData(format!(
                "Not enough data for SMA calculation. Need {} values, have {}.",
                self.period, self.count
            )));
        }

        // Use direct division instead of len() call for better performance
        Ok(self.sum / self.period as f64)
    }

    /// Get the current period
    #[inline]
    pub fn period(&self) -> usize {
        self.period
    }

    /// Reset the SMA, clearing all values
    pub fn reset(&mut self) {
        self.values.clear();
        self.sum = 0.0;
        self.count = 0;
    }

    /// Get current value count
    #[inline]
    pub fn count(&self) -> usize {
        self.count
    }

    /// Check if SMA is ready (has enough values)
    #[inline]
    pub fn is_ready(&self) -> bool {
        self.count >= self.period
    }
}

/// Exponential Moving Average (EMA) implementation with optimizations
#[derive(Debug, Clone)]
pub struct ExponentialMovingAverage {
    period: usize,
    multiplier: f64,
    current_ema: Option<f64>,
    values_seen: usize,
    alpha: f64, // Store alpha separately for better performance
}

impl ExponentialMovingAverage {
    /// Create a new Exponential Moving Average with the specified period
    pub fn new(period: usize) -> Result<Self> {
        if period == 0 {
            return Err(MathError::InvalidInput(
                "Period must be greater than zero".to_string(),
            ));
        }

        let multiplier = 2.0 / (period as f64 + 1.0);
        let alpha = 2.0 / (period as f64 + 1.0);

        Ok(Self {
            period,
            multiplier,
            current_ema: None,
            values_seen: 0,
            alpha,
        })
    }

    /// Update the EMA with a new value - optimized version
    #[inline]
    pub fn update(&mut self, value: f64) -> Result<()> {
        self.values_seen += 1;

        match self.current_ema {
            None => {
                // First value - initialize EMA
                self.current_ema = Some(value);
            }
            Some(current) => {
                // Optimized EMA calculation: EMA = α * value + (1-α) * EMA_prev
                // This is mathematically equivalent but faster than the original formula
                let new_ema = self.alpha * value + (1.0 - self.alpha) * current;
                self.current_ema = Some(new_ema);
            }
        }

        Ok(())
    }

    /// Get the current EMA value - optimized version
    #[inline]
    pub fn value(&self) -> Result<f64> {
        self.current_ema.ok_or_else(|| {
            MathError::InsufficientData(
                "Not enough data for EMA calculation. Need at least 1 value.".to_string(),
            )
        })
    }

    /// Get the current period
    #[inline]
    pub fn period(&self) -> usize {
        self.period
    }

    /// Reset the EMA, clearing all values
    pub fn reset(&mut self) {
        self.current_ema = None;
        self.values_seen = 0;
    }

    /// Get the number of values seen
    #[inline]
    pub fn values_seen(&self) -> usize {
        self.values_seen
    }

    /// Check if EMA is ready (has at least one value)
    #[inline]
    pub fn is_ready(&self) -> bool {
        self.current_ema.is_some()
    }

    /// Get the smoothing factor (alpha)
    #[inline]
    pub fn alpha(&self) -> f64 {
        self.alpha
    }
}

/// Volume-Weighted Moving Average (VWMA) implementation
#[derive(Debug, Clone)]
pub struct VolumeWeightedMovingAverage {
    period: usize,
    price_volume_products: VecDeque<f64>,
    volumes: VecDeque<f64>,
    sum_price_volume: f64,
    sum_volume: f64,
}

impl VolumeWeightedMovingAverage {
    /// Create a new Volume-Weighted Moving Average with the specified period
    pub fn new(period: usize) -> Result<Self> {
        if period == 0 {
            return Err(MathError::InvalidInput(
                "Period must be greater than zero".to_string(),
            ));
        }

        Ok(Self {
            period,
            price_volume_products: VecDeque::with_capacity(period),
            volumes: VecDeque::with_capacity(period),
            sum_price_volume: 0.0,
            sum_volume: 0.0,
        })
    }

    /// Update the VWMA with a new price and volume
    pub fn update(&mut self, price: f64, volume: f64) -> Result<()> {
        if volume < 0.0 {
            return Err(MathError::InvalidInput(
                "Volume cannot be negative".to_string(),
            ));
        }

        let price_volume = price * volume;

        // Add new values
        self.price_volume_products.push_back(price_volume);
        self.volumes.push_back(volume);

        self.sum_price_volume += price_volume;
        self.sum_volume += volume;

        // Remove oldest values if we have more than period values
        if self.price_volume_products.len() > self.period {
            if let Some(old_price_volume) = self.price_volume_products.pop_front() {
                self.sum_price_volume -= old_price_volume;
            }

            if let Some(old_volume) = self.volumes.pop_front() {
                self.sum_volume -= old_volume;
            }
        }

        Ok(())
    }

    /// Get the current VWMA value
    pub fn value(&self) -> Result<f64> {
        if self.price_volume_products.len() < self.period {
            return Err(MathError::InsufficientData(format!(
                "Not enough data for VWMA calculation. Need {} values, have {}.",
                self.period,
                self.price_volume_products.len()
            )));
        }

        if self.sum_volume == 0.0 {
            return Err(MathError::CalculationError(
                "Volume sum is zero, cannot calculate VWMA".to_string(),
            ));
        }

        Ok(self.sum_price_volume / self.sum_volume)
    }

    /// Get the current period
    pub fn period(&self) -> usize {
        self.period
    }

    /// Reset the VWMA, clearing all values
    pub fn reset(&mut self) {
        self.price_volume_products.clear();
        self.volumes.clear();
        self.sum_price_volume = 0.0;
        self.sum_volume = 0.0;
    }
}

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

    #[test]
    fn test_sma_calculation() {
        let mut sma = SimpleMovingAverage::new(3).unwrap();

        // Not enough data yet
        assert!(sma.value().is_err());

        sma.update(2.0).unwrap();
        sma.update(4.0).unwrap();

        // Still not enough data
        assert!(sma.value().is_err());

        sma.update(6.0).unwrap();

        // Now we have enough data
        assert_eq!(sma.value().unwrap(), 4.0); // (2 + 4 + 6) / 3 = 4

        // Add another value
        sma.update(8.0).unwrap();

        // The window slides, dropping the oldest value
        assert_eq!(sma.value().unwrap(), 6.0); // (4 + 6 + 8) / 3 = 6
    }

    #[test]
    fn test_ema_calculation() {
        let mut ema = ExponentialMovingAverage::new(3).unwrap();

        // Test with sample prices
        let prices = [1.0, 2.0, 3.0, 4.0, 5.0];

        let mut results = Vec::new();
        for price in prices {
            ema.update(price).unwrap();
            if ema.is_ready() {
                results.push(ema.value().unwrap());
            }
        }

        // For EMA(3): α = 2/(3+1) = 0.5
        // EMA starts after first value, so we get 4 results from 5 inputs
        // EMA[0] = 1.0 (first value)
        // EMA[1] = 0.5 * 2.0 + 0.5 * 1.0 = 1.5
        // EMA[2] = 0.5 * 3.0 + 0.5 * 1.5 = 2.25
        // EMA[3] = 0.5 * 4.0 + 0.5 * 2.25 = 3.125
        // EMA[4] = 0.5 * 5.0 + 0.5 * 3.125 = 4.0625

        assert_eq!(results.len(), 5); // Should have 5 values (EMA starts immediately)
        let first_ema = results[0];
        assert!(
            (0.9..=1.1).contains(&first_ema),
            "First EMA value: {}",
            first_ema
        );
        let last_ema = results[results.len() - 1];
        assert!(
            (4.0..=4.1).contains(&last_ema),
            "Last EMA value: {}",
            last_ema
        );
    }

    #[test]
    fn test_vwma_calculation() {
        let mut vwma = VolumeWeightedMovingAverage::new(2).unwrap();

        // Not enough data yet
        assert!(vwma.value().is_err());

        vwma.update(10.0, 100.0).unwrap(); // price = 10, volume = 100
        vwma.update(20.0, 200.0).unwrap(); // price = 20, volume = 200

        // Now we have enough data
        // VWMA = (10*100 + 20*200) / (100 + 200) = 16.67
        let expected = (10.0 * 100.0 + 20.0 * 200.0) / (100.0 + 200.0);
        assert!((vwma.value().unwrap() - expected).abs() < 0.001);

        // Add another value
        vwma.update(15.0, 300.0).unwrap(); // price = 15, volume = 300

        // The window slides, dropping the oldest value
        // VWMA = (20*200 + 15*300) / (200 + 300) = 17.0
        let expected = (20.0 * 200.0 + 15.0 * 300.0) / (200.0 + 300.0);
        assert!((vwma.value().unwrap() - expected).abs() < 0.001);
    }
}