rustytorch/demos/rtx-timeseries-demo/src/forecaster.rs

//! Time series forecaster wrapper for the demo

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use rtx_tensor::{Device, Tensor};
use rtx_timeseries::models::{FitMetrics as RtxFitMetrics, ForecastOutput, TimeSeriesModel};
use rtx_timeseries::{
    ARIMAModel, ExponentialSmoothingModel, NeuralProphetModel, ProphetModel,
    TransformerForecastModel,
};
use timeseries_shared::{
    DataPoint, FitMetrics, ForecastComponents, ForecastConfig, ForecastResult, ForecasterStatus,
    ModelType, TimeSeriesData,
};
use tokio::sync::RwLock;
use tracing::{debug, info};

use crate::error::{ForecastError, Result};

/// Statistics tracking for the forecaster
struct ForecasterStats {
    forecasts_generated: AtomicU64,
    data_points_processed: AtomicU64,
}

impl ForecasterStats {
    fn new() -> Self {
        Self {
            forecasts_generated: AtomicU64::new(0),
            data_points_processed: AtomicU64::new(0),
        }
    }
}

/// Wrapper around the different time series model types
enum ModelWrapper {
    Arima(ARIMAModel),
    // Note: SARIMA is handled via ARIMA with seasonal parameters
    Prophet(ProphetModel),
    ExponentialSmoothing(ExponentialSmoothingModel),
    NeuralProphet(NeuralProphetModel),
    Transformer(TransformerForecastModel),
}

impl ModelWrapper {
    fn model_name(&self) -> &'static str {
        match self {
            Self::Arima(_) => "ARIMA",
            Self::Prophet(_) => "Prophet",
            Self::ExponentialSmoothing(_) => "Exponential Smoothing",
            Self::NeuralProphet(_) => "Neural Prophet",
            Self::Transformer(_) => "Transformer",
        }
    }
}

/// Time series forecaster that wraps various model types
pub struct TimeSeriesForecaster {
    model: RwLock<Option<ModelWrapper>>,
    config: RwLock<Option<ForecastConfig>>,
    device: Device,
    stats: ForecasterStats,
    training_data: RwLock<Option<TimeSeriesData>>,
    fit_metrics: RwLock<Option<FitMetrics>>,
}

impl TimeSeriesForecaster {
    /// Create a new forecaster
    #[must_use]
    pub fn new() -> Self {
        let device = select_device(true);
        Self {
            model: RwLock::new(None),
            config: RwLock::new(None),
            device,
            stats: ForecasterStats::new(),
            training_data: RwLock::new(None),
            fit_metrics: RwLock::new(None),
        }
    }

    /// Reset the forecaster to initial state
    pub async fn reset(&self) {
        *self.model.write().await = None;
        *self.config.write().await = None;
        *self.training_data.write().await = None;
        *self.fit_metrics.write().await = None;
    }

    /// Fit a model to the provided data
    pub async fn fit(&self, data: TimeSeriesData, config: ForecastConfig) -> Result<FitMetrics> {
        info!(
            "Fitting {:?} model to {} data points",
            config.model_type,
            data.data.len()
        );

        let start = Instant::now();

        // Convert data to tensors
        let values: Vec<f32> = data.data.iter().map(|p| p.value as f32).collect();
        let timestamps: Vec<f32> = data.data.iter().map(|p| p.timestamp as f32).collect();

        let n = values.len();
        let data_tensor = Tensor::from_vec(values, &[n], &self.device)?;
        let timestamps_tensor = Tensor::from_vec(timestamps, &[n], &self.device)?;

        // Create and fit the model
        let model = self.create_model(&config)?;
        let fit_metrics = self
            .fit_model(model, &data_tensor, &timestamps_tensor, &config)
            .await?;

        // Store training data and fit metrics
        *self.training_data.write().await = Some(data);
        *self.config.write().await = Some(config);
        *self.fit_metrics.write().await = Some(fit_metrics.clone());

        // Update stats
        self.stats
            .data_points_processed
            .fetch_add(n as u64, Ordering::Relaxed);

        let elapsed = start.elapsed().as_secs_f64() * 1000.0;
        debug!("Model fitted in {:.2}ms", elapsed);

        Ok(fit_metrics)
    }

    /// Generate forecasts using the fitted model
    pub async fn forecast(&self) -> Result<ForecastResult> {
        let config_guard = self.config.read().await;
        let config = config_guard.as_ref().ok_or(ForecastError::NotInitialized)?;

        let horizon = config.horizon;
        let confidence_level = config.confidence_level;
        drop(config_guard);

        let model_guard = self.model.read().await;
        let model = model_guard.as_ref().ok_or(ForecastError::NotInitialized)?;

        let start = Instant::now();

        // Generate forecast
        let forecast_output = self
            .generate_forecast(model, horizon, confidence_level)
            .await?;

        // Convert to result format
        let result = self
            .convert_forecast_output(forecast_output, model.model_name(), start)
            .await?;

        // Update stats
        self.stats
            .forecasts_generated
            .fetch_add(1, Ordering::Relaxed);

        Ok(result)
    }

    /// Get the current status
    pub async fn status(&self) -> ForecasterStatus {
        let model_guard = self.model.read().await;
        let initialized = model_guard.is_some();
        let model_name = model_guard
            .as_ref()
            .map(|m: &ModelWrapper| m.model_name().to_string());

        ForecasterStatus {
            initialized,
            model: model_name,
            device: self.device.to_string(),
            forecast_count: self.stats.forecasts_generated.load(Ordering::Relaxed),
            data_points_processed: self.stats.data_points_processed.load(Ordering::Relaxed),
        }
    }

    /// Check if a model is initialized
    pub async fn is_initialized(&self) -> bool {
        self.model.read().await.is_some()
    }

    fn create_model(&self, config: &ForecastConfig) -> Result<ModelWrapper> {
        match config.model_type {
            ModelType::Arima => {
                let default_arima = timeseries_shared::ArimaConfig::default();
                let arima_config = config.arima_config.as_ref().unwrap_or(&default_arima);
                let order = (arima_config.p, arima_config.d, arima_config.q);
                Ok(ModelWrapper::Arima(ARIMAModel::new(order, None)))
            }
            ModelType::Sarima => {
                // SARIMA is implemented via ARIMA with seasonal parameters
                let default_sarima = timeseries_shared::SarimaConfig::default();
                let sarima_config = config.sarima_config.as_ref().unwrap_or(&default_sarima);
                let order = (
                    sarima_config.order.p,
                    sarima_config.order.d,
                    sarima_config.order.q,
                );
                let seasonal = (
                    sarima_config.seasonal_order.p,
                    sarima_config.seasonal_order.d,
                    sarima_config.seasonal_order.q,
                    sarima_config.seasonal_period,
                );
                Ok(ModelWrapper::Arima(ARIMAModel::new(order, Some(seasonal))))
            }
            ModelType::Prophet => {
                let default_prophet = timeseries_shared::ProphetConfig::default();
                let prophet_config = config.prophet_config.as_ref().unwrap_or(&default_prophet);
                let rtx_config = rtx_timeseries::models::ProphetConfig {
                    growth: match prophet_config.growth.as_str() {
                        "logistic" => rtx_timeseries::models::GrowthType::Logistic,
                        _ => rtx_timeseries::models::GrowthType::Linear,
                    },
                    yearly_seasonality: if prophet_config.yearly_seasonality {
                        rtx_timeseries::models::SeasonalityConfig::auto()
                    } else {
                        rtx_timeseries::models::SeasonalityConfig::disabled()
                    },
                    weekly_seasonality: if prophet_config.weekly_seasonality {
                        rtx_timeseries::models::SeasonalityConfig::auto()
                    } else {
                        rtx_timeseries::models::SeasonalityConfig::disabled()
                    },
                    daily_seasonality: if prophet_config.daily_seasonality {
                        rtx_timeseries::models::SeasonalityConfig::auto()
                    } else {
                        rtx_timeseries::models::SeasonalityConfig::disabled()
                    },
                    changepoint_prior_scale: prophet_config.changepoint_prior_scale,
                    ..Default::default()
                };
                Ok(ModelWrapper::Prophet(ProphetModel::with_config(rtx_config)))
            }
            ModelType::ExponentialSmoothing => Ok(ModelWrapper::ExponentialSmoothing(
                ExponentialSmoothingModel::new(
                    rtx_timeseries::models::ExponentialSmoothingConfig::default(),
                ),
            )),
            ModelType::NeuralProphet => Ok(ModelWrapper::NeuralProphet(NeuralProphetModel::new(
                rtx_timeseries::models::NeuralProphetConfig::default(),
            ))),
            ModelType::Transformer => Ok(ModelWrapper::Transformer(TransformerForecastModel::new(
                rtx_timeseries::models::TransformerForecastConfig::default(),
            ))),
        }
    }

    async fn fit_model(
        &self,
        model: ModelWrapper,
        data: &Tensor,
        timestamps: &Tensor,
        _config: &ForecastConfig,
    ) -> Result<FitMetrics> {
        let fit_metrics = match model {
            ModelWrapper::Arima(mut m) => {
                TimeSeriesModel::fit(&mut m, data, timestamps).await?;
                let metrics = m.calculate_fit_metrics(data, timestamps).await?;
                *self.model.write().await = Some(ModelWrapper::Arima(m));
                metrics
            }
            ModelWrapper::Prophet(mut m) => {
                TimeSeriesModel::fit(&mut m, data, timestamps).await?;
                let metrics = m.calculate_fit_metrics(data, timestamps).await?;
                *self.model.write().await = Some(ModelWrapper::Prophet(m));
                metrics
            }
            ModelWrapper::ExponentialSmoothing(mut m) => {
                TimeSeriesModel::fit(&mut m, data, timestamps).await?;
                let metrics = m.calculate_fit_metrics(data, timestamps).await?;
                *self.model.write().await = Some(ModelWrapper::ExponentialSmoothing(m));
                metrics
            }
            ModelWrapper::NeuralProphet(mut m) => {
                TimeSeriesModel::fit(&mut m, data, timestamps).await?;
                let metrics = m.calculate_fit_metrics(data, timestamps).await?;
                *self.model.write().await = Some(ModelWrapper::NeuralProphet(m));
                metrics
            }
            ModelWrapper::Transformer(mut m) => {
                TimeSeriesModel::fit(&mut m, data, timestamps).await?;
                let metrics = m.calculate_fit_metrics(data, timestamps).await?;
                *self.model.write().await = Some(ModelWrapper::Transformer(m));
                metrics
            }
        };

        Ok(convert_fit_metrics(&fit_metrics))
    }

    async fn generate_forecast(
        &self,
        model: &ModelWrapper,
        horizon: usize,
        confidence_level: f64,
    ) -> Result<ForecastOutput> {
        let output = match model {
            ModelWrapper::Arima(m) => m.forecast(horizon, confidence_level).await?,
            ModelWrapper::Prophet(m) => m.forecast(horizon, confidence_level).await?,
            ModelWrapper::ExponentialSmoothing(m) => m.forecast(horizon, confidence_level).await?,
            ModelWrapper::NeuralProphet(m) => m.forecast(horizon, confidence_level).await?,
            ModelWrapper::Transformer(m) => m.forecast(horizon, confidence_level).await?,
        };
        Ok(output)
    }

    async fn convert_forecast_output(
        &self,
        output: ForecastOutput,
        model_name: &str,
        start: Instant,
    ) -> Result<ForecastResult> {
        let training_data = self.training_data.read().await;
        let training_data = training_data
            .as_ref()
            .ok_or(ForecastError::NotInitialized)?;

        let fit_metrics = self.fit_metrics.read().await;
        let metrics = fit_metrics
            .as_ref()
            .ok_or(ForecastError::NotInitialized)?
            .clone();

        // Get the last timestamp from training data
        let last_timestamp = training_data.data.last().map_or(0.0, |p| p.timestamp);

        // Convert tensors to DataPoints
        let mean_vec = output.mean.to_vec()?;
        let lower_vec = output.lower.to_vec()?;
        let upper_vec = output.upper.to_vec()?;

        let forecast: Vec<DataPoint> = mean_vec
            .iter()
            .enumerate()
            .map(|(i, &v)| DataPoint {
                timestamp: last_timestamp + (i + 1) as f64,
                value: f64::from(v),
            })
            .collect();

        let lower_bound: Vec<DataPoint> = lower_vec
            .iter()
            .enumerate()
            .map(|(i, &v)| DataPoint {
                timestamp: last_timestamp + (i + 1) as f64,
                value: f64::from(v),
            })
            .collect();

        let upper_bound: Vec<DataPoint> = upper_vec
            .iter()
            .enumerate()
            .map(|(i, &v)| DataPoint {
                timestamp: last_timestamp + (i + 1) as f64,
                value: f64::from(v),
            })
            .collect();

        // Get fitted values for historical data
        let fitted = training_data.data.clone();

        // Convert components if available
        let components = output.components.map(|c| ForecastComponents {
            trend: c.trend.map(|t| {
                let vec = t.to_vec().unwrap_or_default();
                vec.iter()
                    .enumerate()
                    .map(|(i, &v)| DataPoint {
                        timestamp: i as f64,
                        value: f64::from(v),
                    })
                    .collect()
            }),
            seasonal: c.seasonal.map(|s| {
                let vec = s.to_vec().unwrap_or_default();
                vec.iter()
                    .enumerate()
                    .map(|(i, &v)| DataPoint {
                        timestamp: i as f64,
                        value: f64::from(v),
                    })
                    .collect()
            }),
            residual: c.noise.map(|r| {
                let vec = r.to_vec().unwrap_or_default();
                vec.iter()
                    .enumerate()
                    .map(|(i, &v)| DataPoint {
                        timestamp: i as f64,
                        value: f64::from(v),
                    })
                    .collect()
            }),
        });

        let processing_time_ms = start.elapsed().as_secs_f64() * 1000.0;

        Ok(ForecastResult {
            forecast,
            lower_bound,
            upper_bound,
            fitted,
            components,
            metrics,
            processing_time_ms,
            model: model_name.to_string(),
        })
    }
}

impl Default for TimeSeriesForecaster {
    fn default() -> Self {
        Self::new()
    }
}

fn select_device(use_gpu: bool) -> Device {
    if use_gpu
        && let Ok(device) = Device::try_default()
        && device.is_gpu()
    {
        return device;
    }
    Device::cpu()
}

fn convert_fit_metrics(metrics: &RtxFitMetrics) -> FitMetrics {
    FitMetrics {
        aic: metrics.aic,
        bic: metrics.bic,
        mae: metrics.mae,
        rmse: metrics.rmse,
        mape: metrics.mape,
        r_squared: metrics.r_squared,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use timeseries_shared::ArimaConfig;

    #[tokio::test]
    async fn test_forecaster_creation() {
        let forecaster = TimeSeriesForecaster::new();
        assert!(!forecaster.is_initialized().await);
    }

    #[tokio::test]
    async fn test_forecaster_status() {
        let forecaster = TimeSeriesForecaster::new();
        let status = forecaster.status().await;
        assert!(!status.initialized);
        assert!(status.model.is_none());
    }

    #[tokio::test]
    async fn test_forecast_result_contains_actual_fit_metrics() {
        let forecaster = TimeSeriesForecaster::new();

        let data = create_test_data();
        let config = ForecastConfig {
            model_type: ModelType::Arima,
            horizon: 5,
            confidence_level: 0.95,
            arima_config: Some(ArimaConfig { p: 1, d: 0, q: 1 }),
            sarima_config: None,
            prophet_config: None,
            use_gpu: false,
        };

        let fit_metrics = forecaster
            .fit(data, config)
            .await
            .expect("Failed to fit model");

        assert!(
            fit_metrics.mae > 0.0 || fit_metrics.rmse > 0.0,
            "Fit metrics returned from fit() should contain non-zero values"
        );

        let forecast_result = forecaster
            .forecast()
            .await
            .expect("Failed to generate forecast");

        assert!(
            forecast_result.metrics.mae > 0.0 || forecast_result.metrics.mae == 0.0,
            "Forecast result metrics MAE should be a valid number, got: {}",
            forecast_result.metrics.mae
        );
        assert!(
            forecast_result.metrics.rmse > 0.0 || forecast_result.metrics.rmse == 0.0,
            "Forecast result metrics RMSE should be a valid number, got: {}",
            forecast_result.metrics.rmse
        );

        assert_eq!(
            forecast_result.metrics.mae, fit_metrics.mae,
            "Forecast result MAE should match the fit metrics MAE"
        );
        assert_eq!(
            forecast_result.metrics.rmse, fit_metrics.rmse,
            "Forecast result RMSE should match the fit metrics RMSE"
        );
        assert_eq!(
            forecast_result.metrics.mape, fit_metrics.mape,
            "Forecast result MAPE should match the fit metrics MAPE"
        );
        assert_eq!(
            forecast_result.metrics.aic, fit_metrics.aic,
            "Forecast result AIC should match the fit metrics AIC"
        );
        assert_eq!(
            forecast_result.metrics.bic, fit_metrics.bic,
            "Forecast result BIC should match the fit metrics BIC"
        );
        assert_eq!(
            forecast_result.metrics.r_squared, fit_metrics.r_squared,
            "Forecast result R² should match the fit metrics R²"
        );
    }

    fn create_test_data() -> TimeSeriesData {
        let values = vec![
            100.0, 105.0, 110.0, 108.0, 112.0, 115.0, 120.0, 118.0, 125.0, 130.0, 128.0, 135.0,
            140.0, 138.0, 145.0, 150.0, 148.0, 155.0, 160.0, 158.0, 165.0, 170.0, 168.0, 175.0,
            180.0, 178.0, 185.0, 190.0, 188.0, 195.0,
        ];

        let data = values
            .into_iter()
            .enumerate()
            .map(|(i, value)| DataPoint {
                timestamp: i as f64,
                value,
            })
            .collect();

        TimeSeriesData {
            data,
            name: Some("Test Series".to_string()),
            frequency: Some("daily".to_string()),
        }
    }
}