anofox-forecast 0.15.9

//! Trend-based features for time series.
//!
//! Provides features related to linear trends and autoregressive properties.

/// Result of linear regression.
#[derive(Debug, Clone)]
pub struct LinearTrendResult {
    /// Slope of the fitted line
    pub slope: f64,
    /// Intercept of the fitted line
    pub intercept: f64,
    /// R-squared (coefficient of determination)
    pub r_squared: f64,
    /// Standard error of the slope
    pub stderr: f64,
    /// P-value for the slope (approximate)
    pub p_value: f64,
}

/// Computes linear trend statistics for the time series.
///
/// Fits a linear regression y = slope * x + intercept where x is the index.
pub fn linear_trend(series: &[f64]) -> LinearTrendResult {
    if series.len() < 2 {
        return LinearTrendResult {
            slope: f64::NAN,
            intercept: f64::NAN,
            r_squared: f64::NAN,
            stderr: f64::NAN,
            p_value: f64::NAN,
        };
    }

    let n = series.len() as f64;

    // x values are indices 0, 1, 2, ...
    let sum_x: f64 = (0..series.len()).map(|i| i as f64).sum();
    let sum_y: f64 = series.iter().sum();
    let sum_xy: f64 = series.iter().enumerate().map(|(i, &y)| i as f64 * y).sum();
    let sum_x2: f64 = (0..series.len()).map(|i| (i * i) as f64).sum();

    let mean_x = sum_x / n;
    let mean_y = sum_y / n;

    let ss_xx = sum_x2 - n * mean_x * mean_x;
    let ss_xy = sum_xy - n * mean_x * mean_y;

    if ss_xx.abs() < 1e-10 {
        return LinearTrendResult {
            slope: 0.0,
            intercept: mean_y,
            r_squared: 0.0,
            stderr: f64::NAN,
            p_value: 1.0,
        };
    }

    let slope = ss_xy / ss_xx;
    let intercept = mean_y - slope * mean_x;

    // Compute R-squared
    let ss_yy: f64 = series.iter().map(|&y| (y - mean_y).powi(2)).sum();
    let ss_res: f64 = series
        .iter()
        .enumerate()
        .map(|(i, &y)| {
            let y_pred = slope * i as f64 + intercept;
            (y - y_pred).powi(2)
        })
        .sum();

    let r_squared = if ss_yy.abs() < 1e-10 {
        1.0 // Perfect fit for constant y
    } else {
        1.0 - ss_res / ss_yy
    };

    // Compute standard error of slope
    let mse = if n > 2.0 { ss_res / (n - 2.0) } else { 0.0 };
    let stderr = if ss_xx > 0.0 {
        (mse / ss_xx).sqrt()
    } else {
        f64::NAN
    };

    // Compute approximate p-value using t-distribution
    let t_stat = if stderr > 1e-10 {
        slope / stderr
    } else {
        f64::INFINITY
    };

    // Approximate p-value (2-tailed) using normal approximation for large n
    let p_value = 2.0 * (1.0 - normal_cdf(t_stat.abs()));

    LinearTrendResult {
        slope,
        intercept,
        r_squared,
        stderr,
        p_value,
    }
}

/// Computes aggregated linear trend statistics (tsfresh-compatible).
///
/// Divides the series into chunks of fixed length and computes linear trends for each,
/// then aggregates the results.
///
/// # Arguments
/// * `series` - Input time series
/// * `chunk_len` - Length of each chunk (number of values per chunk)
/// * `agg_func` - Aggregation function: "mean", "var", "std", "min", "max"
/// * `attribute` - Which attribute to aggregate: "slope", "intercept", "rvalue", "stderr", "pvalue"
pub fn agg_linear_trend(series: &[f64], chunk_len: usize, agg_func: &str, attribute: &str) -> f64 {
    if series.is_empty() || chunk_len == 0 || chunk_len > series.len() {
        return f64::NAN;
    }

    // Create chunks of size chunk_len
    let chunks: Vec<&[f64]> = series.chunks(chunk_len).collect();

    // Aggregate each chunk using the specified function
    let aggregated: Vec<f64> = chunks
        .iter()
        .map(|chunk| crate::utils::stats::aggregate(chunk, agg_func))
        .filter(|v| !v.is_nan())
        .collect();

    if aggregated.len() < 2 {
        return f64::NAN;
    }

    // Perform linear regression on indices [0, 1, 2, ...] vs aggregated values
    let x: Vec<f64> = (0..aggregated.len()).map(|i| i as f64).collect();
    let trend = linear_regression(&x, &aggregated);

    match attribute {
        "slope" => trend.slope,
        "intercept" => trend.intercept,
        "rvalue" => trend.r_squared.sqrt(),
        "r_squared" => trend.r_squared,
        "pvalue" => trend.p_value,
        "stderr" => trend.stderr,
        _ => f64::NAN,
    }
}

/// Perform linear regression and return trend result
fn linear_regression(x: &[f64], y: &[f64]) -> LinearTrendResult {
    if x.len() != y.len() || x.len() < 2 {
        return LinearTrendResult {
            slope: f64::NAN,
            intercept: f64::NAN,
            r_squared: f64::NAN,
            stderr: f64::NAN,
            p_value: f64::NAN,
        };
    }

    let n = x.len() as f64;
    let sum_x: f64 = x.iter().sum();
    let sum_y: f64 = y.iter().sum();
    let sum_xy: f64 = x.iter().zip(y.iter()).map(|(&xi, &yi)| xi * yi).sum();
    let sum_x2: f64 = x.iter().map(|&xi| xi * xi).sum();

    let mean_x = sum_x / n;
    let mean_y = sum_y / n;

    let ss_xx = sum_x2 - n * mean_x * mean_x;
    let ss_xy = sum_xy - n * mean_x * mean_y;

    if ss_xx.abs() < 1e-10 {
        return LinearTrendResult {
            slope: 0.0,
            intercept: mean_y,
            r_squared: 0.0,
            stderr: f64::NAN,
            p_value: 1.0,
        };
    }

    let slope = ss_xy / ss_xx;
    let intercept = mean_y - slope * mean_x;

    // Compute R-squared
    let ss_yy: f64 = y.iter().map(|&yi| (yi - mean_y).powi(2)).sum();
    let ss_res: f64 = x
        .iter()
        .zip(y.iter())
        .map(|(&xi, &yi)| {
            let y_pred = slope * xi + intercept;
            (yi - y_pred).powi(2)
        })
        .sum();

    let r_squared = if ss_yy.abs() < 1e-10 {
        1.0
    } else {
        1.0 - ss_res / ss_yy
    };

    // Standard error of slope
    let mse = if n > 2.0 { ss_res / (n - 2.0) } else { 0.0 };
    let stderr = if ss_xx > 0.0 {
        (mse / ss_xx).sqrt()
    } else {
        f64::NAN
    };

    // P-value
    let t_stat = if stderr > 1e-10 {
        slope / stderr
    } else {
        f64::INFINITY
    };
    let p_value = 2.0 * (1.0 - normal_cdf(t_stat.abs()));

    LinearTrendResult {
        slope,
        intercept,
        r_squared,
        stderr,
        p_value,
    }
}

/// Estimates AR coefficients using OLS (matches tsfresh AutoReg).
///
/// Uses Ordinary Least Squares to fit an AR(k) model with intercept.
/// Returns the specified coefficient.
///
/// # Arguments
/// * `series` - Input time series
/// * `k` - AR order (maximum lag)
/// * `coeff` - Which coefficient to return (0=intercept, 1..k=AR coefficients)
pub fn ar_coefficient(series: &[f64], k: usize, coeff: usize) -> f64 {
    if series.len() <= k || k == 0 || coeff > k {
        return f64::NAN;
    }

    let n = series.len();
    let n_obs = n - k;

    if n_obs < k + 2 {
        return f64::NAN;
    }

    let (xtx, xty) = build_ar_normal_equations(series, k, n);
    match solve_linear_system(&xtx, &xty) {
        Some(p) if coeff < p.len() => p[coeff],
        _ => f64::NAN,
    }
}

/// Build X'X and X'y normal equations for AR(k) model with intercept.
#[inline]
fn build_ar_normal_equations(series: &[f64], k: usize, n: usize) -> (Vec<Vec<f64>>, Vec<f64>) {
    let n_params = k + 1;
    let mut xtx = vec![vec![0.0; n_params]; n_params];
    let mut xty = vec![0.0; n_params];

    for t in k..n {
        let y_t = series[t];

        // Accumulate intercept row/column
        xtx[0][0] += 1.0;
        xty[0] += y_t;
        for i in 1..n_params {
            let xi = series[t - i];
            xtx[0][i] += xi;
            xtx[i][0] += xi;
            xty[i] += xi * y_t;
            for j in 1..n_params {
                xtx[i][j] += xi * series[t - j];
            }
        }
    }

    (xtx, xty)
}

/// Solve a linear system Ax = b using Gaussian elimination with partial pivoting.
fn solve_linear_system(a: &[Vec<f64>], b: &[f64]) -> Option<Vec<f64>> {
    let n = b.len();
    if n == 0 || a.len() != n {
        return None;
    }

    let mut aug = build_augmented_matrix(a, b, n);

    if !gaussian_eliminate(&mut aug, n) {
        return None;
    }

    Some(back_substitute_augmented(&aug, n))
}

/// Build augmented matrix [A | b].
#[inline]
fn build_augmented_matrix(a: &[Vec<f64>], b: &[f64], n: usize) -> Vec<Vec<f64>> {
    a.iter()
        .enumerate()
        .map(|(i, row)| {
            let mut r = Vec::with_capacity(n + 1);
            r.extend_from_slice(row);
            r.push(b[i]);
            r
        })
        .collect()
}

/// Gaussian elimination with partial pivoting. Returns false if singular.
#[inline]
fn gaussian_eliminate(aug: &mut [Vec<f64>], n: usize) -> bool {
    for col in 0..n {
        // Find pivot
        let (max_row, max_val) = (col..n).fold((col, aug[col][col].abs()), |(mr, mv), row| {
            let v = aug[row][col].abs();
            if v > mv {
                (row, v)
            } else {
                (mr, mv)
            }
        });

        if max_val < 1e-14 {
            return false;
        }

        aug.swap(col, max_row);

        for row in (col + 1)..n {
            let factor = aug[row][col] / aug[col][col];
            for j in col..=n {
                aug[row][j] -= factor * aug[col][j];
            }
        }
    }
    true
}

/// Back substitution on augmented matrix.
#[inline]
fn back_substitute_augmented(aug: &[Vec<f64>], n: usize) -> Vec<f64> {
    let mut x = vec![0.0; n];
    for i in (0..n).rev() {
        let mut sum = aug[i][n];
        for j in (i + 1)..n {
            sum -= aug[i][j] * x[j];
        }
        x[i] = sum / aug[i][i];
    }
    x
}

/// Estimates AR coefficients using Yule-Walker equations.
///
/// This is an alternative to the OLS-based method.
///
/// # Arguments
/// * `series` - Input time series
/// * `k` - Which coefficient to return (1-indexed)
pub fn ar_coefficient_yule_walker(series: &[f64], k: usize) -> f64 {
    if series.len() <= k || k == 0 {
        return f64::NAN;
    }

    let mean: f64 = series.iter().sum::<f64>() / series.len() as f64;
    let var: f64 = series.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / series.len() as f64;

    if var < 1e-10 {
        return 0.0;
    }

    let acf = compute_acf_for_yule_walker(series, k, mean, var);
    durbin_levinson(&acf, k)
}

/// Compute autocorrelation function values for Yule-Walker.
#[inline]
fn compute_acf_for_yule_walker(series: &[f64], k: usize, mean: f64, var: f64) -> Vec<f64> {
    (0..=k)
        .map(|lag| {
            if lag == 0 {
                1.0
            } else {
                let cov: f64 = series[lag..]
                    .iter()
                    .zip(series.iter())
                    .map(|(&x1, &x2)| (x1 - mean) * (x2 - mean))
                    .sum::<f64>()
                    / series.len() as f64;
                cov / var
            }
        })
        .collect()
}

/// Durbin-Levinson algorithm: solve Yule-Walker equations.
/// Returns the k-th AR coefficient, or NaN on degeneracy.
#[inline]
fn durbin_levinson(acf: &[f64], k: usize) -> f64 {
    let mut phi = vec![0.0; k + 1];
    phi[1] = acf[1];

    for m in 2..=k {
        let num: f64 = acf[m] - (1..m).map(|j| phi[j] * acf[m - j]).sum::<f64>();
        let denom: f64 = 1.0 - (1..m).map(|j| phi[j] * acf[j]).sum::<f64>();

        if denom.abs() < 1e-10 {
            return f64::NAN;
        }

        let new_phi = num / denom;
        let mut new_coeffs = vec![0.0; k + 1];
        new_coeffs[m] = new_phi;
        for j in 1..m {
            new_coeffs[j] = phi[j] - new_phi * phi[m - j];
        }
        phi = new_coeffs;
    }

    phi[k]
}

/// Performs the Augmented Dickey-Fuller test for stationarity.
///
/// Returns the ADF test statistic. More negative values indicate
/// stronger evidence against the unit root hypothesis (i.e., evidence for stationarity).
pub fn augmented_dickey_fuller(series: &[f64]) -> f64 {
    if series.len() < 4 {
        return f64::NAN;
    }

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

    let n = diff.len();

    // Lagged level (y_{t-1})
    let y_lag: Vec<f64> = series[..n].to_vec();

    // Simple regression: diff = alpha + beta * y_lag + error
    let mean_y_lag: f64 = y_lag.iter().sum::<f64>() / n as f64;
    let mean_diff: f64 = diff.iter().sum::<f64>() / n as f64;

    let ss_yy: f64 = y_lag.iter().map(|&y| (y - mean_y_lag).powi(2)).sum();
    let ss_xy: f64 = diff
        .iter()
        .zip(y_lag.iter())
        .map(|(&d, &y)| (d - mean_diff) * (y - mean_y_lag))
        .sum();

    if ss_yy.abs() < 1e-10 {
        return f64::NAN;
    }

    let beta = ss_xy / ss_yy;
    let alpha = mean_diff - beta * mean_y_lag;

    // Compute residuals and standard error
    let residuals: Vec<f64> = diff
        .iter()
        .zip(y_lag.iter())
        .map(|(&d, &y)| d - alpha - beta * y)
        .collect();

    let sse: f64 = residuals.iter().map(|r| r * r).sum();
    let mse = sse / (n - 2) as f64;
    let se_beta = (mse / ss_yy).sqrt();

    if se_beta < 1e-10 {
        return f64::NAN;
    }

    // ADF statistic is beta / se(beta)
    beta / se_beta
}

/// Helper: approximate normal CDF using error function approximation
fn normal_cdf(x: f64) -> f64 {
    0.5 * (1.0 + erf(x / std::f64::consts::SQRT_2))
}

/// Helper: error function approximation
fn erf(x: f64) -> f64 {
    // Abramowitz and Stegun approximation
    let a1 = 0.254829592;
    let a2 = -0.284496736;
    let a3 = 1.421413741;
    let a4 = -1.453152027;
    let a5 = 1.061405429;
    let p = 0.3275911;

    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x = x.abs();

    let t = 1.0 / (1.0 + p * x);
    let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();

    sign * y
}

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

    // ==================== linear_trend ====================

    #[test]
    fn linear_trend_perfect_line() {
        // y = 2x + 1
        let series: Vec<f64> = (0..10).map(|i| 2.0 * i as f64 + 1.0).collect();
        let trend = linear_trend(&series);

        assert_relative_eq!(trend.slope, 2.0, epsilon = 1e-10);
        assert_relative_eq!(trend.intercept, 1.0, epsilon = 1e-10);
        assert_relative_eq!(trend.r_squared, 1.0, epsilon = 1e-10);
    }

    #[test]
    fn linear_trend_no_trend() {
        let series = vec![5.0; 10];
        let trend = linear_trend(&series);

        assert_relative_eq!(trend.slope, 0.0, epsilon = 1e-10);
        assert_relative_eq!(trend.intercept, 5.0, epsilon = 1e-10);
    }

    #[test]
    fn linear_trend_negative_slope() {
        // y = -1.5x + 10
        let series: Vec<f64> = (0..10).map(|i| -1.5 * i as f64 + 10.0).collect();
        let trend = linear_trend(&series);

        assert_relative_eq!(trend.slope, -1.5, epsilon = 1e-10);
        assert_relative_eq!(trend.intercept, 10.0, epsilon = 1e-10);
    }

    #[test]
    fn linear_trend_with_noise() {
        // y = x + noise
        let series = vec![0.1, 1.2, 1.9, 3.1, 4.0, 5.2, 5.9, 7.1, 8.0, 9.1];
        let trend = linear_trend(&series);

        assert!(trend.slope > 0.9 && trend.slope < 1.1);
        assert!(trend.r_squared > 0.99);
    }

    #[test]
    fn linear_trend_short() {
        assert!(linear_trend(&[]).slope.is_nan());
        assert!(linear_trend(&[1.0]).slope.is_nan());
    }

    #[test]
    fn linear_trend_two_points() {
        let series = vec![0.0, 10.0];
        let trend = linear_trend(&series);

        assert_relative_eq!(trend.slope, 10.0, epsilon = 1e-10);
        assert_relative_eq!(trend.intercept, 0.0, epsilon = 1e-10);
    }

    // ==================== agg_linear_trend (tsfresh-compatible with chunk_len) ====================

    #[test]
    fn agg_linear_trend_mean_slope() {
        // 100 values, chunk_len=10 -> 10 chunks
        // Linear series: each chunk's mean increases linearly
        let series: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let agg = agg_linear_trend(&series, 10, "mean", "slope");
        // Chunk means: 4.5, 14.5, 24.5, ... -> slope should be 10
        assert_relative_eq!(agg, 10.0, epsilon = 0.1);
    }

    #[test]
    fn agg_linear_trend_rvalue() {
        // Linear series should have perfect r-value
        let series: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let agg = agg_linear_trend(&series, 10, "mean", "rvalue");
        assert_relative_eq!(agg, 1.0, epsilon = 1e-10);
    }

    #[test]
    fn agg_linear_trend_empty() {
        assert!(agg_linear_trend(&[], 5, "mean", "slope").is_nan());
    }

    #[test]
    fn agg_linear_trend_zero_chunk_len() {
        let series = vec![1.0, 2.0, 3.0];
        assert!(agg_linear_trend(&series, 0, "mean", "slope").is_nan());
    }

    #[test]
    fn agg_linear_trend_chunk_len_too_large() {
        let series = vec![1.0, 2.0, 3.0];
        assert!(agg_linear_trend(&series, 10, "mean", "slope").is_nan());
    }

    // ==================== ar_coefficient (OLS-based, tsfresh-compatible) ====================

    #[test]
    fn ar_coefficient_ar1_process() {
        // Simulate AR(1): y[t] = 0.8 * y[t-1] + noise
        let mut series = vec![0.0; 200];
        for i in 1..200 {
            series[i] = 0.8 * series[i - 1] + (i as f64 * 0.1).sin() * 0.1;
        }
        // coeff=1 is the AR(1) coefficient
        let coef = ar_coefficient(&series, 1, 1);
        // Should be close to 0.8
        assert!(
            coef > 0.6 && coef < 1.0,
            "AR(1) coef should be ~0.8, got {}",
            coef
        );
    }

    #[test]
    fn ar_coefficient_intercept() {
        // For data with trend, intercept should be non-zero
        let series: Vec<f64> = (0..100).map(|i| i as f64 * 0.5 + 10.0).collect();
        let intercept = ar_coefficient(&series, 1, 0);
        assert!(!intercept.is_nan());
    }

    #[test]
    fn ar_coefficient_constant() {
        let series = vec![5.0; 50];
        // For constant series, OLS matrix X'X is singular (all lagged values are identical)
        // This should return NaN since the system can't be solved uniquely
        let intercept = ar_coefficient(&series, 1, 0);
        let ar1 = ar_coefficient(&series, 1, 1);
        // Constant series leads to singular matrix, so NaN is expected
        // (or could return values if numerically stable)
        // Just verify it doesn't panic
        let _ = intercept;
        let _ = ar1;
    }

    #[test]
    fn ar_coefficient_short() {
        assert!(ar_coefficient(&[], 1, 1).is_nan());
        assert!(ar_coefficient(&[1.0], 1, 1).is_nan());
        assert!(ar_coefficient(&[1.0, 2.0], 3, 1).is_nan());
    }

    #[test]
    fn ar_coefficient_zero_k() {
        let series = vec![1.0, 2.0, 3.0, 4.0];
        assert!(ar_coefficient(&series, 0, 0).is_nan());
    }

    #[test]
    fn ar_coefficient_coeff_out_of_range() {
        let series: Vec<f64> = (0..50).map(|i| i as f64).collect();
        // k=2 means coeffs 0,1,2 are valid; 3 is invalid
        assert!(ar_coefficient(&series, 2, 3).is_nan());
    }

    // ==================== ar_coefficient_yule_walker ====================

    #[test]
    fn ar_coefficient_yule_walker_works() {
        let mut series = vec![0.0; 200];
        for i in 1..200 {
            series[i] = 0.8 * series[i - 1] + (i as f64 * 0.1).sin() * 0.1;
        }
        let coef = ar_coefficient_yule_walker(&series, 1);
        assert!(coef > 0.6 && coef < 1.0);
    }

    // ==================== augmented_dickey_fuller ====================

    #[test]
    fn adf_stationary() {
        // Stationary series (mean-reverting)
        let series: Vec<f64> = (0..100).map(|i| (i as f64 * 0.5).sin()).collect();
        let adf = augmented_dickey_fuller(&series);
        // Should be negative for stationary series
        assert!(!adf.is_nan());
    }

    #[test]
    fn adf_unit_root() {
        // Random walk (unit root)
        let mut series = vec![0.0; 100];
        for i in 1..100 {
            series[i] = series[i - 1] + ((i * 7) % 11) as f64 - 5.0;
        }
        let adf = augmented_dickey_fuller(&series);
        assert!(!adf.is_nan());
    }

    #[test]
    fn adf_trending() {
        // Strong upward trend with some noise
        let series: Vec<f64> = (0..100)
            .map(|i| i as f64 * 2.0 + ((i * 7) % 5) as f64 * 0.1)
            .collect();
        let adf = augmented_dickey_fuller(&series);
        assert!(!adf.is_nan());
    }

    #[test]
    fn adf_short() {
        assert!(augmented_dickey_fuller(&[]).is_nan());
        assert!(augmented_dickey_fuller(&[1.0, 2.0]).is_nan());
    }

    #[test]
    fn adf_constant() {
        let series = vec![5.0; 50];
        let adf = augmented_dickey_fuller(&series);
        // Constant series has no variation, may return NaN or extreme value
        assert!(adf.is_nan() || adf.abs() < 1e-10);
    }
}