anofox-forecast 0.15.9

//! Stationarity tests for time series.
//!
//! Provides tests to determine if a time series is stationary.

/// Result of a stationarity test.
#[derive(Debug, Clone)]
pub struct StationarityResult {
    /// Test statistic
    pub statistic: f64,
    /// P-value (approximate)
    pub p_value: f64,
    /// Number of lags used
    pub lags: usize,
    /// Whether series appears stationary
    pub is_stationary: bool,
    /// Critical values at common significance levels
    pub critical_values: CriticalValues,
}

/// Critical values for stationarity tests.
#[derive(Debug, Clone, Default)]
pub struct CriticalValues {
    /// Critical value at 1% significance
    pub cv_1pct: f64,
    /// Critical value at 5% significance
    pub cv_5pct: f64,
    /// Critical value at 10% significance
    pub cv_10pct: f64,
}

/// Augmented Dickey-Fuller test for unit root (non-stationarity).
///
/// Tests null hypothesis that series has a unit root (non-stationary).
/// Rejection implies stationarity.
///
/// # Arguments
/// * `series` - Time series data
/// * `max_lags` - Maximum lags to include (default: (n-1)^(1/3))
///
/// # Returns
/// `StationarityResult` with test statistic and p-value
pub fn adf_test(series: &[f64], max_lags: Option<usize>) -> StationarityResult {
    let n = series.len();

    if n < 4 {
        return StationarityResult {
            statistic: f64::NAN,
            p_value: f64::NAN,
            lags: 0,
            is_stationary: false,
            critical_values: CriticalValues::default(),
        };
    }

    // Default lag selection: (n-1)^(1/3)
    let max_lags = max_lags.unwrap_or_else(|| ((n - 1) as f64).powf(1.0 / 3.0).floor() as usize);
    let max_lags = max_lags.min(n / 2 - 1).max(1);

    // First difference
    let diff: Vec<f64> = series.windows(2).map(|w| w[1] - w[0]).collect();

    // Select optimal lag using AIC
    let (best_lag, _) = select_lag_aic(&diff, &series[..n - 1], max_lags);

    // Run ADF regression: Δy_t = α + β*y_{t-1} + Σγ_i*Δy_{t-i} + ε_t
    let (stat, se) = compute_adf_statistic(&diff, &series[..n - 1], best_lag);

    if se == 0.0 || se.is_nan() {
        return StationarityResult {
            statistic: f64::NAN,
            p_value: f64::NAN,
            lags: best_lag,
            is_stationary: false,
            critical_values: CriticalValues::default(),
        };
    }

    let t_stat = stat / se;

    // Critical values for ADF with constant (MacKinnon approximation)
    let critical_values = CriticalValues {
        cv_1pct: -3.43,
        cv_5pct: -2.86,
        cv_10pct: -2.57,
    };

    // Approximate p-value using MacKinnon tables
    let p_value = adf_p_value(t_stat, n);

    // Series is stationary if we reject null (t_stat < critical value)
    let is_stationary = t_stat < critical_values.cv_5pct;

    StationarityResult {
        statistic: t_stat,
        p_value,
        lags: best_lag,
        is_stationary,
        critical_values,
    }
}

/// Select lag order using AIC.
fn select_lag_aic(diff: &[f64], level: &[f64], max_lags: usize) -> (usize, f64) {
    let mut best_lag = 1;
    let mut best_aic = f64::INFINITY;

    for lag in 1..=max_lags {
        let aic = compute_aic(diff, level, lag);
        if aic < best_aic {
            best_aic = aic;
            best_lag = lag;
        }
    }

    (best_lag, best_aic)
}

/// Compute AIC for a given lag order.
fn compute_aic(diff: &[f64], level: &[f64], lag: usize) -> f64 {
    let n = diff.len();
    if n <= lag + 1 {
        return f64::INFINITY;
    }

    let start = lag;
    let effective_n = n - start;

    if effective_n < 3 {
        return f64::INFINITY;
    }

    // Build design matrix and compute residual sum of squares
    let rss = compute_rss(diff, level, lag);

    if rss <= 0.0 {
        return f64::INFINITY;
    }

    // AIC = n * ln(RSS/n) + 2 * k
    let k = lag + 2; // intercept + level coefficient + lag coefficients
    effective_n as f64 * (rss / effective_n as f64).ln() + 2.0 * k as f64
}

/// Compute residual sum of squares for ADF regression.
fn compute_rss(diff: &[f64], level: &[f64], lag: usize) -> f64 {
    let n = diff.len();
    let start = lag;

    if n <= start + 1 || level.len() <= start {
        return f64::INFINITY;
    }

    // Simple OLS for: Δy_t = α + β*y_{t-1} + Σγ_i*Δy_{t-i}
    let effective_n = n - start;

    // Compute means
    let y_mean: f64 = diff[start..].iter().sum::<f64>() / effective_n as f64;
    let x_mean: f64 = level[start..n].iter().sum::<f64>() / effective_n as f64;

    // Simple regression of diff on level (ignoring lags for AIC simplicity)
    let mut xx = 0.0;
    let mut xy = 0.0;

    for i in start..n {
        let x = level[i] - x_mean;
        let y = diff[i] - y_mean;
        xx += x * x;
        xy += x * y;
    }

    if xx == 0.0 {
        return f64::INFINITY;
    }

    let beta = xy / xx;
    let alpha = y_mean - beta * x_mean;

    // Compute RSS
    let mut rss = 0.0;
    for i in start..n {
        let predicted = alpha + beta * level[i];
        let residual = diff[i] - predicted;
        rss += residual * residual;
    }

    rss
}

/// Compute ADF test statistic and standard error.
fn compute_adf_statistic(diff: &[f64], level: &[f64], lag: usize) -> (f64, f64) {
    let n = diff.len();
    let start = lag;

    if n <= start + 2 || level.len() <= start {
        return (f64::NAN, f64::NAN);
    }

    let effective_n = n - start;

    // Build augmented regression
    // y = diff[start..], X = [1, level[start..], diff_lags]

    // For simplicity, just compute the coefficient on level[t-1]
    let y_mean: f64 = diff[start..].iter().sum::<f64>() / effective_n as f64;
    let x_mean: f64 = level[start..n].iter().sum::<f64>() / effective_n as f64;

    let mut xx = 0.0;
    let mut xy = 0.0;
    let mut yy = 0.0;

    for i in start..n {
        let x = level[i] - x_mean;
        let y = diff[i] - y_mean;
        xx += x * x;
        xy += x * y;
        yy += y * y;
    }

    if xx == 0.0 {
        return (f64::NAN, f64::NAN);
    }

    let beta = xy / xx;

    // Compute residual variance
    let rss = yy - beta * xy;
    let sigma_sq = rss / (effective_n - 2) as f64;

    if sigma_sq <= 0.0 {
        return (f64::NAN, f64::NAN);
    }

    // Standard error of beta
    let se_beta = (sigma_sq / xx).sqrt();

    (beta, se_beta)
}

/// Approximate p-value for ADF test using MacKinnon regression.
fn adf_p_value(t_stat: f64, _n: usize) -> f64 {
    if t_stat.is_nan() {
        return f64::NAN;
    }

    // Threshold/p-value lookup: (upper_bound, p_value)
    // Sorted ascending by threshold for linear scan
    const ADF_TABLE: [(f64, f64); 9] = [
        (-4.00, 0.001),
        (-3.43, 0.01),
        (-2.86, 0.05),
        (-2.57, 0.10),
        (-1.94, 0.20),
        (-1.62, 0.30),
        (-1.28, 0.40),
        (-0.84, 0.50),
        (0.00, 0.70),
    ];

    for &(threshold, p) in &ADF_TABLE {
        if t_stat < threshold {
            return p;
        }
    }

    0.90 + 0.05 * (1.0 - (-t_stat).exp())
}

/// KPSS test for stationarity.
///
/// Tests null hypothesis that series is (trend) stationary.
/// Rejection implies non-stationarity.
///
/// # Arguments
/// * `series` - Time series data
/// * `lags` - Number of lags for HAC variance (default: 4*(n/100)^0.25)
///
/// # Returns
/// `StationarityResult` with test statistic and p-value
pub fn kpss_test(series: &[f64], lags: Option<usize>) -> StationarityResult {
    let n = series.len();

    if n < 4 {
        return StationarityResult {
            statistic: f64::NAN,
            p_value: f64::NAN,
            lags: 0,
            is_stationary: false,
            critical_values: CriticalValues::default(),
        };
    }

    // Default lag: 4 * (n/100)^0.25
    let lags = lags.unwrap_or_else(|| (4.0 * (n as f64 / 100.0).powf(0.25)).floor() as usize);
    let lags = lags.min(n / 2).max(1);

    // Demean the series (level stationarity)
    let mean: f64 = series.iter().sum::<f64>() / n as f64;
    let residuals: Vec<f64> = series.iter().map(|&x| x - mean).collect();

    // Compute cumulative sum of residuals
    let mut cumsum = vec![0.0; n];
    cumsum[0] = residuals[0];
    for i in 1..n {
        cumsum[i] = cumsum[i - 1] + residuals[i];
    }

    // Compute numerator: sum of squared cumulative sums
    let numerator: f64 = cumsum.iter().map(|&s| s * s).sum::<f64>() / (n * n) as f64;

    // Compute HAC variance estimator (Bartlett kernel)
    let mut variance = residuals.iter().map(|&r| r * r).sum::<f64>() / n as f64;

    for j in 1..=lags {
        let weight = 1.0 - j as f64 / (lags + 1) as f64;
        let autocovar: f64 = residuals
            .iter()
            .skip(j)
            .zip(residuals.iter())
            .map(|(&a, &b)| a * b)
            .sum::<f64>()
            / n as f64;
        variance += 2.0 * weight * autocovar;
    }

    if variance <= 0.0 {
        return StationarityResult {
            statistic: f64::NAN,
            p_value: f64::NAN,
            lags,
            is_stationary: true,
            critical_values: CriticalValues::default(),
        };
    }

    let stat = numerator / variance;

    // Critical values for KPSS level stationarity
    let critical_values = CriticalValues {
        cv_1pct: 0.739,
        cv_5pct: 0.463,
        cv_10pct: 0.347,
    };

    // Approximate p-value
    let p_value = kpss_p_value(stat);

    // Series is stationary if we fail to reject null (stat < critical value)
    let is_stationary = stat < critical_values.cv_5pct;

    StationarityResult {
        statistic: stat,
        p_value,
        lags,
        is_stationary,
        critical_values,
    }
}

/// Approximate p-value for KPSS test.
fn kpss_p_value(stat: f64) -> f64 {
    if stat.is_nan() {
        return f64::NAN;
    }

    // Simplified approximation based on critical values
    if stat < 0.347 {
        0.10 + 0.90 * (1.0 - stat / 0.347)
    } else if stat < 0.463 {
        0.05 + 0.05 * (0.463 - stat) / (0.463 - 0.347)
    } else if stat < 0.739 {
        0.01 + 0.04 * (0.739 - stat) / (0.739 - 0.463)
    } else {
        0.01 * (1.0 - (stat - 0.739).min(1.0))
    }
}

/// Combined stationarity test using both ADF and KPSS.
///
/// # Returns
/// A tuple of (adf_result, kpss_result, conclusion)
/// where conclusion is:
/// - "stationary" if ADF rejects AND KPSS fails to reject
/// - "non_stationary" if ADF fails to reject AND KPSS rejects
/// - "inconclusive" otherwise
pub fn test_stationarity(series: &[f64]) -> (StationarityResult, StationarityResult, &'static str) {
    let adf = adf_test(series, None);
    let kpss = kpss_test(series, None);

    let conclusion = if adf.is_stationary && kpss.is_stationary {
        "stationary"
    } else if !adf.is_stationary && !kpss.is_stationary {
        "non_stationary"
    } else {
        "inconclusive"
    };

    (adf, kpss, conclusion)
}

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

    // ==================== adf_test ====================

    #[test]
    fn adf_stationary_series() {
        // White noise should be stationary
        let series: Vec<f64> = (0..200)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = adf_test(&series, Some(5));

        assert!(!result.statistic.is_nan());
        assert!(result.statistic < 0.0); // Should be negative
    }

    #[test]
    fn adf_random_walk() {
        // Random walk should be non-stationary
        let mut series = vec![0.0; 200];
        for i in 1..200 {
            series[i] = series[i - 1] + ((i * 17) % 19) as f64 / 10.0 - 0.9;
        }

        let result = adf_test(&series, Some(5));

        assert!(!result.statistic.is_nan());
        // Just verify we get a valid result
        assert!(result.p_value >= 0.0 && result.p_value <= 1.0);
    }

    #[test]
    fn adf_trending_series() {
        // Series with strong trend (add small noise to avoid numerical issues)
        let series: Vec<f64> = (0..200)
            .map(|i| i as f64 * 0.5 + ((i * 13) % 7) as f64 * 0.01)
            .collect();

        let result = adf_test(&series, Some(5));

        assert!(!result.statistic.is_nan());
        // Trending series should fail to reject null
        assert!(!result.is_stationary);
    }

    #[test]
    fn adf_short_series() {
        let series = vec![1.0, 2.0, 3.0];
        let result = adf_test(&series, Some(1));

        assert!(result.statistic.is_nan());
    }

    #[test]
    fn adf_empty() {
        let result = adf_test(&[], None);
        assert!(result.statistic.is_nan());
    }

    #[test]
    fn adf_critical_values() {
        let series: Vec<f64> = (0..100)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = adf_test(&series, None);

        assert!(result.critical_values.cv_1pct < result.critical_values.cv_5pct);
        assert!(result.critical_values.cv_5pct < result.critical_values.cv_10pct);
    }

    // ==================== kpss_test ====================

    #[test]
    fn kpss_stationary_series() {
        // White noise should be stationary
        let series: Vec<f64> = (0..200)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, Some(10));

        assert!(!result.statistic.is_nan());
        assert!(result.statistic > 0.0);
        // Should fail to reject null (stationary)
        assert!(result.is_stationary);
    }

    #[test]
    fn kpss_trending_series() {
        // Series with strong trend should reject stationarity
        let series: Vec<f64> = (0..200).map(|i| i as f64 * 0.5).collect();

        let result = kpss_test(&series, Some(10));

        assert!(!result.statistic.is_nan());
        // Should reject null (not stationary)
        assert!(!result.is_stationary);
    }

    #[test]
    fn kpss_random_walk() {
        // Random walk
        let mut series = vec![0.0; 200];
        for i in 1..200 {
            series[i] = series[i - 1] + ((i * 17) % 19) as f64 / 10.0 - 0.9;
        }

        let result = kpss_test(&series, Some(10));

        assert!(!result.statistic.is_nan());
    }

    #[test]
    fn kpss_short_series() {
        let series = vec![1.0, 2.0, 3.0];
        let result = kpss_test(&series, Some(1));

        assert!(result.statistic.is_nan());
    }

    #[test]
    fn kpss_empty() {
        let result = kpss_test(&[], None);
        assert!(result.statistic.is_nan());
    }

    #[test]
    fn kpss_critical_values() {
        let series: Vec<f64> = (0..100)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, None);

        // KPSS critical values should increase
        assert!(result.critical_values.cv_10pct < result.critical_values.cv_5pct);
        assert!(result.critical_values.cv_5pct < result.critical_values.cv_1pct);
    }

    // ==================== test_stationarity ====================

    #[test]
    fn combined_test_stationary() {
        // White noise should be conclusively stationary
        let series: Vec<f64> = (0..200)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let (adf, kpss, conclusion) = test_stationarity(&series);

        assert!(!adf.statistic.is_nan());
        assert!(!kpss.statistic.is_nan());
        // Conclusion depends on specific values
        assert!(conclusion == "stationary" || conclusion == "inconclusive");
    }

    #[test]
    fn combined_test_trending() {
        // Strong trend should be non-stationary (add small noise to avoid numerical issues)
        let series: Vec<f64> = (0..200)
            .map(|i| i as f64 * 0.5 + ((i * 13) % 7) as f64 * 0.01)
            .collect();

        let (adf, kpss, conclusion) = test_stationarity(&series);

        assert!(!adf.statistic.is_nan());
        assert!(!kpss.statistic.is_nan());
        // Should be non-stationary or inconclusive
        assert!(conclusion == "non_stationary" || conclusion == "inconclusive");
    }

    #[test]
    fn combined_test_short() {
        let series = vec![1.0, 2.0, 3.0];

        let (adf, kpss, _) = test_stationarity(&series);

        assert!(adf.statistic.is_nan());
        assert!(kpss.statistic.is_nan());
    }

    // ==================== adf_test edge cases ====================

    #[test]
    fn adf_constant_series() {
        // Constant series has zero variance in differences
        let series = vec![5.0; 100];
        let result = adf_test(&series, Some(1));

        // Constant series -> all diffs zero -> NaN or degenerate
        assert!(result.statistic.is_nan() || result.statistic.is_finite());
    }

    #[test]
    fn adf_single_value() {
        let result = adf_test(&[42.0], None);
        assert!(result.statistic.is_nan());
        assert!(!result.is_stationary);
    }

    #[test]
    fn adf_two_values() {
        let result = adf_test(&[1.0, 2.0], None);
        assert!(result.statistic.is_nan());
    }

    #[test]
    fn adf_four_values_minimum() {
        // Exactly 4 values is the minimum for non-NaN
        let series = vec![1.0, 3.0, 2.0, 4.0];
        let result = adf_test(&series, Some(1));
        // Should produce some result (may or may not be NaN depending on regression)
        assert!(result.lags >= 1);
    }

    #[test]
    fn adf_large_lag_clamped() {
        let series: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let result = adf_test(&series, Some(100));
        // max_lags should be clamped to n/2 - 1
        assert!(result.lags <= 9);
    }

    #[test]
    fn adf_p_value_range() {
        let series: Vec<f64> = (0..200)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = adf_test(&series, None);
        // p-value should be in valid range
        if !result.p_value.is_nan() {
            assert!(result.p_value >= 0.0 && result.p_value <= 1.0);
        }
    }

    #[test]
    fn adf_critical_values_ordering() {
        // 1% CV < 5% CV < 10% CV (all negative, becoming less negative)
        let series: Vec<f64> = (0..100)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = adf_test(&series, None);
        assert!(result.critical_values.cv_1pct < result.critical_values.cv_5pct);
        assert!(result.critical_values.cv_5pct < result.critical_values.cv_10pct);
        // All should be negative for ADF
        assert!(result.critical_values.cv_1pct < 0.0);
        assert!(result.critical_values.cv_5pct < 0.0);
        assert!(result.critical_values.cv_10pct < 0.0);
    }

    // ==================== kpss_test edge cases ====================

    #[test]
    fn kpss_constant_series() {
        let series = vec![5.0; 100];
        let result = kpss_test(&series, Some(5));

        // Constant series has zero variance -> NaN or degenerate statistic
        // After demeaning, residuals are all zero
        assert!(result.statistic.is_nan() || result.statistic == 0.0 || result.is_stationary);
    }

    #[test]
    fn kpss_single_value() {
        let result = kpss_test(&[42.0], None);
        assert!(result.statistic.is_nan());
    }

    #[test]
    fn kpss_two_values() {
        let result = kpss_test(&[1.0, 2.0], None);
        assert!(result.statistic.is_nan());
    }

    #[test]
    fn kpss_lag_selection_default() {
        let series: Vec<f64> = (0..100)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, None);
        // Default lag: 4 * (100/100)^0.25 = 4
        assert!(result.lags >= 1);
    }

    #[test]
    fn kpss_critical_values_ordering() {
        let series: Vec<f64> = (0..100)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, None);
        // KPSS critical values should be positive and increasing
        assert!(result.critical_values.cv_10pct > 0.0);
        assert!(result.critical_values.cv_10pct < result.critical_values.cv_5pct);
        assert!(result.critical_values.cv_5pct < result.critical_values.cv_1pct);
    }

    #[test]
    fn kpss_p_value_range() {
        let series: Vec<f64> = (0..200)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, None);
        if !result.p_value.is_nan() {
            assert!(result.p_value >= 0.0 && result.p_value <= 1.0);
        }
    }

    #[test]
    fn kpss_large_lag_clamped() {
        let series: Vec<f64> = (0..20)
            .map(|i| ((i * 17 + 13) % 97) as f64 / 50.0 - 1.0)
            .collect();

        let result = kpss_test(&series, Some(100));
        // Lags should be clamped to n/2
        assert!(result.lags <= 10);
    }

    // ==================== kpss_p_value internal ====================

    #[test]
    fn kpss_p_value_boundaries() {
        // stat < 0.347 -> p > 0.10
        let p1 = kpss_p_value(0.1);
        assert!(p1 > 0.10);

        // stat == 0.347 -> p ≈ 0.10
        let p2 = kpss_p_value(0.347);
        assert!((p2 - 0.10).abs() < 0.01);

        // stat == 0.463 -> p ≈ 0.05
        let p3 = kpss_p_value(0.463);
        assert!((p3 - 0.05).abs() < 0.01);

        // stat == 0.739 -> p ≈ 0.01
        let p4 = kpss_p_value(0.739);
        assert!((p4 - 0.01).abs() < 0.01);

        // stat > 0.739 -> p < 0.01
        let p5 = kpss_p_value(1.0);
        assert!(p5 < 0.01);

        // NaN input -> NaN output
        let p_nan = kpss_p_value(f64::NAN);
        assert!(p_nan.is_nan());
    }

    // ==================== adf_p_value internal ====================

    #[test]
    fn adf_p_value_boundaries() {
        // Very negative stat -> small p-value
        let p1 = adf_p_value(-5.0, 100);
        assert!(p1 < 0.01);

        // Stat near 0 -> large p-value
        let p2 = adf_p_value(0.0, 100);
        assert!(p2 > 0.5);

        // Positive stat -> very large p-value
        let p3 = adf_p_value(2.0, 100);
        assert!(p3 > 0.5);

        // NaN input -> NaN output
        let p_nan = adf_p_value(f64::NAN, 100);
        assert!(p_nan.is_nan());
    }

    // ==================== test_stationarity edge cases ====================

    #[test]
    fn combined_test_constant() {
        let series = vec![42.0; 100];
        let (adf, kpss, _conclusion) = test_stationarity(&series);

        // Just verify it doesn't panic and returns valid results
        let _ = adf.statistic;
        let _ = kpss.statistic;
        assert!(
            _conclusion == "stationary"
                || _conclusion == "non_stationary"
                || _conclusion == "inconclusive"
        );
    }

    #[test]
    fn combined_test_empty() {
        let (adf, kpss, _) = test_stationarity(&[]);
        assert!(adf.statistic.is_nan());
        assert!(kpss.statistic.is_nan());
    }
}