anofox-forecast 0.8.0

//! MSTL Forecaster - Decomposition-based forecasting with multiple seasonality.
//!
//! MSTLForecaster uses Multiple Seasonal-Trend decomposition using LOESS (MSTL)
//! to decompose a time series, then forecasts each component separately and
//! combines them.
//!
//! This matches the statsforecast MSTL implementation which uses MSTL decomposition
//! followed by trend forecasting (default: AutoETS).

use crate::core::{Forecast, TimeSeries};
use crate::error::{ForecastError, Result};
use crate::models::explain::{Explainable, ForecastExplanation};
use crate::models::exponential::{AutoETS, AutoETSConfig, SimpleExponentialSmoothing};
use crate::models::inspect::{Explanation, Inspectable, MstlExplanation};
use crate::models::{validate_series_complete, Forecaster};
use crate::seasonality::{MSTLResult, MSTL};
use crate::utils::ols::OLSResult;
use std::collections::HashMap;

/// Method for forecasting the deseasonalized (trend + remainder) component.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TrendForecastMethod {
    /// Use AutoETS for trend forecasting (default, most accurate).
    #[default]
    AutoETS,
    /// Use Simple Exponential Smoothing (faster).
    SES,
    /// Use linear extrapolation.
    Linear,
    /// Use last value (naive).
    Naive,
}

/// Method for forecasting seasonal components.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum SeasonalForecastMethod {
    /// Repeat the last seasonal cycle (default, matches statsforecast).
    #[default]
    Naive,
    /// Average of all seasonal cycles.
    Average,
}

/// MSTL-based forecaster for time series with multiple seasonalities.
///
/// The MSTL forecaster:
/// 1. Decomposes the series using MSTL into trend, multiple seasonal components, and remainder
/// 2. Forecasts the deseasonalized series (trend + remainder) using the specified method
/// 3. Projects seasonal components forward using naive repetition
/// 4. Combines forecasts: forecast = trend_forecast + sum(seasonal_forecasts)
///
/// # Example
/// ```
/// use anofox_forecast::models::mstl_forecaster::MSTLForecaster;
/// use anofox_forecast::models::Forecaster;
/// use anofox_forecast::core::TimeSeries;
/// use chrono::{TimeZone, Utc};
///
/// // Create hourly data with daily seasonality
/// let timestamps: Vec<_> = (0..100).map(|i| Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap() + chrono::Duration::hours(i)).collect();
/// let values: Vec<f64> = (0..100).map(|i| {
///     50.0 + 0.1 * i as f64  // trend
///         + 5.0 * (2.0 * std::f64::consts::PI * i as f64 / 24.0).sin()  // daily
/// }).collect();
/// let ts = TimeSeries::univariate(timestamps, values).unwrap();
///
/// let mut model = MSTLForecaster::new(vec![24]);
/// model.fit(&ts).unwrap();
/// let forecast = model.predict(24).unwrap();
/// ```
#[derive(Debug, Clone)]
pub struct MSTLForecaster {
    /// Seasonal periods for MSTL decomposition.
    seasonal_periods: Vec<usize>,
    /// Number of MSTL iterations.
    mstl_iterations: usize,
    /// Use robust fitting in MSTL.
    robust: bool,
    /// Method for forecasting trend component.
    trend_method: TrendForecastMethod,
    /// Method for forecasting seasonal components.
    seasonal_method: SeasonalForecastMethod,
    /// Decomposition result.
    decomposition: Option<MSTLResult>,
    /// Fitted trend forecaster.
    trend_forecaster: Option<Box<dyn TrendForecasterTrait>>,
    /// Original series length.
    n: usize,
    /// Fitted values.
    fitted: Option<Vec<f64>>,
    /// Residuals.
    residuals: Option<Vec<f64>>,
    /// Residual variance for confidence intervals.
    residual_variance: Option<f64>,
    /// OLS result from pre-regression (when fit with exogenous regressors).
    ols_result: Option<OLSResult>,
    /// Names of exogenous regressors used during fitting.
    exog_name_list: Vec<String>,
    /// Training input values, retained for the issue #106 decomposition contract.
    training_values: Option<Vec<f64>>,
    /// Training-time exogenous regressors, retained for the residual-Ridge shim.
    training_regressors: Option<std::collections::HashMap<String, Vec<f64>>>,
    /// Per-row sum of `decomposition.seasonal_components`, materialised once
    /// at fit time so `seasonal_component()` can return a `&[f64]`.
    seasonal_sum: Option<Vec<f64>>,
}

/// Internal trait for trend forecasters (to allow different types).
trait TrendForecasterTrait: std::fmt::Debug + Send + Sync {
    fn predict(&self, horizon: usize) -> Result<Vec<f64>>;
    fn clone_box(&self) -> Box<dyn TrendForecasterTrait>;
}

impl Clone for Box<dyn TrendForecasterTrait> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

/// Wrapper for AutoETS trend forecaster.
#[derive(Debug, Clone)]
struct AutoETSTrendForecaster {
    model: AutoETS,
}

impl TrendForecasterTrait for AutoETSTrendForecaster {
    fn predict(&self, horizon: usize) -> Result<Vec<f64>> {
        let forecast = self.model.predict(horizon)?;
        Ok(forecast.primary().to_vec())
    }

    fn clone_box(&self) -> Box<dyn TrendForecasterTrait> {
        Box::new(self.clone())
    }
}

/// Wrapper for SES trend forecaster.
#[derive(Debug, Clone)]
struct SESTrendForecaster {
    model: SimpleExponentialSmoothing,
}

impl TrendForecasterTrait for SESTrendForecaster {
    fn predict(&self, horizon: usize) -> Result<Vec<f64>> {
        let forecast = self.model.predict(horizon)?;
        Ok(forecast.primary().to_vec())
    }

    fn clone_box(&self) -> Box<dyn TrendForecasterTrait> {
        Box::new(self.clone())
    }
}

/// Linear trend forecaster.
#[derive(Debug, Clone)]
struct LinearTrendForecaster {
    intercept: f64,
    slope: f64,
    n: usize,
}

impl TrendForecasterTrait for LinearTrendForecaster {
    fn predict(&self, horizon: usize) -> Result<Vec<f64>> {
        let mut forecasts = Vec::with_capacity(horizon);
        for h in 1..=horizon {
            let t = (self.n + h) as f64;
            forecasts.push(self.intercept + self.slope * t);
        }
        Ok(forecasts)
    }

    fn clone_box(&self) -> Box<dyn TrendForecasterTrait> {
        Box::new(self.clone())
    }
}

/// Naive trend forecaster (last value).
#[derive(Debug, Clone)]
struct NaiveTrendForecaster {
    last_value: f64,
}

impl TrendForecasterTrait for NaiveTrendForecaster {
    fn predict(&self, horizon: usize) -> Result<Vec<f64>> {
        Ok(vec![self.last_value; horizon])
    }

    fn clone_box(&self) -> Box<dyn TrendForecasterTrait> {
        Box::new(self.clone())
    }
}

impl MSTLForecaster {
    /// Create a new MSTL forecaster with the given seasonal periods.
    ///
    /// # Arguments
    /// * `seasonal_periods` - Vector of seasonal periods (e.g., [24, 168] for daily/weekly in hourly data)
    pub fn new(seasonal_periods: Vec<usize>) -> Self {
        Self {
            seasonal_periods,
            mstl_iterations: 2,
            robust: false,
            trend_method: TrendForecastMethod::AutoETS,
            seasonal_method: SeasonalForecastMethod::Naive,
            decomposition: None,
            trend_forecaster: None,
            n: 0,
            fitted: None,
            residuals: None,
            residual_variance: None,
            ols_result: None,
            exog_name_list: Vec::new(),
            training_values: None,
            training_regressors: None,
            seasonal_sum: None,
        }
    }

    /// Set the number of MSTL iterations.
    pub fn with_iterations(mut self, iterations: usize) -> Self {
        self.mstl_iterations = iterations;
        self
    }

    /// Enable robust fitting in MSTL decomposition.
    pub fn robust(mut self) -> Self {
        self.robust = true;
        self
    }

    /// Set the trend forecasting method.
    pub fn with_trend_method(mut self, method: TrendForecastMethod) -> Self {
        self.trend_method = method;
        self
    }

    /// Set the seasonal forecasting method.
    pub fn with_seasonal_method(mut self, method: SeasonalForecastMethod) -> Self {
        self.seasonal_method = method;
        self
    }

    /// Get the decomposition result.
    pub fn decomposition(&self) -> Option<&MSTLResult> {
        self.decomposition.as_ref()
    }

    /// Get seasonal periods.
    pub fn seasonal_periods(&self) -> &[usize] {
        &self.seasonal_periods
    }

    /// Core prediction logic shared by `predict()` and `predict_with_exog()`.
    fn predict_base(&self, horizon: usize) -> Result<Forecast> {
        let decomposition = self
            .decomposition
            .as_ref()
            .ok_or(ForecastError::FitRequired { model: None })?;
        let trend_forecaster = self
            .trend_forecaster
            .as_ref()
            .ok_or(ForecastError::FitRequired { model: None })?;

        if horizon == 0 {
            return Ok(Forecast::new());
        }

        // Forecast deseasonalized component
        let trend_forecast = trend_forecaster.predict(horizon)?;

        // Project each seasonal component
        let mut seasonal_forecasts: Vec<Vec<f64>> = Vec::new();
        for (idx, seasonal) in decomposition.seasonal_components.iter().enumerate() {
            let period = decomposition.seasonal_periods[idx];
            let seasonal_forecast = self.project_seasonal(seasonal, period, horizon);
            seasonal_forecasts.push(seasonal_forecast);
        }

        // Combine forecasts
        let mut forecasts = trend_forecast;
        for seasonal_forecast in &seasonal_forecasts {
            for (i, &s) in seasonal_forecast.iter().enumerate() {
                forecasts[i] += s;
            }
        }

        Ok(Forecast::from_values(forecasts))
    }

    /// Project a seasonal component forward.
    fn project_seasonal(&self, seasonal: &[f64], period: usize, horizon: usize) -> Vec<f64> {
        match self.seasonal_method {
            SeasonalForecastMethod::Naive => {
                // Repeat the last cycle
                let last_cycle_start = seasonal.len().saturating_sub(period);
                let last_cycle = &seasonal[last_cycle_start..];
                (0..horizon)
                    .map(|h| last_cycle[h % last_cycle.len()])
                    .collect()
            }
            SeasonalForecastMethod::Average => {
                // Average all cycles
                let mut avg_cycle = vec![0.0; period];
                let mut counts = vec![0usize; period];
                for (i, &s) in seasonal.iter().enumerate() {
                    avg_cycle[i % period] += s;
                    counts[i % period] += 1;
                }
                for i in 0..period {
                    if counts[i] > 0 {
                        avg_cycle[i] /= counts[i] as f64;
                    }
                }
                (0..horizon).map(|h| avg_cycle[h % period]).collect()
            }
        }
    }

    /// Fit a linear trend model.
    fn fit_linear(values: &[f64]) -> (f64, f64) {
        let n = values.len();
        if n == 0 {
            return (0.0, 0.0);
        }

        let x_mean = (n - 1) as f64 / 2.0;
        let y_mean = values.iter().sum::<f64>() / n as f64;

        let mut ss_xx = 0.0;
        let mut ss_xy = 0.0;
        for (i, &y) in values.iter().enumerate() {
            let x = i as f64;
            ss_xx += (x - x_mean).powi(2);
            ss_xy += (x - x_mean) * (y - y_mean);
        }

        let slope = if ss_xx > 0.0 { ss_xy / ss_xx } else { 0.0 };
        let intercept = y_mean - slope * x_mean;

        (intercept, slope)
    }
}

impl Default for MSTLForecaster {
    fn default() -> Self {
        Self::new(vec![12])
    }
}

impl Forecaster for MSTLForecaster {
    fn fit(&mut self, series: &TimeSeries) -> Result<()> {
        validate_series_complete(series)?;
        let values = series.primary_values();
        self.n = values.len();

        if self.seasonal_periods.is_empty() {
            return Err(ForecastError::InvalidParameter(
                "At least one seasonal period is required".to_string(),
            ));
        }

        for &period in &self.seasonal_periods {
            if period < 2 {
                return Err(ForecastError::InvalidParameter(format!(
                    "seasonal period must be >= 2, got {}",
                    period
                )));
            }
        }

        let max_period = *self.seasonal_periods.iter().max().unwrap_or(&1);
        if values.len() < 2 * max_period {
            return Err(ForecastError::InsufficientData {
                needed: 2 * max_period,
                got: values.len(),
                hint: Some(format!(
                    "MSTL requires at least 2 * max_period = {} observations for seasonal decomposition",
                    2 * max_period
                )),
            });
        }

        // Create MSTL decomposer
        let mut mstl =
            MSTL::new(self.seasonal_periods.clone()).with_iterations(self.mstl_iterations);
        if self.robust {
            mstl = mstl.robust();
        }

        // Check for exogenous regressors
        let regressors = series.all_regressors();

        // Decompose (with or without pre-regression)
        let decomposition = if regressors.is_empty() {
            self.ols_result = None;
            self.exog_name_list.clear();
            mstl.decompose(values).ok_or_else(|| {
                ForecastError::ComputationError("MSTL decomposition failed".to_string())
            })?
        } else {
            let result = mstl
                .decompose_with_regressors(values, &regressors)
                .ok_or_else(|| {
                    ForecastError::ComputationError(
                        "MSTL decomposition with regressors failed".to_string(),
                    )
                })?;
            self.ols_result = result.regressor_coefficients.clone();
            self.exog_name_list = regressors.keys().cloned().collect();
            self.exog_name_list.sort();
            result
        };

        // Create deseasonalized series (trend + remainder)
        let deseasonalized: Vec<f64> = decomposition
            .trend
            .iter()
            .zip(decomposition.remainder.iter())
            .map(|(t, r)| t + r)
            .collect();

        // Create a temporary time series for the deseasonalized component
        let deseas_ts =
            TimeSeries::univariate(series.timestamps().to_vec(), deseasonalized.clone())?;

        // Fit trend forecaster based on method
        let trend_forecaster: Box<dyn TrendForecasterTrait> = match self.trend_method {
            TrendForecastMethod::AutoETS => {
                let config = AutoETSConfig::non_seasonal();
                let mut model = AutoETS::with_config(config);
                model.fit(&deseas_ts)?;
                Box::new(AutoETSTrendForecaster { model })
            }
            TrendForecastMethod::SES => {
                let mut model = SimpleExponentialSmoothing::new(0.3);
                model.fit(&deseas_ts)?;
                Box::new(SESTrendForecaster { model })
            }
            TrendForecastMethod::Linear => {
                let (intercept, slope) = Self::fit_linear(&deseasonalized);
                Box::new(LinearTrendForecaster {
                    intercept,
                    slope,
                    n: self.n,
                })
            }
            TrendForecastMethod::Naive => {
                let last_value = *deseasonalized.last().unwrap_or(&0.0);
                Box::new(NaiveTrendForecaster { last_value })
            }
        };

        // Compute fitted values (reconstruct from components)
        let regressor_effect = decomposition.regressor_effect.as_ref();
        let fitted: Vec<f64> = (0..self.n)
            .map(|i| {
                let mut val = decomposition.trend[i] + decomposition.remainder[i];
                for seasonal in &decomposition.seasonal_components {
                    val += seasonal[i];
                }
                if let Some(effect) = regressor_effect {
                    val += effect[i];
                }
                val
            })
            .collect();

        // Compute residuals
        let residuals: Vec<f64> = values
            .iter()
            .zip(fitted.iter())
            .map(|(y, f)| y - f)
            .collect();

        // Compute residual variance
        if residuals.len() > 1 {
            let variance = crate::simd::sum_of_squares(&residuals) / residuals.len() as f64;
            self.residual_variance = Some(variance);
        }

        // Precompute the per-row sum of all seasonal components so
        // `seasonal_component()` can return a borrowed slice.
        let n = decomposition.trend.len();
        let mut seasonal_sum = vec![0.0_f64; n];
        for comp in &decomposition.seasonal_components {
            for (s, c) in seasonal_sum.iter_mut().zip(comp.iter()) {
                *s += *c;
            }
        }

        // Retain training values and regressors per issue #106.
        self.training_values = Some(values.to_vec());
        self.training_regressors = if regressors.is_empty() {
            None
        } else {
            Some(regressors.clone())
        };
        self.seasonal_sum = Some(seasonal_sum);

        self.decomposition = Some(decomposition);
        self.trend_forecaster = Some(trend_forecaster);
        self.fitted = Some(fitted);
        self.residuals = Some(residuals);

        Ok(())
    }

    fn predict(&self, horizon: usize) -> Result<Forecast> {
        if self.ols_result.is_some() {
            return Err(ForecastError::InvalidParameter(
                "Model was fit with exogenous regressors. Use predict_with_exog() instead."
                    .to_string(),
            ));
        }
        self.predict_base(horizon)
    }

    fn predict_with_intervals(&self, horizon: usize, confidence: f64) -> Result<Forecast> {
        let forecast = self.predict(horizon)?;
        let variance = self.residual_variance.unwrap_or(0.0);

        if horizon == 0 || variance <= 0.0 {
            return Ok(forecast);
        }

        let z = crate::utils::stats::quantile_normal((1.0 + confidence) / 2.0);
        let se = variance.sqrt();

        let preds = forecast.primary();
        let mut lower = Vec::with_capacity(horizon);
        let mut upper = Vec::with_capacity(horizon);

        // Simple constant variance intervals (could be improved with fan-out)
        for h in 0..horizon {
            // Variance increases with horizon (simple approximation)
            let h_factor = (1.0 + 0.1 * h as f64).sqrt();
            lower.push(preds[h] - z * se * h_factor);
            upper.push(preds[h] + z * se * h_factor);
        }

        Ok(Forecast::from_values_with_intervals(
            preds.to_vec(),
            lower,
            upper,
        ))
    }

    fn fitted_values(&self) -> Option<&[f64]> {
        self.fitted.as_deref()
    }

    fn fitted_values_with_intervals(&self, level: f64) -> Option<Forecast> {
        let fitted = self.fitted.as_ref()?;
        let variance = self.residual_variance?;

        if variance <= 0.0 {
            return Some(Forecast::from_values(fitted.clone()));
        }

        let z = crate::utils::stats::quantile_normal((1.0 + level) / 2.0);
        let sigma = variance.sqrt();

        let lower: Vec<f64> = fitted.iter().map(|&f| f - z * sigma).collect();
        let upper: Vec<f64> = fitted.iter().map(|&f| f + z * sigma).collect();

        Some(Forecast::from_values_with_intervals(
            fitted.clone(),
            lower,
            upper,
        ))
    }

    fn residuals(&self) -> Option<&[f64]> {
        self.residuals.as_deref()
    }

    fn trend_component(&self) -> Result<&[f64]> {
        self.decomposition
            .as_ref()
            .map(|d| d.trend.as_slice())
            .ok_or(ForecastError::FitRequired {
                model: Some("MSTLForecaster".into()),
            })
    }

    /// For MSTL, the decomposition "residual" is the STL remainder — what's
    /// left of training values after extracting trend and seasonal. The
    /// post-fit `residuals` field (`training - fitted_values`) is near
    /// zero because the decomposition exactly reconstructs the training
    /// series, so it isn't the right input for the issue #106 invariant.
    fn residual_component(&self) -> Result<Vec<f64>> {
        self.decomposition
            .as_ref()
            .map(|d| d.remainder.clone())
            .ok_or(ForecastError::FitRequired {
                model: Some("MSTLForecaster".into()),
            })
    }

    fn seasonal_component(&self) -> Result<&[f64]> {
        if self.seasonal_periods.is_empty() {
            return Err(ForecastError::InvalidParameter(
                "MSTLForecaster fit has no seasonal contribution".into(),
            ));
        }
        self.seasonal_sum
            .as_deref()
            .ok_or(ForecastError::FitRequired {
                model: Some("MSTLForecaster".into()),
            })
    }

    fn training_values(&self) -> Result<&[f64]> {
        self.training_values
            .as_deref()
            .ok_or(ForecastError::FitRequired {
                model: Some("MSTLForecaster".into()),
            })
    }

    fn training_regressors(&self) -> Option<&HashMap<String, Vec<f64>>> {
        self.training_regressors.as_ref()
    }

    fn name(&self) -> &str {
        "MSTLForecaster"
    }

    fn supports_exog(&self) -> bool {
        true
    }

    fn has_exog(&self) -> bool {
        self.ols_result.is_some()
    }

    fn exog_names(&self) -> Option<&[String]> {
        if self.exog_name_list.is_empty() {
            None
        } else {
            Some(&self.exog_name_list)
        }
    }

    fn exog_coefficients(&self) -> Option<&OLSResult> {
        self.ols_result.as_ref()
    }

    fn predict_with_exog(
        &self,
        horizon: usize,
        future_regressors: &HashMap<String, Vec<f64>>,
    ) -> Result<Forecast> {
        let ols = self.ols_result.as_ref().ok_or_else(|| {
            ForecastError::InvalidParameter(
                "Model was not fit with exogenous regressors".to_string(),
            )
        })?;

        // Get base forecast (trend + seasonal)
        let base = self.predict_base(horizon)?;

        if horizon == 0 {
            return Ok(base);
        }

        // Compute future regressor effect
        let future_effect = ols.predict(future_regressors)?;
        if future_effect.len() != horizon {
            return Err(ForecastError::DimensionMismatch {
                expected: horizon,
                got: future_effect.len(),
            });
        }

        // Add regressor effect to base forecast
        let combined: Vec<f64> = base
            .primary()
            .iter()
            .zip(future_effect.iter())
            .map(|(b, e)| b + e)
            .collect();

        Ok(Forecast::from_values(combined))
    }
}

impl Inspectable for MSTLForecaster {
    fn explanation(&self) -> Result<Explanation> {
        let decomposition =
            self.decomposition
                .as_ref()
                .ok_or_else(|| ForecastError::FitRequired {
                    model: Some("MSTLForecaster".to_string()),
                })?;
        let fitted = self
            .fitted
            .clone()
            .ok_or_else(|| ForecastError::FitRequired {
                model: Some("MSTLForecaster".to_string()),
            })?;
        let residuals = self
            .residuals
            .clone()
            .ok_or_else(|| ForecastError::FitRequired {
                model: Some("MSTLForecaster".to_string()),
            })?;
        let seasonal_component = self.seasonal_sum.clone().unwrap_or_default();

        Ok(Explanation::Mstl(MstlExplanation {
            seasonal_periods: self.seasonal_periods.clone(),
            iterations: self.mstl_iterations,
            fitted_values: fitted,
            trend_component: decomposition.trend.clone(),
            seasonal_component,
            residuals,
        }))
    }
}

impl Explainable for MSTLForecaster {
    fn explain(&self, horizon: usize) -> Result<ForecastExplanation> {
        let decomposition = self
            .decomposition
            .as_ref()
            .ok_or(ForecastError::FitRequired { model: None })?;
        let trend_forecaster = self
            .trend_forecaster
            .as_ref()
            .ok_or(ForecastError::FitRequired { model: None })?;
        if horizon == 0 {
            return Ok(ForecastExplanation {
                level: vec![],
                trend: None,
                seasonal: None,
                residual: None,
                named_components: vec![],
            });
        }
        let trend_forecast = trend_forecaster.predict(horizon)?;
        let mut named_components = Vec::new();
        for (idx, seasonal) in decomposition.seasonal_components.iter().enumerate() {
            let period = decomposition.seasonal_periods[idx];
            let sf = self.project_seasonal(seasonal, period, horizon);
            named_components.push((format!("seasonal_{}", period), sf));
        }
        Ok(ForecastExplanation {
            level: trend_forecast,
            trend: None,
            seasonal: None,
            residual: None,
            named_components,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{Duration, TimeZone, Utc};

    fn make_timestamps(n: usize) -> Vec<chrono::DateTime<Utc>> {
        let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        (0..n).map(|i| base + Duration::hours(i as i64)).collect()
    }

    fn make_multi_seasonal_series(n: usize, periods: &[usize]) -> TimeSeries {
        let timestamps = make_timestamps(n);
        let values: Vec<f64> = (0..n)
            .map(|i| {
                let trend = 50.0 + 0.1 * i as f64;
                let mut seasonal = 0.0;
                for (idx, &period) in periods.iter().enumerate() {
                    let amplitude = 5.0 / (idx + 1) as f64;
                    seasonal +=
                        amplitude * (2.0 * std::f64::consts::PI * i as f64 / period as f64).sin();
                }
                trend + seasonal
            })
            .collect();
        TimeSeries::univariate(timestamps, values).unwrap()
    }

    #[test]
    fn mstl_forecaster_basic() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_multiple_seasonalities() {
        let ts = make_multi_seasonal_series(200, &[12, 24]);
        let mut model = MSTLForecaster::new(vec![12, 24]);
        model.fit(&ts).unwrap();

        let forecast = model.predict(24).unwrap();
        assert_eq!(forecast.horizon(), 24);

        // Decomposition should have 2 seasonal components
        let decomp = model.decomposition().unwrap();
        assert_eq!(decomp.seasonal_components.len(), 2);
    }

    #[test]
    fn mstl_forecaster_with_ses() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]).with_trend_method(TrendForecastMethod::SES);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_with_linear() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model =
            MSTLForecaster::new(vec![12]).with_trend_method(TrendForecastMethod::Linear);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_with_naive() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]).with_trend_method(TrendForecastMethod::Naive);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);

        // Naive produces flat trend component
        // (though seasonal pattern still varies)
    }

    #[test]
    fn mstl_forecaster_robust() {
        // Generate series with outliers
        let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let timestamps: Vec<_> = (0..100).map(|i| base + Duration::hours(i)).collect();
        let mut values: Vec<f64> = (0..100)
            .map(|i| {
                let trend = 50.0 + 0.5 * i as f64;
                let seasonal = 10.0 * (2.0 * std::f64::consts::PI * (i % 12) as f64 / 12.0).sin();
                trend + seasonal
            })
            .collect();
        // Add outliers
        values[30] = 200.0;
        values[60] = -50.0;
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = MSTLForecaster::new(vec![12]).robust();
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_confidence_intervals() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        let forecast = model.predict_with_intervals(12, 0.95).unwrap();
        assert!(forecast.has_lower());
        assert!(forecast.has_upper());

        let lower = forecast.lower_series(0).unwrap();
        let upper = forecast.upper_series(0).unwrap();
        let preds = forecast.primary();

        for i in 0..12 {
            assert!(
                lower[i] <= preds[i],
                "Lower bound {} should be <= prediction {} at index {}",
                lower[i],
                preds[i],
                i
            );
            assert!(
                upper[i] >= preds[i],
                "Upper bound {} should be >= prediction {} at index {}",
                upper[i],
                preds[i],
                i
            );
        }
    }

    #[test]
    fn mstl_forecaster_fitted_residuals() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        assert!(model.fitted_values().is_some());
        assert!(model.residuals().is_some());
        assert_eq!(model.fitted_values().unwrap().len(), 100);
        assert_eq!(model.residuals().unwrap().len(), 100);
    }

    #[test]
    fn mstl_forecaster_insufficient_data() {
        let ts = make_multi_seasonal_series(20, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        assert!(model.fit(&ts).is_err());
    }

    #[test]
    fn mstl_forecaster_empty_periods() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![]);
        assert!(model.fit(&ts).is_err());
    }

    #[test]
    fn mstl_forecaster_requires_fit() {
        let model = MSTLForecaster::new(vec![12]);
        assert!(matches!(
            model.predict(5),
            Err(ForecastError::FitRequired { .. })
        ));
    }

    #[test]
    fn mstl_forecaster_zero_horizon() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        let forecast = model.predict(0).unwrap();
        assert_eq!(forecast.horizon(), 0);
    }

    #[test]
    fn mstl_forecaster_name() {
        let model = MSTLForecaster::new(vec![12]);
        assert_eq!(model.name(), "MSTLForecaster");
    }

    #[test]
    fn mstl_forecaster_seasonal_average_method() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model =
            MSTLForecaster::new(vec![12]).with_seasonal_method(SeasonalForecastMethod::Average);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_with_iterations() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]).with_iterations(3);
        model.fit(&ts).unwrap();

        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    fn make_series_with_regressor(n: usize, periods: &[usize]) -> (TimeSeries, Vec<f64>) {
        use crate::core::CalendarAnnotations;

        let timestamps = make_timestamps(n);
        let x: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin() * 10.0).collect();
        let values: Vec<f64> = (0..n)
            .map(|i| {
                let trend = 50.0 + 0.1 * i as f64;
                let mut seasonal = 0.0;
                for (idx, &period) in periods.iter().enumerate() {
                    let amplitude = 5.0 / (idx + 1) as f64;
                    seasonal +=
                        amplitude * (2.0 * std::f64::consts::PI * i as f64 / period as f64).sin();
                }
                trend + seasonal + 2.0 * x[i]
            })
            .collect();

        let calendar = CalendarAnnotations::new().with_regressor("x".to_string(), x.clone());
        let mut ts = TimeSeries::univariate(timestamps, values).unwrap();
        ts.set_calendar(calendar);
        (ts, x)
    }

    #[test]
    fn mstl_forecaster_supports_exog() {
        let model = MSTLForecaster::new(vec![12]);
        assert!(model.supports_exog());
        assert!(!model.has_exog());
        assert!(model.exog_names().is_none());
    }

    #[test]
    fn mstl_forecaster_fit_with_regressors() {
        let (ts, _x) = make_series_with_regressor(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        assert!(model.has_exog());
        assert_eq!(model.exog_names(), Some(&["x".to_string()][..]));
    }

    #[test]
    fn mstl_forecaster_predict_guards_exog() {
        let (ts, _x) = make_series_with_regressor(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        // predict() should error when exog was used
        let result = model.predict(12);
        assert!(result.is_err());
    }

    #[test]
    fn mstl_forecaster_predict_with_exog() {
        let (ts, _x) = make_series_with_regressor(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        // Provide future regressor values
        let future_x: Vec<f64> = (100..112).map(|i| (i as f64 * 0.1).sin() * 10.0).collect();
        let mut future_regressors = HashMap::new();
        future_regressors.insert("x".to_string(), future_x);

        let forecast = model.predict_with_exog(12, &future_regressors).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_exog_improves_fit() {
        // Fit with and without exog on a series that has a strong regressor effect
        let (ts_with_exog, _x) = make_series_with_regressor(100, &[12]);

        // Fit with exog (pre-regression)
        let mut model_exog = MSTLForecaster::new(vec![12]);
        model_exog.fit(&ts_with_exog).unwrap();

        // Fit without exog (just uses the raw values)
        let ts_no_exog = make_multi_seasonal_series(100, &[12]);
        let mut model_plain = MSTLForecaster::new(vec![12]);
        model_plain.fit(&ts_no_exog).unwrap();

        // Both should produce fitted values
        assert!(model_exog.fitted_values().is_some());
        assert!(model_plain.fitted_values().is_some());
    }

    #[test]
    fn mstl_forecaster_predict_with_exog_missing_regressor() {
        let (ts, _x) = make_series_with_regressor(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        // Missing regressor
        let future_regressors = HashMap::new();
        let result = model.predict_with_exog(12, &future_regressors);
        assert!(result.is_err());
    }

    #[test]
    fn constant_series_produces_constant_forecast() {
        let timestamps = make_timestamps(40);
        let values = vec![5.0; 40];
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = MSTLForecaster::new(vec![4]);
        model.fit(&ts).unwrap();

        let forecast = model.predict(8).unwrap();
        let preds = forecast.primary();

        assert!(
            preds.iter().all(|v| v.is_finite()),
            "All predictions must be finite, got: {:?}",
            preds
        );
        for &p in preds {
            assert!((p - 5.0).abs() < 0.5, "Expected ~5.0, got {}", p);
        }
    }

    #[test]
    fn mstl_forecaster_no_exog_predict_works() {
        // When fit without exog, predict() should work normally
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();

        assert!(!model.has_exog());
        let forecast = model.predict(12).unwrap();
        assert_eq!(forecast.horizon(), 12);
    }

    #[test]
    fn mstl_forecaster_rejects_period_zero() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![0]);
        assert!(matches!(
            model.fit(&ts),
            Err(ForecastError::InvalidParameter(ref msg)) if msg.contains("seasonal period must be >= 2")
        ));
    }

    #[test]
    fn mstl_forecaster_rejects_period_one() {
        let ts = make_multi_seasonal_series(100, &[12]);
        let mut model = MSTLForecaster::new(vec![1]);
        assert!(matches!(
            model.fit(&ts),
            Err(ForecastError::InvalidParameter(ref msg)) if msg.contains("seasonal period must be >= 2")
        ));
    }

    #[test]
    fn mstl_forecaster_rejects_any_period_below_two() {
        let ts = make_multi_seasonal_series(200, &[12, 24]);
        let mut model = MSTLForecaster::new(vec![12, 1]);
        assert!(matches!(
            model.fit(&ts),
            Err(ForecastError::InvalidParameter(ref msg)) if msg.contains("seasonal period must be >= 2")
        ));
    }

    // ── Issue #106 — Decomposable trait additions ────────────────────────

    #[test]
    fn mstl_decomposition_invariant_trend_plus_seasonal_plus_residual_equals_fitted() {
        let ts = make_multi_seasonal_series(200, &[24]);
        let mut model = MSTLForecaster::new(vec![24]);
        model.fit(&ts).unwrap();

        let trend = model.trend_component().unwrap();
        let seasonal = model.seasonal_component().unwrap();
        let residual = model.residual_component().unwrap();
        let fitted = model.fitted_values().unwrap();

        assert_eq!(trend.len(), fitted.len());
        assert_eq!(seasonal.len(), fitted.len());
        assert_eq!(residual.len(), fitted.len());

        for i in 0..fitted.len() {
            let reconstructed = trend[i] + seasonal[i] + residual[i];
            // fitted = train - residual, and trend + seasonal == fitted by MSTL definition
            // → trend + seasonal + residual ≈ training value, not fitted itself.
            // The relevant invariant on the issue's wording is residual ≈ train - fitted,
            // which we check explicitly below.
            let _ = reconstructed; // silence unused-on-pretty-print
        }

        let training = model.training_values().unwrap();
        for i in 0..fitted.len() {
            let sum = trend[i] + seasonal[i] + residual[i];
            assert!(
                (sum - training[i]).abs() < 1e-9,
                "decomposition must sum to training at i={}: trend={} + seasonal={} + residual={} = {}, training={}",
                i,
                trend[i],
                seasonal[i],
                residual[i],
                sum,
                training[i]
            );
        }
    }

    #[test]
    fn mstl_training_values_retained_at_fit_time() {
        let ts = make_multi_seasonal_series(60, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();
        let training = model.training_values().unwrap();
        assert_eq!(training, ts.primary_values());
    }

    #[test]
    fn mstl_training_regressors_retained_when_fit_with_regs() {
        let (ts, x) = make_series_with_regressor(60, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();
        let regs = model
            .training_regressors()
            .expect("regressors should be retained");
        assert_eq!(regs.len(), 1);
        assert_eq!(regs.get("x").unwrap(), &x);
    }

    #[test]
    fn mstl_training_regressors_is_none_without_regs() {
        let ts = make_multi_seasonal_series(60, &[12]);
        let mut model = MSTLForecaster::new(vec![12]);
        model.fit(&ts).unwrap();
        assert!(model.training_regressors().is_none());
    }

    #[test]
    fn mstl_trend_component_requires_fit() {
        let model = MSTLForecaster::new(vec![12]);
        assert!(matches!(
            model.trend_component(),
            Err(ForecastError::FitRequired { .. })
        ));
    }
}