use std::collections::VecDeque;
use crate::{
error::{FinError, FinErrorType},
traits::{Apply, Current, Evaluate},
};
use indicato_rs_proc::{Apply, Evaluate};
use crate::traits::{Executable, ExecutionContext, IoState};
fn calculate_sma(input: f64, period: usize, values: &mut VecDeque<f64>) -> f64 {
values.push_back(input);
if values.len() > period {
values.pop_front();
}
values.iter().sum::<f64>() / values.len() as f64
}
#[derive(Apply, Evaluate)]
pub struct SimpleMovingAverage {
period: usize,
values: VecDeque<f64>,
}
impl IoState for SimpleMovingAverage {
type Input = f64;
type Output = f64;
}
impl SimpleMovingAverage {
pub fn new(period: usize) -> Result<Self, FinError> {
match period {
0 => Err(FinError::new(
FinErrorType::InvalidInput,
"Period must be greater than 0",
)),
_ => Ok(Self {
period,
values: VecDeque::with_capacity(period + 1),
}),
}
}
}
impl Executable for SimpleMovingAverage {
fn execute(
&mut self,
input: Self::Input,
execution_context: &ExecutionContext,
) -> Self::Output {
match execution_context {
ExecutionContext::Apply => calculate_sma(input, self.period, &mut self.values),
ExecutionContext::Evaluate => {
let mut values = self.values.clone();
calculate_sma(input, self.period, &mut values)
}
}
}
}
impl Current for SimpleMovingAverage {
fn current(&self) -> Self::Output {
if self.values.is_empty() {
0.0
} else {
self.values.iter().sum::<f64>() / self.values.len() as f64
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_apply() {
let mut sma = SimpleMovingAverage::new(3).unwrap();
assert_eq!(sma.apply(1.0), 1.0);
assert_eq!(sma.apply(2.0), 1.5);
assert_eq!(sma.apply(3.0), 2.0);
assert_eq!(sma.apply(4.0), 3.0);
assert_eq!(sma.apply(5.0), 4.0);
}
#[test]
fn test_evaluate() {
let mut sma = SimpleMovingAverage::new(3).unwrap();
assert_eq!(sma.apply(1.0), 1.0);
assert_eq!(sma.apply(2.0), 1.5);
assert_eq!(sma.apply(3.0), 2.0);
assert_eq!(sma.apply(4.0), 3.0);
assert_eq!(sma.evaluate(5.0), 4.0);
assert_eq!(sma.apply(5.0), 4.0);
}
}