foxhunt/crates/ml/src/preprocessing.rs

//! Data Preprocessing Module
//!
//! Transforms raw OHLCV price data into stationary, normalized features suitable
//! for machine learning models. This module addresses the core challenges identified
//! in Wave 14:
//!
//! - **Non-stationarity**: Raw prices fail ADF test (p=0.1987)
//! - **Extreme volatility**: 177x price range causes gradient explosions
//! - **Fat tails**: Kurtosis=346.6 indicates severe outliers
//!
//! ## Solution Strategy
//!
//! 1. **Log Returns**: Transform prices to log(P_t / P_{t-1}) for stationarity
//! 2. **Windowed Normalization**: Apply rolling z-score normalization
//! 3. **Outlier Clipping**: Clip extreme values to +/- N sigma
//!
//! ## Expected Impact
//!
//! - 50-70% reduction in gradient explosions
//! - Improved stationarity (ADF p-value < 0.05)
//! - Reduced kurtosis (< 10.0)
//! - Bounded feature range for stable training
//!
//! ## Usage Example
//!
//! ```rust,no_run
//! use ml::preprocessing::{preprocess_prices_cpu, PreprocessConfig};
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let prices = vec![100.0_f32, 105.0, 103.0, 110.0];
//! let config = PreprocessConfig {
//!     window_size: 3,
//!     clip_sigma: 3.0,
//!     use_log_returns: true,
//! };
//! let preprocessed = preprocess_prices_cpu(&prices, config)?;
//! # Ok(())
//! # }
//! ```
//!
//! ## References
//!
//! - Wave 14 Agent 23: Root cause analysis (non-stationarity)
//! - Wave 14 Agent 25: Solution consensus (100% paper validation)
//! - Wave 14 Agent 28: TDD implementation (this module)

use crate::MLError;

/// Preprocessing configuration
///
/// Controls the behavior of the preprocessing pipeline:
/// - `window_size`: Rolling window for normalization (typically 60-240 bars)
/// - `clip_sigma`: Standard deviations for outlier clipping (typically 2.0-4.0)
/// - `use_log_returns`: Use log returns vs simple returns (recommended: true)
#[derive(Debug, Clone, Copy)]
pub struct PreprocessConfig {
    /// Rolling window size for normalization (default: 120)
    pub window_size: i64,

    /// Outlier clipping threshold in standard deviations (default: 3.0)
    pub clip_sigma: f64,

    /// Use log returns instead of simple returns (default: true)
    pub use_log_returns: bool,
}

impl Default for PreprocessConfig {
    fn default() -> Self {
        Self {
            window_size: 120, // 2 hours for 1-minute bars
            clip_sigma: 3.0,  // Clip beyond +/-3 sigma
            use_log_returns: true,
        }
    }
}

/// Compute log returns: log(P_t / P_{t-1})
///
/// Transforms raw prices into stationary log returns. The first value is set to 0.0
/// as a placeholder (no previous price to compare against).
///
/// # Arguments
///
/// * `prices` - Slice of price values
///
/// # Returns
///
/// * `Ok(Vec<f32>)` - Log returns, first value is 0.0
/// * `Err(MLError)` - If input has fewer than 2 prices
pub fn compute_log_returns(prices: &[f32]) -> Result<Vec<f32>, MLError> {
    if prices.len() < 2 {
        return Err(MLError::InvalidInput(
            "Need at least 2 prices to compute returns".to_owned(),
        ));
    }

    let mut returns = Vec::with_capacity(prices.len());
    returns.push(0.0); // First value placeholder

    for i in 1..prices.len() {
        let prev = prices[i - 1];
        let curr = prices[i];
        if prev > 0.0 && curr > 0.0 {
            returns.push((curr / prev).ln());
        } else {
            returns.push(0.0);
        }
    }

    Ok(returns)
}

/// Apply windowed z-score normalization
///
/// Normalizes data using a rolling window approach:
/// - For each position, compute mean and std of the preceding window
/// - Transform value to z-score: (x - mean) / std
///
/// # Arguments
///
/// * `data` - Slice of data to normalize
/// * `window_size` - Size of rolling window (e.g., 120 for 2-hour window)
///
/// # Returns
///
/// * `Ok(Vec<f32>)` - Normalized data
/// * `Err(MLError)` - If data length is less than window_size
pub fn windowed_normalize(data: &[f32], window_size: i64) -> Result<Vec<f32>, MLError> {
    let n = data.len() as i64;

    if n < window_size {
        return Err(MLError::InvalidInput(format!(
            "Data length ({}) must be >= window_size ({})",
            n, window_size
        )));
    }

    let mut normalized = Vec::with_capacity(n as usize);

    for i in 0..n as usize {
        let start = if i + 1 >= window_size as usize {
            i + 1 - window_size as usize
        } else {
            0
        };

        let window = &data[start..=i];

        // Compute mean
        let mean: f32 = window.iter().sum::<f32>() / window.len() as f32;

        // Compute std (population std)
        let variance: f32 =
            window.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / window.len() as f32;
        let std = variance.sqrt();

        // Compute z-score with epsilon for numerical stability
        let eps = 1e-8;
        let z_score = if std > eps {
            (data[i] - mean) / std
        } else {
            0.0 // If std is too small, return 0 (no signal)
        };

        normalized.push(z_score);
    }

    Ok(normalized)
}

/// Clip outliers to +/- N sigma.
///
/// Caps extreme values at `mean +/- N * std` to prevent gradient explosions
/// from fat-tailed distributions.
///
/// # Arguments
///
/// * `data` - Slice of data to clip
/// * `n_sigma` - Number of standard deviations for clipping threshold
///
/// # Returns
///
/// * `Ok(Vec<f32>)` - Clipped data
/// * `Err(MLError)` - If data is empty
pub fn clip_outliers(data: &[f32], n_sigma: f64) -> Result<Vec<f32>, MLError> {
    if data.is_empty() {
        return Err(MLError::InvalidInput("Empty data".to_owned()));
    }

    let mean: f32 = data.iter().sum::<f32>() / data.len() as f32;
    let variance: f32 =
        data.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
    let std = variance.sqrt();

    let lower_bound = mean - (n_sigma as f32) * std;
    let upper_bound = mean + (n_sigma as f32) * std;

    Ok(data
        .iter()
        .map(|&x| x.clamp(lower_bound, upper_bound))
        .collect())
}

/// Compute clipping bounds from a reference dataset (typically training data only).
///
/// Returns `(lower_bound, upper_bound)` as `mean +/- n_sigma * std`.
pub fn compute_clip_bounds(reference_data: &[f32], n_sigma: f64) -> Result<(f64, f64), MLError> {
    if reference_data.is_empty() {
        return Err(MLError::InvalidInput("Empty reference data".to_owned()));
    }

    let mean: f32 = reference_data.iter().sum::<f32>() / reference_data.len() as f32;
    let variance: f32 = reference_data
        .iter()
        .map(|&x| (x - mean).powi(2))
        .sum::<f32>()
        / reference_data.len() as f32;
    let std = variance.sqrt();

    let lower = f64::from(mean) - n_sigma * f64::from(std);
    let upper = f64::from(mean) + n_sigma * f64::from(std);

    Ok((lower, upper))
}

/// Clip outliers using pre-computed bounds (leakage-free).
pub fn clip_outliers_with_bounds(data: &[f32], lower: f64, upper: f64) -> Result<Vec<f32>, MLError> {
    Ok(data
        .iter()
        .map(|&x| x.clamp(lower as f32, upper as f32))
        .collect())
}

/// Full preprocessing pipeline (CPU-resident, Vec-based).
///
/// Applies the complete preprocessing sequence:
/// 1. Compute log returns from prices
/// 2. Apply windowed normalization
/// 3. Clip outliers
pub fn preprocess_prices_cpu(
    close_prices: &[f32],
    config: PreprocessConfig,
) -> Result<Vec<f32>, MLError> {
    // Step 1: Compute returns
    let returns = if config.use_log_returns {
        compute_log_returns(close_prices)?
    } else {
        compute_simple_returns(close_prices)?
    };

    // Step 2: Windowed normalization
    let normalized = windowed_normalize(&returns, config.window_size)?;

    // Step 3: Clip outliers
    clip_outliers(&normalized, config.clip_sigma)
}

/// Full preprocessing pipeline with optional pre-computed clipping bounds.
pub fn preprocess_prices_with_bounds_cpu(
    close_prices: &[f32],
    config: PreprocessConfig,
    clip_bounds: Option<(f64, f64)>,
) -> Result<Vec<f32>, MLError> {
    let returns = if config.use_log_returns {
        compute_log_returns(close_prices)?
    } else {
        compute_simple_returns(close_prices)?
    };

    let normalized = windowed_normalize(&returns, config.window_size)?;

    match clip_bounds {
        Some((lower, upper)) => clip_outliers_with_bounds(&normalized, lower, upper),
        None => clip_outliers(&normalized, config.clip_sigma),
    }
}

/// Compute simple returns: (P_t - P_{t-1}) / P_{t-1}
fn compute_simple_returns(prices: &[f32]) -> Result<Vec<f32>, MLError> {
    if prices.len() < 2 {
        return Err(MLError::InvalidInput(
            "Need at least 2 prices to compute returns".to_owned(),
        ));
    }

    let mut returns = Vec::with_capacity(prices.len());
    returns.push(0.0); // First value placeholder

    for i in 1..prices.len() {
        let prev = prices[i - 1];
        if prev.abs() > 1e-10 {
            returns.push((prices[i] - prev) / prev);
        } else {
            returns.push(0.0);
        }
    }

    Ok(returns)
}

// Legacy aliases for backward compatibility
/// Full preprocessing pipeline (alias for `preprocess_prices_cpu`).
pub fn preprocess_prices(
    close_prices: &[f32],
    config: PreprocessConfig,
) -> Result<Vec<f32>, MLError> {
    preprocess_prices_cpu(close_prices, config)
}

/// Full preprocessing pipeline with bounds (alias for `preprocess_prices_with_bounds_cpu`).
pub fn preprocess_prices_with_bounds(
    close_prices: &[f32],
    config: PreprocessConfig,
    clip_bounds: Option<(f64, f64)>,
) -> Result<Vec<f32>, MLError> {
    preprocess_prices_with_bounds_cpu(close_prices, config, clip_bounds)
}

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

    #[test]
    fn test_config_default() {
        let config = PreprocessConfig::default();
        assert_eq!(config.window_size, 120);
        assert!((config.clip_sigma - 3.0).abs() < 1e-6);
        assert!(config.use_log_returns);
    }

    #[test]
    fn test_compute_log_returns_basic() {
        let prices = [100.0_f32, 110.0, 105.0];
        let returns = compute_log_returns(&prices).expect("Failed to compute log returns");

        assert_eq!(returns.len(), 3);
        assert!((returns[0] - 0.0).abs() < 1e-6, "First value should be 0");
        assert!((returns[1] - (110.0_f32 / 100.0).ln()).abs() < 1e-5);
    }

    #[test]
    fn test_windowed_normalize_basic() {
        let data = [1.0_f32, 2.0, 3.0, 4.0, 5.0];
        let normalized = windowed_normalize(&data, 3).expect("Failed to normalize");

        assert_eq!(normalized.len(), 5);
        // All values should be finite
        for &v in &normalized {
            assert!(v.is_finite(), "Normalized values contain NaN/Inf");
        }
    }

    #[test]
    fn test_clip_outliers_basic() {
        let data = [10.0_f32, 11.0, 12.0, 13.0, 14.0, 50.0, -20.0];
        let clipped = clip_outliers(&data, 1.5).expect("Failed to clip");

        assert_eq!(clipped.len(), 7);
        assert!(clipped[5] < 50.0, "Positive outlier (50.0) should be clipped, got {}", clipped[5]);
        assert!(clipped[6] > -20.0, "Negative outlier (-20.0) should be clipped, got {}", clipped[6]);
        assert!((clipped[0] - 10.0).abs() < 0.1, "Normal value should not change much");
        assert!((clipped[1] - 11.0).abs() < 0.1, "Normal value should not change much");
    }

    #[test]
    fn test_compute_clip_bounds_basic() {
        let data = [1.0_f32, 2.0, 3.0, 4.0, 5.0];
        let (lower, upper) = compute_clip_bounds(&data, 2.0).expect("Failed to compute clip bounds");

        assert!(lower < 1.0, "Lower bound should be below 1.0, got {}", lower);
        assert!(upper > 5.0, "Upper bound should be above 5.0, got {}", upper);
        let mid = (lower + upper) / 2.0;
        assert!(
            (mid - 3.0).abs() < 0.01,
            "Bounds should be symmetric around mean 3.0, midpoint={}",
            mid
        );
    }

    #[test]
    fn test_clip_outliers_with_bounds_clips_correctly() {
        let data = [0.01_f32, 10.0, -8.0, 0.02];
        let clipped = clip_outliers_with_bounds(&data, -3.0, 3.0).expect("Failed to clip with bounds");

        assert_eq!(clipped.len(), 4);
        assert!((clipped[0] - 0.01).abs() < 1e-6);
        assert!((clipped[1] - 3.0).abs() < 1e-6);
        assert!((clipped[2] - (-3.0)).abs() < 1e-6);
        assert!((clipped[3] - 0.02).abs() < 1e-6);
    }

    #[test]
    fn test_clip_outliers_with_bounds_no_op_when_in_range() {
        let data = [1.0_f32, 2.0, 3.0];
        let clipped = clip_outliers_with_bounds(&data, -100.0, 100.0).expect("Failed to clip with bounds");

        for (i, (&orig, &clipped_v)) in data.iter().zip(clipped.iter()).enumerate() {
            assert!((orig - clipped_v).abs() < 1e-6, "Value {} changed: {} -> {}", i, orig, clipped_v);
        }
    }

    #[test]
    fn test_compute_bounds_then_apply_to_other_split() {
        let train = [0.1_f32, -0.05, 0.2, -0.1, 0.15];
        let val = [0.3_f32, -0.5, 0.01];

        let (lo, hi) = compute_clip_bounds(&train, 2.0).expect("Failed to compute clip bounds");
        let clipped_val = clip_outliers_with_bounds(&val, lo, hi).expect("Failed to clip val data");

        assert_eq!(clipped_val.len(), 3);
        for &v in &clipped_val {
            assert!(f64::from(v) >= lo - 1e-6, "Val {} below lower bound {}", v, lo);
            assert!(f64::from(v) <= hi + 1e-6, "Val {} above upper bound {}", v, hi);
        }
    }

    #[test]
    fn test_preprocess_prices_with_bounds_none_matches_original() {
        let mut prices_vec = Vec::with_capacity(150);
        let mut price = 100.0_f32;
        for i in 0..150 {
            price += (i as f32 * 0.1).sin() * 0.5;
            prices_vec.push(price);
        }

        let config = PreprocessConfig {
            window_size: 20,
            clip_sigma: 3.0,
            use_log_returns: true,
        };

        let original = preprocess_prices(&prices_vec, config).expect("preprocess_prices failed");
        let with_none = preprocess_prices_with_bounds(&prices_vec, config, None)
            .expect("with_bounds(None) failed");

        assert_eq!(original.len(), with_none.len());
        for (i, (&a, &b)) in original.iter().zip(with_none.iter()).enumerate() {
            assert!((a - b).abs() < 1e-6, "Mismatch at index {}: {} vs {}", i, a, b);
        }
    }

    #[test]
    fn test_preprocess_prices_with_bounds_uses_provided_bounds() {
        let mut prices_vec = Vec::with_capacity(150);
        let mut price = 100.0_f32;
        for i in 0..150 {
            price += (i as f32 * 0.1).sin() * 0.5;
            prices_vec.push(price);
        }

        let config = PreprocessConfig {
            window_size: 20,
            clip_sigma: 3.0,
            use_log_returns: true,
        };

        let tight_bounds = Some((-0.001, 0.001));
        let clipped = preprocess_prices_with_bounds(&prices_vec, config, tight_bounds)
            .expect("with_bounds failed");

        for &v in &clipped {
            assert!(v >= -0.001 - 1e-6, "Min {} outside tight bounds", v);
            assert!(v <= 0.001 + 1e-6, "Max {} outside tight bounds", v);
        }
    }
}