use alloc::{vec, vec::Vec};
use super::mstl_impl::mstl;
use super::{Error, Float, MstlResult, StlParams};
#[derive(Clone, Debug)]
pub struct MstlParams {
iterate: usize,
lambda: Option<f32>,
swin: Option<Vec<usize>>,
stl_params: StlParams,
}
impl MstlParams {
pub fn new() -> Self {
Self {
iterate: 2,
lambda: None,
swin: None,
stl_params: StlParams::new(),
}
}
pub fn iterations(&mut self, iterations: usize) -> &mut Self {
self.iterate = iterations;
self
}
pub fn lambda(&mut self, lambda: f32) -> &mut Self {
self.lambda = Some(lambda);
self
}
pub fn seasonal_lengths(&mut self, lengths: &[usize]) -> &mut Self {
self.swin = Some(lengths.to_vec());
self
}
pub fn stl_params(&mut self, stl_params: StlParams) -> &mut Self {
self.stl_params = stl_params;
self
}
pub fn fit<T: Float>(&self, series: &[T], periods: &[usize]) -> Result<MstlResult<T>, Error> {
if periods.is_empty() {
return Err(Error::Parameter("periods must not be empty"));
}
if periods.iter().any(|&v| v < 2) {
return Err(Error::Parameter("period must be at least 2"));
}
for np in periods {
if series.len() / 2 < *np {
return Err(Error::Series("series has less than two periods"));
}
}
if let Some(lambda) = self.lambda {
if !(0.0..=1.0).contains(&lambda) {
return Err(Error::Parameter("lambda must be between 0 and 1"));
}
}
if let Some(swin) = &self.swin {
if swin.len() != periods.len() {
return Err(Error::Parameter(
"seasonal_lengths must have the same length as periods",
));
}
}
let n = series.len();
let mut seasonal = Vec::with_capacity(periods.len());
for _ in 0..periods.len() {
seasonal.push(vec![T::zero(); n]);
}
let mut trend = vec![T::zero(); n];
let mut remainder = vec![T::zero(); n];
let mut weights = vec![T::zero(); n];
let mut work = vec![T::zero(); (n + 2 * periods.iter().max().unwrap()) * 5];
mstl(
series,
periods,
self.iterate,
self.lambda,
&self.swin,
&self.stl_params,
&mut seasonal,
&mut trend,
&mut remainder,
&mut weights,
&mut work,
)?;
Ok(MstlResult {
seasonal,
trend,
remainder,
})
}
}
impl Default for MstlParams {
fn default() -> Self {
Self::new()
}
}