use std::collections::VecDeque;
use std::fmt;
use std::time::Duration;
use crate::indicators::AdaptiveTimeDetector;
use crate::{errors::Result, Next, NextBatch, Reset};
use chrono::{DateTime, Utc};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
const MAX_WINDOW_SIZE: usize = 500;
const KEEP_OLDEST: usize = 10;
const KEEP_RECENT: usize = 100;
#[doc(alias = "SMA")]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct SimpleMovingAverage {
duration: Duration,
window: VecDeque<(DateTime<Utc>, f64)>,
sum: f64,
detector: AdaptiveTimeDetector,
#[cfg_attr(feature = "serde", serde(skip))]
cached_window: Option<i64>,
}
impl SimpleMovingAverage {
pub fn get_window(&self) -> VecDeque<(DateTime<Utc>, f64)> {
self.window.clone()
}
pub fn new(duration: Duration) -> Result<Self> {
if duration.as_secs() == 0 && duration.subsec_nanos() == 0 {
return Err(crate::errors::TaError::InvalidParameter);
}
Ok(Self {
duration,
window: VecDeque::new(),
sum: 0.0,
detector: AdaptiveTimeDetector::new(duration),
cached_window: None,
})
}
pub fn get_internal_state(&self) -> (Duration, VecDeque<(DateTime<Utc>, f64)>, f64) {
(self.duration, self.window.clone(), self.sum)
}
fn remove_old_data(&mut self, current_time: DateTime<Utc>) {
let dur_nanos = *self
.cached_window
.get_or_insert_with(|| self.duration.as_nanos() as i64);
let cutoff_nanos = current_time.timestamp_nanos_opt().unwrap_or(i64::MIN) - dur_nanos;
while self.window.front().map_or(false, |(time, _)| {
time.timestamp_nanos_opt().unwrap_or(i64::MIN) <= cutoff_nanos
}) {
if let Some((_, value)) = self.window.pop_front() {
self.sum -= value;
}
}
}
fn thin_window(&mut self) {
if self.window.len() <= MAX_WINDOW_SIZE {
return;
}
let len = self.window.len();
let middle_start = KEEP_OLDEST;
let middle_end = len.saturating_sub(KEEP_RECENT);
if middle_end <= middle_start {
return;
}
let mut new_window = VecDeque::with_capacity(MAX_WINDOW_SIZE);
let mut new_sum = 0.0;
for i in 0..middle_start.min(len) {
let entry = self.window[i];
new_sum += entry.1;
new_window.push_back(entry);
}
let mut keep = true;
for i in middle_start..middle_end {
if keep {
let entry = self.window[i];
new_sum += entry.1;
new_window.push_back(entry);
}
keep = !keep;
}
for i in middle_end..len {
let entry = self.window[i];
new_sum += entry.1;
new_window.push_back(entry);
}
self.window = new_window;
self.sum = new_sum;
}
}
impl Next<f64> for SimpleMovingAverage {
type Output = f64;
fn next(&mut self, (timestamp, value): (DateTime<Utc>, f64)) -> Self::Output {
let should_replace = self.detector.should_replace(timestamp);
self.remove_old_data(timestamp);
if should_replace && !self.window.is_empty() {
if let Some((_, old_value)) = self.window.pop_back() {
self.sum -= old_value;
}
}
self.window.push_back((timestamp, value));
self.sum += value;
self.thin_window();
if self.window.is_empty() {
0.0
} else {
self.sum / self.window.len() as f64
}
}
}
impl NextBatch<f64> for SimpleMovingAverage {}
impl Reset for SimpleMovingAverage {
fn reset(&mut self) {
self.window.clear();
self.sum = 0.0;
self.detector.reset();
}
}
impl Default for SimpleMovingAverage {
fn default() -> Self {
Self::new(Duration::from_secs(14 * 24 * 60 * 60)).unwrap() }
}
impl fmt::Display for SimpleMovingAverage {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "SMA({}s)", self.duration.as_secs())
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{TimeZone, Utc};
#[test]
fn test_new() {
assert!(SimpleMovingAverage::new(Duration::from_secs(0)).is_err());
assert!(SimpleMovingAverage::new(Duration::from_secs(1)).is_ok());
}
#[test]
fn test_next() {
let duration = Duration::from_secs(4);
let mut sma = SimpleMovingAverage::new(duration).unwrap();
let start_time = Utc::now();
let elapsed_time = chrono::Duration::seconds(1);
assert_eq!(sma.next((start_time, 4.0)), 4.0);
assert_eq!(sma.next((start_time + elapsed_time, 5.0)), 4.5);
assert_eq!(sma.next((start_time + elapsed_time * 2, 6.0)), 5.0);
assert_eq!(sma.next((start_time + elapsed_time * 3, 6.0)), 5.25);
assert_eq!(sma.next((start_time + elapsed_time * 4, 6.0)), 5.75);
assert_eq!(sma.next((start_time + elapsed_time * 5, 6.0)), 6.0);
assert_eq!(sma.next((start_time + elapsed_time * 6, 2.0)), 5.0);
let chrono_duration = chrono::Duration::from_std(duration).unwrap();
assert_eq!(
sma.next((start_time + elapsed_time * 6 + chrono_duration, 2.0)),
2.0
);
}
#[test]
fn test_reset() {
let duration = Duration::from_secs(4);
let mut sma = SimpleMovingAverage::new(duration).unwrap();
let start_time = Utc::now();
let elapsed_time = chrono::Duration::seconds(1);
assert_eq!(sma.next((start_time, 4.0)), 4.0);
assert_eq!(sma.next((start_time + elapsed_time, 5.0)), 4.5);
assert_eq!(sma.next((start_time + elapsed_time * 2, 6.0)), 5.0);
sma.reset();
assert_eq!(sma.next((start_time + elapsed_time * 3, 99.0)), 99.0);
}
#[test]
fn test_default() {
let _sma = SimpleMovingAverage::default();
}
#[test]
fn test_display() {
let indicator = SimpleMovingAverage::new(Duration::from_secs(7)).unwrap();
assert_eq!(format!("{}", indicator), "SMA(7s)");
}
}