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, PreprocessConfig};
//! use candle_core::{Tensor, Device};
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Load price data
//! let prices = Tensor::from_slice(&[100.0_f32, 105.0, 103.0, 110.0], (4,), &Device::Cpu)?;
//!
//! // Configure preprocessing
//! let config = PreprocessConfig {
//!     window_size: 120,    // 2-hour window for 1-minute bars
//!     clip_sigma: 3.0,     // Clip outliers beyond ±3σ
//!     use_log_returns: true,
//! };
//!
//! // Apply full pipeline
//! let preprocessed = preprocess_prices(&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;
use candle_core::Tensor;

/// 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σ
            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` - Tensor of shape [N] containing price series
///
/// # Returns
///
/// * `Ok(Tensor)` - Log returns of shape [N], first value is 0.0
/// * `Err(MLError)` - If tensor operations fail
///
/// # Mathematical Formula
///
/// r_t = log(P_t / P_{t-1})
///
/// where r_t is the log return at time t and P_t is the price at time t.
///
/// # Example
///
/// ```rust,no_run
/// use ml::preprocessing::compute_log_returns;
/// use candle_core::{Tensor, Device};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let prices = Tensor::from_slice(&[100.0_f32, 105.0, 103.0], (3,), &Device::Cpu)?;
/// let returns = compute_log_returns(&prices)?;
/// // returns ≈ [0.0, 0.04879, -0.01942]
/// # Ok(())
/// # }
/// ```
pub fn compute_log_returns(prices: &Tensor) -> Result<Tensor, MLError> {
    let n = prices.dims()[0];

    if n < 2 {
        return Err(MLError::InvalidInput(
            "Need at least 2 prices to compute returns".to_owned(),
        ));
    }

    // Get shifted tensors: prices[:-1] and prices[1:]
    let prev_prices = prices.narrow(0, 0, n - 1).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to narrow prev_prices: {}", e))
    })?;

    let curr_prices = prices.narrow(0, 1, n - 1).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to narrow curr_prices: {}", e))
    })?;

    // Compute log(P_t / P_{t-1}) = log(P_t) - log(P_{t-1})
    let log_curr = curr_prices.log().map_err(|e| {
        MLError::TensorOperationError(format!("Failed to compute log of current prices: {}", e))
    })?;

    let log_prev = prev_prices.log().map_err(|e| {
        MLError::TensorOperationError(format!("Failed to compute log of previous prices: {}", e))
    })?;

    let returns = log_curr.sub(&log_prev).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to compute log returns: {}", e))
    })?;

    // Prepend 0.0 for first value (placeholder)
    let first_zero = Tensor::zeros((1,), returns.dtype(), returns.device()).map_err(|e| {
        MLError::TensorCreationError {
            operation: "create_first_zero".to_owned(),
            reason: e.to_string(),
        }
    })?;

    let result = Tensor::cat(&[&first_zero, &returns], 0).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to concatenate returns: {}", e))
    })?;

    Ok(result)
}

/// 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
///
/// This approach handles non-stationary volatility by adapting to local statistics.
///
/// # Arguments
///
/// * `data` - Tensor of shape [N] containing data to normalize
/// * `window_size` - Size of rolling window (e.g., 120 for 2-hour window)
///
/// # Returns
///
/// * `Ok(Tensor)` - Normalized data of shape [N]
/// * `Err(MLError)` - If tensor operations fail
///
/// # Example
///
/// ```rust,no_run
/// use ml::preprocessing::windowed_normalize;
/// use candle_core::{Tensor, Device};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let returns = Tensor::from_slice(&[0.01_f32, 0.02, 0.10, 0.15], (4,), &Device::Cpu)?;
/// let normalized = windowed_normalize(&returns, 3)?;
/// // Each window has mean≈0, std≈1
/// # Ok(())
/// # }
/// ```
pub fn windowed_normalize(data: &Tensor, window_size: i64) -> Result<Tensor, MLError> {
    let n = data.dims()[0] as i64;

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

    let device = data.device();

    // Convert to Vec for easier processing
    let data_vec: Vec<f32> = data.to_vec1().map_err(|e| {
        MLError::TensorOperationError(format!("Failed to convert data to vec: {}", e))
    })?;

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

    for i in 0..n as usize {
        // Define window: max(0, i - window_size + 1) to i (inclusive)
        let start = if i + 1 >= window_size as usize {
            i + 1 - window_size as usize
        } else {
            0
        };

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

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

        // Compute std (sample std with Bessel's correction)
        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_vec[i] - mean) / std
        } else {
            0.0 // If std is too small, return 0 (no signal)
        };

        normalized.push(z_score);
    }

    // Convert back to tensor
    let result = Tensor::from_slice(&normalized, (n as usize,), device).map_err(|e| {
        MLError::TensorCreationError {
            operation: "create_normalized_tensor".to_owned(),
            reason: e.to_string(),
        }
    })?;

    Ok(result)
}

/// Clip outliers to ±N sigma
///
/// Caps extreme values at mean ± N standard deviations to prevent
/// gradient explosions from fat-tailed distributions.
///
/// # Arguments
///
/// * `data` - Tensor of shape [N] containing data to clip
/// * `n_sigma` - Number of standard deviations for clipping threshold
///
/// # Returns
///
/// * `Ok(Tensor)` - Clipped data of shape [N]
/// * `Err(MLError)` - If tensor operations fail
///
/// # Example
///
/// ```rust,no_run
/// use ml::preprocessing::clip_outliers;
/// use candle_core::{Tensor, Device};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let returns = Tensor::from_slice(&[0.01_f32, 10.0, -8.0, 0.02], (4,), &Device::Cpu)?;
/// let clipped = clip_outliers(&returns, 3.0)?;
/// // Extreme values (10.0, -8.0) will be clipped to ±3σ
/// # Ok(())
/// # }
/// ```
pub fn clip_outliers(data: &Tensor, n_sigma: f64) -> Result<Tensor, MLError> {
    // Compute mean and std
    let mean = data
        .mean_all()
        .map_err(|e| MLError::TensorOperationError(format!("Failed to compute mean: {}", e)))?
        .to_scalar::<f32>()
        .map_err(|e| {
            MLError::TensorOperationError(format!("Failed to convert mean to scalar: {}", e))
        })?;

    // Use var(0) without keepdim to get a scalar
    let variance = data
        .var(0)
        .map_err(|e| MLError::TensorOperationError(format!("Failed to compute variance: {}", e)))?;

    let std = variance
        .sqrt()
        .map_err(|e| MLError::TensorOperationError(format!("Failed to compute std: {}", e)))?
        .to_scalar::<f32>()
        .map_err(|e| {
            MLError::TensorOperationError(format!("Failed to convert std to scalar: {}", e))
        })?;

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

    // Clip using candle's clamp operation
    let clipped = data
        .clamp(lower_bound as f64, upper_bound as f64)
        .map_err(|e| MLError::TensorOperationError(format!("Failed to clamp data: {}", e)))?;

    Ok(clipped)
}

/// Full preprocessing pipeline
///
/// Applies the complete preprocessing sequence:
/// 1. Compute log returns from prices
/// 2. Apply windowed normalization
/// 3. Clip outliers
///
/// This produces stationary, normalized features ready for ML training.
///
/// # Arguments
///
/// * `close_prices` - Tensor of shape [N] containing close prices
/// * `config` - Preprocessing configuration
///
/// # Returns
///
/// * `Ok(Tensor)` - Preprocessed features of shape [N]
/// * `Err(MLError)` - If any preprocessing step fails
///
/// # Example
///
/// ```rust,no_run
/// use ml::preprocessing::{preprocess_prices, PreprocessConfig};
/// use candle_core::{Tensor, Device};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let prices = Tensor::from_slice(&[100.0_f32, 105.0, 103.0, 110.0], (4,), &Device::Cpu)?;
/// let config = PreprocessConfig::default();
/// let preprocessed = preprocess_prices(&prices, config)?;
/// # Ok(())
/// # }
/// ```
pub fn preprocess_prices(
    close_prices: &Tensor,
    config: PreprocessConfig,
) -> Result<Tensor, MLError> {
    // Step 1: Compute log returns
    let returns = if config.use_log_returns {
        compute_log_returns(close_prices)?
    } else {
        // Simple returns: (P_t - P_{t-1}) / P_{t-1}
        let n = close_prices.dims()[0];
        let prev_prices = close_prices.narrow(0, 0, n - 1).map_err(|e| {
            MLError::TensorOperationError(format!("Failed to narrow prev_prices: {}", e))
        })?;

        let curr_prices = close_prices.narrow(0, 1, n - 1).map_err(|e| {
            MLError::TensorOperationError(format!("Failed to narrow curr_prices: {}", e))
        })?;

        let simple_returns = curr_prices
            .sub(&prev_prices)
            .map_err(|e| {
                MLError::TensorOperationError(format!("Failed to compute price diff: {}", e))
            })?
            .div(&prev_prices)
            .map_err(|e| {
                MLError::TensorOperationError(format!("Failed to compute simple returns: {}", e))
            })?;

        // Prepend 0.0 for first value
        let first_zero = Tensor::zeros((1,), simple_returns.dtype(), simple_returns.device())
            .map_err(|e| MLError::TensorCreationError {
                operation: "create_first_zero".to_owned(),
                reason: e.to_string(),
            })?;

        Tensor::cat(&[&first_zero, &simple_returns], 0).map_err(|e| {
            MLError::TensorOperationError(format!("Failed to concatenate simple returns: {}", e))
        })?
    };

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

    // Step 3: Clip outliers
    let clipped = clip_outliers(&normalized, config.clip_sigma)?;

    Ok(clipped)
}

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

    #[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 = Tensor::from_slice(&[100.0_f32, 110.0, 105.0], (3,), &Device::Cpu)
            .expect("Failed to create tensor");

        let returns = compute_log_returns(&prices).expect("Failed to compute log returns");
        let returns_vec: Vec<f32> = returns.to_vec1().expect("Failed to convert to vec");

        assert_eq!(returns_vec.len(), 3);
        assert!((returns_vec[0] - 0.0).abs() < 1e-6); // First value is 0
        assert!((returns_vec[1] - (110.0_f32 / 100.0).ln()).abs() < 1e-5); // log(110/100)
    }

    #[test]
    fn test_windowed_normalize_basic() {
        let data = Tensor::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0], (5,), &Device::Cpu)
            .expect("Failed to create tensor");

        let normalized = windowed_normalize(&data, 3).expect("Failed to normalize");
        let normalized_vec: Vec<f32> = normalized.to_vec1().expect("Failed to convert to vec");

        assert_eq!(normalized_vec.len(), 5);
        // All values should be finite
        for val in normalized_vec {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_clip_outliers_basic() {
        // Use data where outliers will actually be clipped with 1.5 sigma
        // Normal data: [10.0, 11.0, 12.0, 13.0, 14.0] with mean ~12.0, std ~1.4
        // Add outliers: 50.0 and -20.0 which are far beyond 1.5 sigma
        let data = Tensor::from_slice(&[10.0_f32, 11.0, 12.0, 13.0, 14.0, 50.0, -20.0], (7,), &Device::Cpu)
            .expect("Failed to create tensor");

        let clipped = clip_outliers(&data, 1.5).expect("Failed to clip");
        let clipped_vec: Vec<f32> = clipped.to_vec1().expect("Failed to convert to vec");

        assert_eq!(clipped_vec.len(), 7);
        // Extreme values should be clipped (50.0 and -20.0 are far beyond 1.5 sigma)
        assert!(clipped_vec[5] < 50.0, "Positive outlier (50.0) should be clipped");
        assert!(clipped_vec[6] > -20.0, "Negative outlier (-20.0) should be clipped");
        // Normal values should be unchanged
        assert!((clipped_vec[0] - 10.0).abs() < 0.1, "Normal value should not change much");
        assert!((clipped_vec[1] - 11.0).abs() < 0.1, "Normal value should not change much");
    }
}