rustytorch/crates/models/rtx-timeseries/tests/prophet_tests.rs

//! Comprehensive tests for Prophet models

use approx::assert_abs_diff_eq;
use rtx_tensor::{Device, Tensor};
use rtx_timeseries::{
    TimeSeriesError,
    models::{
        GrowthType, ProphetConfig, ProphetModel, SeasonalityConfig, SeasonalityMode,
        TimeSeriesModel, TypedTimeSeriesModel,
    },
};
use tokio_test;

#[tokio::test]
async fn test_prophet_model_creation() -> Result<(), Box<dyn std::error::Error>> {
    let model = ProphetModel::new();
    assert_eq!(model.get_config().growth, GrowthType::Linear);
    assert_eq!(model.get_config().n_changepoints, 25);
    assert!(model.get_config().yearly_seasonality.enabled);
    assert!(model.is_fitted().is_err());
    Ok(())
}

#[tokio::test]
async fn test_prophet_model_with_custom_config() -> Result<(), Box<dyn std::error::Error>> {
    let config = ProphetConfig {
        growth: GrowthType::Logistic,
        n_changepoints: 15,
        changepoint_range: 0.9,
        changepoint_prior_scale: 0.1,
        yearly_seasonality: SeasonalityConfig::custom(8, 5.0, 365.25),
        weekly_seasonality: SeasonalityConfig::disabled(),
        daily_seasonality: SeasonalityConfig::disabled(),
        holidays_prior_scale: 5.0,
        seasonality_prior_scale: 5.0,
        seasonality_mode: SeasonalityMode::Multiplicative,
        interval_width: 0.9,
        uncertainty_samples: 500,
        quantum_changepoint_detection: true,
        mcmc_samples: 100,
    };

    let model = ProphetModel::with_config(config.clone());
    assert_eq!(model.get_config().growth, GrowthType::Logistic);
    assert_eq!(model.get_config().n_changepoints, 15);
    assert_eq!(
        model.get_config().seasonality_mode,
        SeasonalityMode::Multiplicative
    );
    assert!(model.get_config().quantum_changepoint_detection);
    assert!(!model.get_config().weekly_seasonality.enabled);
    Ok(())
}

#[tokio::test]
async fn test_seasonality_config() -> Result<(), Box<dyn std::error::Error>> {
    let auto_config = SeasonalityConfig::auto();
    assert!(auto_config.enabled);
    assert_eq!(auto_config.fourier_order, 10);
    assert_eq!(auto_config.prior_scale, 10.0);
    assert!(auto_config.period.is_none());

    let disabled_config = SeasonalityConfig::disabled();
    assert!(!disabled_config.enabled);
    assert_eq!(disabled_config.fourier_order, 0);

    let custom_config = SeasonalityConfig::custom(6, 15.0, 30.0);
    assert!(custom_config.enabled);
    assert_eq!(custom_config.fourier_order, 6);
    assert_eq!(custom_config.prior_scale, 15.0);
    assert_eq!(custom_config.period, Some(30.0));
    Ok(())
}

#[tokio::test]
async fn test_changepoint_detector() -> Result<(), Box<dyn std::error::Error>> {
    use rtx_timeseries::models::prophet::ChangepointDetector;

    let device = Device::cpu();
    let timestamps = Tensor::arange(0, 100, &device)?;
    let data = Tensor::ones(&[100], &device)?;

    let detector = ChangepointDetector::new(5, 0.8, 0.05);
    let changepoints = detector
        .detect_changepoints(&timestamps, &data)
        .await
        .unwrap();

    assert_eq!(changepoints.len(), 5);

    // Changepoints should be in chronological order
    for i in 1..changepoints.len() {
        assert!(changepoints[i] > changepoints[i - 1]);
    }

    // All changepoints should be within the range
    let start_time = timestamps.get(&[0]).unwrap() as f64;
    let end_time = timestamps.get(&[79]).unwrap() as f64; // 80% of data (0.8 range)

    for &cp in &changepoints {
        assert!(cp >= start_time);
        assert!(cp <= end_time);
    }
    Ok(())
}

#[tokio::test]
async fn test_changepoint_detector_quantum_enhanced() -> Result<(), Box<dyn std::error::Error>> {
    use rtx_timeseries::models::prophet::ChangepointDetector;

    let device = Device::cpu();
    let timestamps = Tensor::arange(0, 50, &device)?;
    let data = Tensor::ones(&[50], &device)?;

    let detector = ChangepointDetector::new(3, 0.8, 0.05).with_quantum();
    assert!(detector.quantum_enhanced);

    let changepoints = detector
        .detect_changepoints(&timestamps, &data)
        .await
        .unwrap();
    assert_eq!(changepoints.len(), 3);
    Ok(())
}

#[tokio::test]
async fn test_prophet_fitting_simple_trend() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create simple linear trend data
    let data_vec: Vec<f32> = (0..50).map(|i| i as f32 * 0.5 + 10.0).collect();
    let data = Tensor::from_vec(data_vec, &[50], &device)?;
    let timestamps = Tensor::arange(0, 50, &device)?;

    let mut model = ProphetModel::new();
    let result = TimeSeriesModel::fit(&mut model, &data, &timestamps).await;

    assert!(result.is_ok(), "Prophet fitting failed: {:?}", result.err());
    assert!(model.is_fitted().is_ok());

    // Check that components were fitted
    let params = model.get_parameters().unwrap();
    assert!(params.contains_key("growth_rate"));
    assert!(params.contains_key("offset"));
    Ok(())
}

#[tokio::test]
async fn test_prophet_fitting_seasonal_data() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create data with trend and seasonality
    let mut data_vec = Vec::new();
    for i in 0..365 {
        let t = i as f32;
        let trend = 0.01 * t + 100.0;
        let yearly_seasonal = 5.0 * (2.0 * std::f32::consts::PI * t / 365.25).sin();
        let weekly_seasonal = 2.0 * (2.0 * std::f32::consts::PI * t / 7.0).sin();
        let noise = 0.5 * (rand::random::<f32>() - 0.5);
        data_vec.push(trend + yearly_seasonal + weekly_seasonal + noise);
    }

    let data = Tensor::from_vec(data_vec, &[365], &device)?;
    let timestamps = Tensor::arange(0, 365, &device)?;

    let mut model = ProphetModel::new();
    let result = TimeSeriesModel::fit(&mut model, &data, &timestamps).await;

    assert!(result.is_ok());
    assert!(model.is_fitted().is_ok());

    // Verify fit quality
    let fit_metrics = model
        .calculate_fit_metrics(&data, &timestamps)
        .await
        .unwrap();
    assert!(fit_metrics.r_squared > 0.0); // Should capture some variance
    Ok(())
}

#[tokio::test]
async fn test_prophet_forecasting() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create predictable trend data
    let data_vec: Vec<f32> = (1..=30).map(|i| i as f32 * 0.2 + 10.0).collect();
    let data = Tensor::from_vec(data_vec, &[30], &device)?;
    let timestamps = Tensor::arange(1, 31, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    // Generate forecasts
    let forecast = model.forecast(7, 0.95).await.unwrap();

    assert_eq!(forecast.len(), 7);
    assert_eq!(forecast.mean.shape()[0], 7);
    assert_eq!(forecast.lower.shape()[0], 7);
    assert_eq!(forecast.upper.shape()[0], 7);

    // Check that confidence intervals are reasonable
    for i in 0..7 {
        let mean_val = forecast.mean.get(&[i]).unwrap();
        let lower_val = forecast.lower.get(&[i]).unwrap();
        let upper_val = forecast.upper.get(&[i]).unwrap();

        assert!(lower_val < mean_val);
        assert!(mean_val < upper_val);
        assert!(upper_val - lower_val > 0.0); // Non-zero uncertainty
    }

    // Check forecast components
    assert!(forecast.components.is_some());
    let components = forecast.components.unwrap();
    assert!(components.trend.is_some());
    Ok(())
}

#[tokio::test]
async fn test_prophet_forecasting_different_growth_types() -> Result<(), Box<dyn std::error::Error>>
{
    let device = Device::cpu();

    let data = Tensor::from_vec(
        vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
        &[10],
        &device,
    )?;
    let timestamps = Tensor::arange(1, 11, &device)?;

    // Test linear growth
    let config_linear = ProphetConfig {
        growth: GrowthType::Linear,
        ..Default::default()
    };
    let mut model_linear = ProphetModel::with_config(config_linear);
    model_linear.fit(&data, &timestamps).await.unwrap();
    let forecast_linear = model_linear.forecast(3, 0.95).await.unwrap();

    // Test logistic growth
    let config_logistic = ProphetConfig {
        growth: GrowthType::Logistic,
        ..Default::default()
    };
    let mut model_logistic = ProphetModel::with_config(config_logistic);
    model_logistic.fit(&data, &timestamps).await.unwrap();
    let forecast_logistic = model_logistic.forecast(3, 0.95).await.unwrap();

    // Test flat growth
    let config_flat = ProphetConfig {
        growth: GrowthType::Flat,
        ..Default::default()
    };
    let mut model_flat = ProphetModel::with_config(config_flat);
    model_flat.fit(&data, &timestamps).await.unwrap();
    let forecast_flat = model_flat.forecast(3, 0.95).await.unwrap();

    // All should produce valid forecasts
    assert_eq!(forecast_linear.len(), 3);
    assert_eq!(forecast_logistic.len(), 3);
    assert_eq!(forecast_flat.len(), 3);

    // Flat growth should have approximately constant forecasts
    let flat_mean_1 = forecast_flat.mean.get(&[0]).unwrap();
    let flat_mean_3 = forecast_flat.mean.get(&[2]).unwrap();
    assert_abs_diff_eq!(flat_mean_1, flat_mean_3, epsilon = 0.1);
    Ok(())
}

#[tokio::test]
async fn test_prophet_seasonality_modes() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create data with multiplicative seasonality
    let mut data_vec = Vec::new();
    for i in 0..100 {
        let t = i as f32;
        let trend = 10.0 + 0.1 * t;
        let seasonal_factor = 1.0 + 0.2 * (2.0 * std::f32::consts::PI * t / 12.0).sin();
        data_vec.push(trend * seasonal_factor);
    }

    let data = Tensor::from_vec(data_vec, &[100], &device)?;
    let timestamps = Tensor::arange(0, 100, &device)?;

    // Test additive seasonality
    let config_additive = ProphetConfig {
        seasonality_mode: SeasonalityMode::Additive,
        ..Default::default()
    };
    let mut model_additive = ProphetModel::with_config(config_additive);
    let result_additive = model_additive.fit(&data, &timestamps).await;
    assert!(result_additive.is_ok());

    // Test multiplicative seasonality
    let config_multiplicative = ProphetConfig {
        seasonality_mode: SeasonalityMode::Multiplicative,
        ..Default::default()
    };
    let mut model_multiplicative = ProphetModel::with_config(config_multiplicative);
    let result_multiplicative = model_multiplicative.fit(&data, &timestamps).await;
    assert!(result_multiplicative.is_ok());

    // Both should fit successfully
    assert!(model_additive.is_fitted().is_ok());
    assert!(model_multiplicative.is_fitted().is_ok());
    Ok(())
}

#[tokio::test]
async fn test_prophet_holiday_effects() -> Result<(), Box<dyn std::error::Error>> {
    use chrono::{DateTime, Utc};
    use rtx_timeseries::models::prophet::HolidayEffects;

    let device = Device::cpu();
    let data = Tensor::from_vec(
        (0..100).map(|i| 10.0 + i as f32 * 0.1).collect(),
        &[100],
        &device,
    )?;
    let timestamps = Tensor::arange(0, 100, &device)?;

    let mut model = ProphetModel::new();

    // Add some holiday effects
    let holidays = HolidayEffects {
        dates: vec![
            DateTime::<Utc>::from_timestamp(25 * 86400, 0).unwrap(),
            DateTime::<Utc>::from_timestamp(50 * 86400, 0).unwrap(),
            DateTime::<Utc>::from_timestamp(75 * 86400, 0).unwrap(),
        ],
        effects: vec![5.0, -3.0, 2.0],
        names: vec![
            "Holiday1".to_string(),
            "Holiday2".to_string(),
            "Holiday3".to_string(),
        ],
    };

    model.add_holidays(holidays);

    let result = TimeSeriesModel::fit(&mut model, &data, &timestamps).await;
    assert!(result.is_ok());
    Ok(())
}

#[tokio::test]
async fn test_prophet_residuals() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    let data = Tensor::from_vec(
        vec![10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5],
        &[10],
        &device,
    )?;
    let timestamps = Tensor::arange(0, 10, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    let residuals = model.residuals(&data, &timestamps).await.unwrap();

    assert_eq!(residuals.shape()[0], 10);

    // Check that residuals are reasonable for this linear data
    let residuals_vec: Vec<f32> = (0..10).map(|i| residuals.get(&[i]).unwrap()).collect();

    let mean_abs_residual = residuals_vec.iter().map(|r| r.abs()).sum::<f32>() / 10.0;
    assert!(mean_abs_residual < 1.0); // Should be small for linear trend
    Ok(())
}

#[tokio::test]
async fn test_prophet_fit_metrics() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create data with trend and some noise
    let data_vec: Vec<f32> = (0..50)
        .map(|i| i as f32 * 0.3 + 10.0 + 0.2 * (rand::random::<f32>() - 0.5))
        .collect();
    let data = Tensor::from_vec(data_vec, &[50], &device)?;
    let timestamps = Tensor::arange(0, 50, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    let metrics = model
        .calculate_fit_metrics(&data, &timestamps)
        .await
        .unwrap();

    // Check that metrics are reasonable
    assert!(metrics.aic.is_finite());
    assert!(metrics.bic.is_finite());
    assert!(metrics.r_squared >= 0.0 && metrics.r_squared <= 1.0);
    assert!(metrics.rmse >= 0.0);
    assert!(metrics.mae >= 0.0);
    assert!(metrics.mape >= 0.0);

    // For trend data, should have decent R²
    assert!(metrics.r_squared > 0.5);
    Ok(())
}

#[tokio::test]
async fn test_prophet_parameter_getting() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    let data = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0], &[5], &device)?;
    let timestamps = Tensor::arange(0, 5, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    // Test parameter getting
    let params = model.get_parameters().unwrap();
    assert!(params.contains_key("growth_rate"));
    assert!(params.contains_key("offset"));

    // Check that parameters are reasonable
    let growth_rate = params["growth_rate"];
    assert!(growth_rate.is_finite());

    let offset = params["offset"];
    assert!(offset.is_finite());
    Ok(())
}

#[tokio::test]
async fn test_prophet_model_cloning() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    let data = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0], &[5], &device)?;
    let timestamps = Tensor::arange(0, 5, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    // Test model cloning
    let cloned_model = model.clone_boxed();
    assert!(cloned_model.is_ok());

    let cloned = cloned_model.unwrap();
    assert!(cloned.is_fitted().is_ok());

    // Verify parameters match
    let original_params = model.get_parameters().unwrap();
    let cloned_params = cloned.get_parameters().unwrap();

    assert_eq!(original_params.len(), cloned_params.len());
    for (key, value) in original_params {
        assert_abs_diff_eq!(cloned_params[&key], value, epsilon = 1e-6);
    }
    Ok(())
}

#[tokio::test]
async fn test_prophet_validation_errors() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Test mismatched data and timestamps
    let data = Tensor::from_vec(vec![1.0, 2.0, 3.0], &[3], &device)?;
    let timestamps = Tensor::from_vec(vec![1.0, 2.0], &[2], &device)?; // Wrong length

    let mut model = ProphetModel::new();
    let result = TimeSeriesModel::fit(&mut model, &data, &timestamps).await;

    assert!(result.is_err());
    match result.err().unwrap() {
        TimeSeriesError::ValidationError(_) => {} // Expected
        other => panic!("Expected ValidationError, got {:?}", other),
    }
    Ok(())
}

#[tokio::test]
async fn test_prophet_insufficient_data() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Test with insufficient data points
    let data = Tensor::from_vec(vec![1.0], &[1], &device)?; // Only 1 point
    let timestamps = Tensor::from_vec(vec![1.0], &[1], &device)?;

    let mut model = ProphetModel::new();
    let result = TimeSeriesModel::fit(&mut model, &data, &timestamps).await;

    assert!(result.is_err());
    match result.err().unwrap() {
        TimeSeriesError::ValidationError(_) => {} // Expected
        other => panic!("Expected ValidationError, got {:?}", other),
    }
    Ok(())
}

#[tokio::test]
async fn test_prophet_forecasting_unfitted_model() -> Result<(), Box<dyn std::error::Error>> {
    let model = ProphetModel::new();

    let result = model.forecast(5, 0.95).await;
    assert!(result.is_err());
    match result.err().unwrap() {
        TimeSeriesError::ModelStateError(_) => {} // Expected
        other => panic!("Expected ModelStateError, got {:?}", other),
    }
    Ok(())
}

#[tokio::test]
async fn test_prophet_mcmc_sampler() -> Result<(), Box<dyn std::error::Error>> {
    use rtx_timeseries::models::prophet::MCMCSampler;

    let sampler = MCMCSampler::new(1000);
    assert_eq!(sampler.n_samples, 1000);
    assert_eq!(sampler.n_warmup, 500); // Half of samples
    assert_eq!(sampler.adapt_delta, 0.8);
    assert_eq!(sampler.max_treedepth, 10);
    Ok(())
}

#[tokio::test]
async fn test_prophet_fourier_seasonality_fitting() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create data with known seasonal pattern
    let mut data_vec = Vec::new();
    for i in 0..365 {
        let t = i as f32;
        let seasonal = 3.0 * (2.0 * std::f32::consts::PI * t / 365.25).sin();
        data_vec.push(10.0 + seasonal);
    }

    let data = Tensor::from_vec(data_vec, &[365], &device)?;
    let timestamps = Tensor::arange(0, 365, &device)?;

    let mut model = ProphetModel::new();
    let result = model
        .fit_fourier_seasonality(&timestamps, &data, 365.25, 5)
        .await;

    assert!(result.is_ok());
    let seasonal_component = result.unwrap();

    assert_eq!(seasonal_component.fourier_order, 5);
    assert_eq!(seasonal_component.period, 365.25);
    assert_eq!(seasonal_component.cosine_coeffs.len(), 5);
    assert_eq!(seasonal_component.sine_coeffs.len(), 5);
    Ok(())
}

#[tokio::test]
async fn test_prophet_trend_estimation() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Create data with known linear trend
    let data_vec: Vec<f32> = (0..50).map(|i| 5.0 + i as f32 * 0.3).collect();
    let data = Tensor::from_vec(data_vec, &[50], &device)?;
    let timestamps = Tensor::arange(0, 50, &device)?;

    let mut model = ProphetModel::new();
    let growth_rate = model
        .estimate_growth_rate(&timestamps, &data)
        .await
        .unwrap();

    // Should be close to 0.3 (the true growth rate)
    assert_abs_diff_eq!(growth_rate, 0.3, epsilon = 0.1);
    Ok(())
}

#[tokio::test]
async fn test_prophet_trend_generation() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    let data = Tensor::from_vec(vec![10.0, 11.0, 12.0, 13.0, 14.0], &[5], &device)?;
    let timestamps = Tensor::arange(0, 5, &device)?;

    let mut model = ProphetModel::new();
    TimeSeriesModel::fit(&mut model, &data, &timestamps)
        .await
        .unwrap();

    // Generate trend for same timestamps
    let trend = model.generate_trend(&timestamps).await.unwrap();

    assert_eq!(trend.shape()[0], 5);

    // Trend should be approximately linear for this data
    let trend_vec: Vec<f32> = (0..5).map(|i| trend.get(&[i]).unwrap()).collect();

    // Check that trend is generally increasing (allowing for some estimation error)
    for i in 1..5 {
        // Allow some tolerance for estimation differences
        assert!(trend_vec[i] >= trend_vec[i - 1] - 1.0);
    }
    Ok(())
}