use std::collections::VecDeque;
use crate::error::{Error, Result};
use crate::indicators::rolling_moments::ShiftedTrend;
use crate::traits::Indicator;
#[derive(Debug, Clone)]
pub struct LinRegIntercept {
period: usize,
window: VecDeque<f64>,
sum_x: f64,
denom: f64,
trend: ShiftedTrend,
}
impl LinRegIntercept {
pub fn new(period: usize) -> Result<Self> {
if period < 2 {
return Err(Error::InvalidPeriod {
message: "linear regression intercept needs period >= 2",
});
}
if period > crate::error::MAX_PERIOD {
return Err(Error::InvalidPeriod {
message: crate::error::PERIOD_ABOVE_MAX,
});
}
let n = period as f64;
let sum_x = n * (n - 1.0) / 2.0;
let sum_xx = (n - 1.0) * n * (2.0 * n - 1.0) / 6.0;
Ok(Self {
period,
window: VecDeque::with_capacity(period),
sum_x,
denom: n * sum_xx - sum_x * sum_x,
trend: ShiftedTrend::new(),
})
}
pub const fn period(&self) -> usize {
self.period
}
}
impl Indicator for LinRegIntercept {
type Input = f64;
type Output = f64;
#[inline]
fn update(&mut self, value: f64) -> Option<f64> {
if !value.is_finite() {
return None;
}
if self.window.len() == self.period {
let front = self.window.pop_front().expect("non-empty");
self.trend.slide(front);
}
let index = self.window.len();
self.window.push_back(value);
self.trend.push(value, index);
if self.trend.needs_reseed(self.period) {
self.trend.reseed(self.window.iter().copied());
}
if self.window.len() < self.period {
return None;
}
let n = self.period as f64;
let slope = (n * self.trend.sum_xy() - self.sum_x * self.trend.sum_y()) / self.denom;
let intercept = (self.trend.sum_y() - slope * self.sum_x) / n + self.trend.offset();
Some(intercept)
}
fn reset(&mut self) {
self.window.clear();
self.trend.reset();
}
#[inline]
fn warmup_period(&self) -> usize {
self.period
}
#[inline]
fn is_ready(&self) -> bool {
self.window.len() == self.period
}
#[inline]
fn name(&self) -> &'static str {
"LINEARREG_INTERCEPT"
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::traits::BatchExt;
use approx::assert_relative_eq;
#[test]
fn rejects_short_period() {
assert!(matches!(
LinRegIntercept::new(1),
Err(Error::InvalidPeriod { .. })
));
}
#[test]
fn accessors_report_config() {
let lr = LinRegIntercept::new(5).unwrap();
assert_eq!(lr.period(), 5);
assert_eq!(lr.name(), "LINEARREG_INTERCEPT");
assert_eq!(lr.warmup_period(), 5);
assert!(!lr.is_ready());
}
#[test]
fn reference_value() {
let mut lr = LinRegIntercept::new(3).unwrap();
let out: Vec<Option<f64>> = lr.batch(&[1.0, 2.0, 9.0]);
assert!(out[0].is_none());
assert!(out[1].is_none());
assert_relative_eq!(out[2].unwrap(), 0.0, epsilon = 1e-9);
assert!(lr.is_ready());
}
#[test]
fn slides_and_tracks_a_shifted_line() {
let mut lr = LinRegIntercept::new(3).unwrap();
let out: Vec<Option<f64>> = lr.batch(&[1.0, 10.0, 12.0, 14.0]);
assert_relative_eq!(out[3].unwrap(), 10.0, epsilon = 1e-9);
}
#[test]
fn reset_clears_state() {
let mut lr = LinRegIntercept::new(3).unwrap();
let _ = lr.batch(&[1.0, 2.0, 9.0]);
assert!(lr.is_ready());
lr.reset();
assert!(!lr.is_ready());
assert_eq!(lr.update(1.0), None);
}
}