anofox-forecast 0.4.2

//! Bootstrap methods for forecast uncertainty estimation.
//!
//! Provides residual bootstrap and block bootstrap methods for generating
//! empirical confidence intervals when analytical formulas are unavailable
//! or unreliable.

use crate::core::{Forecast, TimeSeries};
use crate::error::Result;
use crate::models::Forecaster;
use rand::prelude::*;
use rand::SeedableRng;

#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Configuration for bootstrap interval estimation.
#[derive(Debug, Clone)]
pub struct BootstrapConfig {
    /// Number of bootstrap samples to generate.
    pub n_samples: usize,
    /// Block size for block bootstrap (None for residual bootstrap).
    pub block_size: Option<usize>,
    /// Random seed for reproducibility (None for random).
    pub seed: Option<u64>,
}

impl Default for BootstrapConfig {
    fn default() -> Self {
        Self {
            n_samples: 1000,
            block_size: None,
            seed: None,
        }
    }
}

impl BootstrapConfig {
    /// Create a new bootstrap config with specified number of samples.
    pub fn new(n_samples: usize) -> Self {
        Self {
            n_samples,
            ..Default::default()
        }
    }

    /// Use block bootstrap with specified block size.
    /// Preserves autocorrelation structure better than residual bootstrap.
    pub fn with_block_size(mut self, block_size: usize) -> Self {
        self.block_size = Some(block_size);
        self
    }

    /// Set random seed for reproducibility.
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }
}

/// Result of bootstrap interval estimation.
#[derive(Debug, Clone)]
pub struct BootstrapResult {
    /// Lower bounds of confidence intervals (one per horizon step).
    pub lower: Vec<f64>,
    /// Upper bounds of confidence intervals (one per horizon step).
    pub upper: Vec<f64>,
    /// Confidence level used.
    pub level: f64,
    /// Number of bootstrap samples used.
    pub n_samples: usize,
}

/// Resample residuals with replacement (residual bootstrap).
fn resample_residuals(residuals: &[f64], rng: &mut impl Rng) -> Vec<f64> {
    let n = residuals.len();
    (0..n).map(|_| residuals[rng.gen_range(0..n)]).collect()
}

/// Resample using block bootstrap (preserves autocorrelation).
fn resample_blocks(residuals: &[f64], block_size: usize, rng: &mut impl Rng) -> Vec<f64> {
    let n = residuals.len();
    if block_size == 0 || block_size > n {
        return resample_residuals(residuals, rng);
    }

    let mut result = Vec::with_capacity(n);
    let n_blocks = n / block_size + 1;

    for _ in 0..n_blocks {
        let start = rng.gen_range(0..=(n - block_size));
        for j in 0..block_size {
            if result.len() >= n {
                break;
            }
            result.push(residuals[start + j]);
        }
    }

    result.truncate(n);
    result
}

/// Generate bootstrap forecast intervals.
///
/// Uses residual bootstrap: resamples fitted residuals and generates
/// new synthetic series, fits the model, and collects forecast distributions.
///
/// # Arguments
/// * `model` - A fitted forecaster with residuals
/// * `series` - The original time series
/// * `horizon` - Forecast horizon
/// * `level` - Confidence level (e.g., 0.95 for 95% intervals)
/// * `config` - Bootstrap configuration
///
/// # Returns
/// Bootstrap result with lower and upper bounds for each horizon step.
///
/// # Example
/// ```ignore
/// use anofox_forecast::utils::bootstrap::{bootstrap_intervals, BootstrapConfig};
/// use anofox_forecast::models::baseline::Naive;
///
/// let mut model = Naive::new();
/// model.fit(&series).unwrap();
///
/// let config = BootstrapConfig::new(500).with_seed(42);
/// let result = bootstrap_intervals(&model, &series, 10, 0.95, &config).unwrap();
/// ```
pub fn bootstrap_intervals<M>(
    model: &M,
    series: &TimeSeries,
    horizon: usize,
    level: f64,
    config: &BootstrapConfig,
) -> Result<BootstrapResult>
where
    M: Forecaster + Clone + Send + Sync,
{
    let residuals = model
        .residuals()
        .ok_or(crate::error::ForecastError::FitRequired)?;

    let fitted = model
        .fitted_values()
        .ok_or(crate::error::ForecastError::FitRequired)?;

    // Filter out NaN residuals
    let valid_residuals: Vec<f64> = residuals.iter().copied().filter(|r| !r.is_nan()).collect();

    if valid_residuals.is_empty() {
        return Err(crate::error::ForecastError::ComputationError(
            "No valid residuals for bootstrap".to_string(),
        ));
    }

    // Initialize RNG
    let mut rng: StdRng = match config.seed {
        Some(seed) => StdRng::seed_from_u64(seed),
        None => StdRng::from_entropy(),
    };

    // Collect forecast samples for each horizon step
    let forecast_samples = collect_bootstrap_samples(
        model,
        series,
        fitted,
        &valid_residuals,
        horizon,
        config,
        &mut rng,
    );

    // Extract confidence interval bounds from samples
    let (lower, upper) = extract_quantile_bounds(&forecast_samples, level);

    Ok(BootstrapResult {
        lower,
        upper,
        level,
        n_samples: config.n_samples,
    })
}

/// Run a single bootstrap sample: resample, fit, predict.
fn run_bootstrap_sample<M: Forecaster + Clone>(
    model: &M,
    fitted: &[f64],
    original_values: &[f64],
    valid_residuals: &[f64],
    timestamps: &[chrono::DateTime<chrono::Utc>],
    horizon: usize,
    block_size: Option<usize>,
    rng: &mut impl Rng,
) -> Option<Vec<f64>> {
    let resampled = match block_size {
        Some(bs) => resample_blocks(valid_residuals, bs, rng),
        None => resample_residuals(valid_residuals, rng),
    };

    let synthetic_values: Vec<f64> = fitted
        .iter()
        .zip(resampled.iter().cycle())
        .enumerate()
        .map(|(i, (f, r))| {
            let v = f + r;
            if v.is_finite() {
                v
            } else {
                original_values[i]
            }
        })
        .collect();

    let ts = TimeSeries::univariate(timestamps.to_vec(), synthetic_values).ok()?;
    let mut bootstrap_model = model.clone();
    bootstrap_model.fit(&ts).ok()?;
    let forecast = bootstrap_model.predict(horizon).ok()?;
    let values: Vec<f64> = forecast.primary().to_vec();
    if values.iter().all(|v| v.is_finite()) {
        Some(values)
    } else {
        // Filter to only finite values per step in the caller
        Some(values)
    }
}

/// Generate bootstrap forecast samples by resampling residuals and re-fitting.
#[cfg(not(feature = "parallel"))]
fn collect_bootstrap_samples<M: Forecaster + Clone + Send + Sync>(
    model: &M,
    series: &TimeSeries,
    fitted: &[f64],
    valid_residuals: &[f64],
    horizon: usize,
    config: &BootstrapConfig,
    rng: &mut impl Rng,
) -> Vec<Vec<f64>> {
    let original_values = series.primary_values();
    let timestamps = series.timestamps();
    let mut forecast_samples: Vec<Vec<f64>> = vec![Vec::with_capacity(config.n_samples); horizon];

    for _ in 0..config.n_samples {
        if let Some(values) = run_bootstrap_sample(
            model,
            fitted,
            original_values,
            valid_residuals,
            timestamps,
            horizon,
            config.block_size,
            rng,
        ) {
            for (h, &val) in values.iter().enumerate() {
                if val.is_finite() {
                    forecast_samples[h].push(val);
                }
            }
        }
    }

    forecast_samples
}

/// Generate bootstrap forecast samples in parallel.
#[cfg(feature = "parallel")]
fn collect_bootstrap_samples<M: Forecaster + Clone + Send + Sync>(
    model: &M,
    series: &TimeSeries,
    fitted: &[f64],
    valid_residuals: &[f64],
    horizon: usize,
    config: &BootstrapConfig,
    _rng: &mut impl Rng,
) -> Vec<Vec<f64>> {
    let original_values = series.primary_values();
    let timestamps = series.timestamps();
    let base_seed = config.seed.unwrap_or(0);

    // Each sample gets its own deterministic RNG
    let sample_results: Vec<Option<Vec<f64>>> = (0..config.n_samples)
        .into_par_iter()
        .map(|i| {
            let mut rng = StdRng::seed_from_u64(base_seed.wrapping_add(i as u64));
            run_bootstrap_sample(
                model,
                fitted,
                original_values,
                valid_residuals,
                timestamps,
                horizon,
                config.block_size,
                &mut rng,
            )
        })
        .collect();

    // Gather results into per-horizon vectors
    let mut forecast_samples: Vec<Vec<f64>> = vec![Vec::with_capacity(config.n_samples); horizon];
    for values in sample_results.into_iter().flatten() {
        for (h, &val) in values.iter().enumerate() {
            if val.is_finite() {
                forecast_samples[h].push(val);
            }
        }
    }

    forecast_samples
}

/// Extract lower/upper confidence bounds from sorted forecast samples.
fn extract_quantile_bounds(forecast_samples: &[Vec<f64>], level: f64) -> (Vec<f64>, Vec<f64>) {
    let alpha = (1.0 - level) / 2.0;
    let mut lower = Vec::with_capacity(forecast_samples.len());
    let mut upper = Vec::with_capacity(forecast_samples.len());

    for samples in forecast_samples {
        if samples.is_empty() {
            lower.push(f64::NAN);
            upper.push(f64::NAN);
            continue;
        }

        let mut sorted = samples.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let n = sorted.len();

        let lower_idx = ((alpha * n as f64).floor() as usize).min(n - 1);
        let upper_idx = (((1.0 - alpha) * n as f64).floor() as usize).min(n - 1);

        lower.push(sorted[lower_idx]);
        upper.push(sorted[upper_idx]);
    }

    (lower, upper)
}

/// Compute bootstrap forecast with intervals, returning a Forecast object.
///
/// Combines the point forecast from the original model with bootstrap intervals.
pub fn bootstrap_forecast<M>(
    model: &M,
    series: &TimeSeries,
    horizon: usize,
    level: f64,
    config: &BootstrapConfig,
) -> Result<Forecast>
where
    M: Forecaster + Clone + Send + Sync,
{
    let point_forecast = model.predict(horizon)?;
    let bootstrap_result = bootstrap_intervals(model, series, horizon, level, config)?;

    // Combine point forecast with bootstrap intervals
    Ok(Forecast::from_values_with_intervals(
        point_forecast.primary().to_vec(),
        bootstrap_result.lower,
        bootstrap_result.upper,
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::baseline::Naive;
    use crate::models::exponential::SimpleExponentialSmoothing;
    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()
    }

    #[test]
    fn bootstrap_config_default() {
        let config = BootstrapConfig::default();
        assert_eq!(config.n_samples, 1000);
        assert!(config.block_size.is_none());
        assert!(config.seed.is_none());
    }

    #[test]
    fn bootstrap_config_builder() {
        let config = BootstrapConfig::new(500).with_block_size(10).with_seed(42);

        assert_eq!(config.n_samples, 500);
        assert_eq!(config.block_size, Some(10));
        assert_eq!(config.seed, Some(42));
    }

    #[test]
    fn resample_residuals_length() {
        let residuals = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let mut rng = StdRng::seed_from_u64(42);
        let resampled = resample_residuals(&residuals, &mut rng);
        assert_eq!(resampled.len(), residuals.len());
    }

    #[test]
    fn resample_blocks_length() {
        let residuals = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
        let mut rng = StdRng::seed_from_u64(42);
        let resampled = resample_blocks(&residuals, 3, &mut rng);
        assert_eq!(resampled.len(), residuals.len());
    }

    #[test]
    fn bootstrap_intervals_naive() {
        let timestamps = make_timestamps(50);
        let values: Vec<f64> = (0..50).map(|i| 10.0 + (i as f64 * 0.3).sin()).collect();
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = Naive::new();
        model.fit(&ts).unwrap();

        let config = BootstrapConfig::new(100).with_seed(42);
        let result = bootstrap_intervals(&model, &ts, 5, 0.95, &config).unwrap();

        assert_eq!(result.lower.len(), 5);
        assert_eq!(result.upper.len(), 5);
        assert_eq!(result.level, 0.95);

        // Lower should be less than upper
        for i in 0..5 {
            assert!(
                result.lower[i] <= result.upper[i],
                "Lower {} > Upper {} at horizon {}",
                result.lower[i],
                result.upper[i],
                i
            );
        }
    }

    #[test]
    fn bootstrap_intervals_ses() {
        let timestamps = make_timestamps(50);
        let values: Vec<f64> = (0..50).map(|i| 10.0 + (i as f64 * 0.3).sin()).collect();
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = SimpleExponentialSmoothing::auto();
        model.fit(&ts).unwrap();

        let config = BootstrapConfig::new(100).with_seed(123);
        let result = bootstrap_intervals(&model, &ts, 5, 0.90, &config).unwrap();

        assert_eq!(result.lower.len(), 5);
        assert_eq!(result.upper.len(), 5);
        assert_eq!(result.level, 0.90);
    }

    #[test]
    fn bootstrap_forecast_contains_intervals() {
        let timestamps = make_timestamps(50);
        let values: Vec<f64> = (0..50).map(|i| 10.0 + i as f64 * 0.5).collect();
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = Naive::new();
        model.fit(&ts).unwrap();

        let config = BootstrapConfig::new(50).with_seed(42);
        let forecast = bootstrap_forecast(&model, &ts, 5, 0.95, &config).unwrap();

        assert!(forecast.has_lower());
        assert!(forecast.has_upper());
        assert_eq!(forecast.horizon(), 5);
    }

    #[test]
    fn bootstrap_reproducible_with_seed() {
        let timestamps = make_timestamps(50);
        let values: Vec<f64> = (0..50).map(|i| 10.0 + i as f64 * 0.5).collect();
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = Naive::new();
        model.fit(&ts).unwrap();

        let config = BootstrapConfig::new(50).with_seed(42);

        let result1 = bootstrap_intervals(&model, &ts, 5, 0.95, &config).unwrap();
        let result2 = bootstrap_intervals(&model, &ts, 5, 0.95, &config).unwrap();

        for i in 0..5 {
            assert!(
                (result1.lower[i] - result2.lower[i]).abs() < 1e-10,
                "Results should be reproducible with seed"
            );
        }
    }

    #[test]
    fn bootstrap_block_vs_residual() {
        let timestamps = make_timestamps(100);
        let values: Vec<f64> = (0..100)
            .map(|i| 10.0 + (i as f64 * 0.1).sin() * 5.0)
            .collect();
        let ts = TimeSeries::univariate(timestamps, values).unwrap();

        let mut model = Naive::new();
        model.fit(&ts).unwrap();

        let residual_config = BootstrapConfig::new(50).with_seed(42);
        let block_config = BootstrapConfig::new(50).with_block_size(5).with_seed(42);

        let residual_result = bootstrap_intervals(&model, &ts, 5, 0.95, &residual_config).unwrap();
        let block_result = bootstrap_intervals(&model, &ts, 5, 0.95, &block_config).unwrap();

        // Both should produce valid intervals
        assert_eq!(residual_result.lower.len(), 5);
        assert_eq!(block_result.lower.len(), 5);

        // Results may differ due to different resampling strategies
        // Just verify both are valid
        for i in 0..5 {
            assert!(residual_result.lower[i] <= residual_result.upper[i]);
            assert!(block_result.lower[i] <= block_result.upper[i]);
        }
    }
}