foxhunt/crates/ml-features/src/feature_extraction.rs

//! Feature extraction for ML models
//!
//! This module extracts 256-dimensional feature vectors from OHLCV bars.
//! Phase 1: Implements 15 core features (5 OHLCV + 10 technical indicators)
//! Phase 2-4: Will add 241 additional engineered features

pub use crate::OHLCVBar;
use crate::MLError;
#[cfg(test)]
use chrono::Utc;

/// Feature extractor with technical indicators
#[derive(Debug)]
pub struct FeatureExtractor {
    /// RSI period (default 14)
    rsi_period: usize,
    /// EMA fast period (default 12)
    ema_fast_period: usize,
    /// EMA slow period (default 26)
    ema_slow_period: usize,
    /// Bollinger Bands period (default 20)
    bb_period: usize,
    /// Bollinger Bands std dev multiplier
    bb_std_dev: f64,
    /// ATR period (default 14)
    atr_period: usize,
}

impl Default for FeatureExtractor {
    fn default() -> Self {
        Self {
            rsi_period: 14,
            ema_fast_period: 12,
            ema_slow_period: 26,
            bb_period: 20,
            bb_std_dev: 2.0,
            atr_period: 14,
        }
    }
}

impl FeatureExtractor {
    /// Create new feature extractor with default parameters
    pub fn new() -> Self {
        Self::default()
    }

    /// Extract 15 core features from OHLCV bars
    ///
    /// Features:
    /// 1-5: OHLCV (open, high, low, close, volume)
    /// 6: RSI (Relative Strength Index)
    /// 7-8: EMA fast, slow
    /// 9-11: MACD (line, signal, histogram)
    /// 12-14: Bollinger Bands (upper, middle, lower)
    /// 15: ATR (Average True Range)
    pub fn extract_features(&self, bars: &[OHLCVBar]) -> Result<Vec<Vec<f32>>, MLError> {
        if bars.is_empty() {
            return Err(MLError::InsufficientData(
                "Cannot extract features from empty bars".to_owned(),
            ));
        }

        let min_required = self
            .bb_period
            .max(self.ema_slow_period)
            .max(self.rsi_period);
        if bars.len() < min_required {
            return Err(MLError::InsufficientData(format!(
                "Need at least {} bars for feature extraction, got {}",
                min_required,
                bars.len()
            )));
        }

        let mut features = Vec::new();

        // Calculate technical indicators
        let rsi_values = self.calculate_rsi(bars);
        let ema_fast = self.calculate_ema(bars, self.ema_fast_period);
        let ema_slow = self.calculate_ema(bars, self.ema_slow_period);
        let (macd_line, macd_signal, macd_hist) = self.calculate_macd(bars);
        let (bb_upper, bb_middle, bb_lower) = self.calculate_bollinger_bands(bars);
        let atr = self.calculate_atr(bars);

        // Extract features for each bar
        for (i, bar) in bars.iter().enumerate() {
            let feature_vec = vec![
                // Features 1-5: OHLCV (normalized)
                bar.open as f32,
                bar.high as f32,
                bar.low as f32,
                bar.close as f32,
                bar.volume as f32,
                // Feature 6: RSI
                rsi_values.get(i).copied().unwrap_or(50.0) as f32,
                // Features 7-8: EMA
                ema_fast.get(i).copied().unwrap_or(bar.close) as f32,
                ema_slow.get(i).copied().unwrap_or(bar.close) as f32,
                // Features 9-11: MACD
                macd_line.get(i).copied().unwrap_or(0.0) as f32,
                macd_signal.get(i).copied().unwrap_or(0.0) as f32,
                macd_hist.get(i).copied().unwrap_or(0.0) as f32,
                // Features 12-14: Bollinger Bands
                bb_upper.get(i).copied().unwrap_or(bar.high) as f32,
                bb_middle.get(i).copied().unwrap_or(bar.close) as f32,
                bb_lower.get(i).copied().unwrap_or(bar.low) as f32,
                // Feature 15: ATR
                atr.get(i).copied().unwrap_or(0.0) as f32,
            ];

            features.push(feature_vec);
        }

        Ok(features)
    }

    /// Calculate RSI (Relative Strength Index)
    fn calculate_rsi(&self, bars: &[OHLCVBar]) -> Vec<f64> {
        let mut rsi_values = Vec::with_capacity(bars.len());
        let period = self.rsi_period;

        if bars.len() < period + 1 {
            return vec![50.0; bars.len()];
        }

        let mut gains = Vec::new();
        let mut losses = Vec::new();

        // Calculate price changes
        for i in 1..bars.len() {
            let change = bars[i].close - bars[i - 1].close;
            if change > 0.0 {
                gains.push(change);
                losses.push(0.0);
            } else {
                gains.push(0.0);
                losses.push(change.abs());
            }
        }

        // Calculate RSI
        for i in 0..bars.len() {
            if i < period {
                rsi_values.push(50.0); // Neutral value for warmup
                continue;
            }

            let start_idx = i.saturating_sub(period);
            let avg_gain: f64 = gains[start_idx..i].iter().sum::<f64>() / period as f64;
            let avg_loss: f64 = losses[start_idx..i].iter().sum::<f64>() / period as f64;

            let rsi = if avg_loss == 0.0 {
                100.0
            } else {
                let rs = avg_gain / avg_loss;
                100.0 - (100.0 / (1.0 + rs))
            };

            rsi_values.push(rsi);
        }

        rsi_values
    }

    /// Calculate EMA (Exponential Moving Average)
    fn calculate_ema(&self, bars: &[OHLCVBar], period: usize) -> Vec<f64> {
        let mut ema_values = Vec::with_capacity(bars.len());
        if bars.is_empty() {
            return ema_values;
        }

        let multiplier = 2.0 / (period as f64 + 1.0);
        let mut ema = bars[0].close;
        ema_values.push(ema);

        for bar in bars.iter().skip(1) {
            ema = (bar.close - ema) * multiplier + ema;
            ema_values.push(ema);
        }

        ema_values
    }

    /// Calculate MACD (Moving Average Convergence Divergence)
    fn calculate_macd(&self, bars: &[OHLCVBar]) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
        let ema_fast = self.calculate_ema(bars, 12);
        let ema_slow = self.calculate_ema(bars, 26);

        let mut macd_line = Vec::with_capacity(bars.len());
        for i in 0..bars.len() {
            macd_line.push(ema_fast[i] - ema_slow[i]);
        }

        // Calculate signal line (9-period EMA of MACD line)
        let signal_period = 9;
        let multiplier = 2.0 / (signal_period as f64 + 1.0);
        let mut signal = Vec::with_capacity(bars.len());

        if !macd_line.is_empty() {
            let mut ema = macd_line[0];
            signal.push(ema);

            for &macd_val in macd_line.iter().skip(1) {
                ema = (macd_val - ema) * multiplier + ema;
                signal.push(ema);
            }
        }

        // Calculate histogram
        let histogram: Vec<f64> = macd_line
            .iter()
            .zip(signal.iter())
            .map(|(m, s)| m - s)
            .collect();

        (macd_line, signal, histogram)
    }

    /// Calculate Bollinger Bands
    fn calculate_bollinger_bands(&self, bars: &[OHLCVBar]) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
        let period = self.bb_period;
        let std_dev = self.bb_std_dev;

        let mut upper = Vec::with_capacity(bars.len());
        let mut middle = Vec::with_capacity(bars.len());
        let mut lower = Vec::with_capacity(bars.len());

        for i in 0..bars.len() {
            if i < period {
                middle.push(bars[i].close);
                upper.push(bars[i].high);
                lower.push(bars[i].low);
                continue;
            }

            let start_idx = i - period;
            let prices: Vec<f64> = bars[start_idx..=i].iter().map(|b| b.close).collect();

            let sma: f64 = prices.iter().sum::<f64>() / prices.len() as f64;
            let variance: f64 =
                prices.iter().map(|&p| (p - sma).powi(2)).sum::<f64>() / prices.len() as f64;
            let std = variance.sqrt();

            middle.push(sma);
            upper.push(sma + std_dev * std);
            lower.push(sma - std_dev * std);
        }

        (upper, middle, lower)
    }

    /// Calculate ATR (Average True Range)
    fn calculate_atr(&self, bars: &[OHLCVBar]) -> Vec<f64> {
        let mut atr_values = Vec::with_capacity(bars.len());
        let period = self.atr_period;

        if bars.is_empty() {
            return atr_values;
        }

        let mut true_ranges = Vec::new();
        atr_values.push(0.0); // First bar has no ATR

        for i in 1..bars.len() {
            let high_low = bars[i].high - bars[i].low;
            let high_close = (bars[i].high - bars[i - 1].close).abs();
            let low_close = (bars[i].low - bars[i - 1].close).abs();

            let tr = high_low.max(high_close).max(low_close);
            true_ranges.push(tr);

            if i < period {
                atr_values.push(0.0);
            } else {
                let start_idx = i.saturating_sub(period);
                let atr: f64 = true_ranges[start_idx..].iter().sum::<f64>()
                    / (true_ranges.len() - start_idx) as f64;
                atr_values.push(atr);
            }
        }

        atr_values
    }
}

/// Extract ML features from OHLCV bars (convenience function)
pub fn extract_ml_features(bars: &[OHLCVBar]) -> Result<Vec<Vec<f32>>, MLError> {
    let extractor = FeatureExtractor::new();
    extractor.extract_features(bars)
}

/// Compute ATR (Average True Range) for a slice of bars
///
/// This is a public helper function for other modules (e.g., `regime_adaptive.rs`)
/// that need ATR calculation.
///
/// ## Arguments
/// - `bars`: Slice of OHLCV bars
/// - `period`: ATR period (typically 14)
///
/// ## Returns
/// - ATR value for the most recent bar, or 0.0 if insufficient data
///
/// ## Example
/// ```rust
/// use ml::features::feature_extraction::{compute_atr, OHLCVBar};
///
/// let bars = vec![/* OHLCV bars */];
/// let atr = compute_atr(&bars, 14);
/// ```
pub fn compute_atr(bars: &[OHLCVBar], period: usize) -> f64 {
    if bars.len() < period + 1 {
        return 0.0;
    }

    let mut true_ranges = Vec::with_capacity(bars.len() - 1);

    // Calculate true range for each bar (starting from bar 1)
    for i in 1..bars.len() {
        let high_low = bars[i].high - bars[i].low;
        let high_close = (bars[i].high - bars[i - 1].close).abs();
        let low_close = (bars[i].low - bars[i - 1].close).abs();

        let tr = high_low.max(high_close).max(low_close);
        true_ranges.push(tr);
    }

    // Calculate ATR for the most recent period
    let start_idx = true_ranges.len().saturating_sub(period);
    let atr: f64 =
        true_ranges[start_idx..].iter().sum::<f64>() / period.min(true_ranges.len()) as f64;

    atr
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states, clippy::manual_range_contains)]
mod tests {
    use super::*;

    fn create_test_bars(count: usize) -> Vec<OHLCVBar> {
        let base_time = Utc::now();
        (0..count)
            .map(|i| OHLCVBar {
                timestamp: base_time + chrono::Duration::seconds(i as i64),
                open: 100.0 + i as f64,
                high: 105.0 + i as f64,
                low: 95.0 + i as f64,
                close: 102.0 + i as f64,
                volume: 1000.0 + i as f64 * 10.0,
            })
            .collect()
    }

    #[test]
    fn test_feature_extraction() {
        let bars = create_test_bars(100);
        let features = extract_ml_features(&bars).unwrap();

        assert_eq!(features.len(), bars.len());
        for feature_vec in &features {
            assert_eq!(feature_vec.len(), 15, "Expected 15 features per bar");

            // Check for finite values
            for &value in feature_vec {
                assert!(value.is_finite(), "Feature should be finite");
            }
        }
    }

    #[test]
    fn test_insufficient_data() {
        let bars = create_test_bars(5);
        let result = extract_ml_features(&bars);
        assert!(result.is_err());
    }

    #[test]
    fn test_rsi_calculation() {
        let extractor = FeatureExtractor::new();
        let bars = create_test_bars(50);
        let rsi = extractor.calculate_rsi(&bars);

        assert_eq!(rsi.len(), bars.len());
        for &value in &rsi {
            assert!(value >= 0.0 && value <= 100.0, "RSI should be in [0, 100]");
        }
    }

    #[test]
    fn test_ema_calculation() {
        let extractor = FeatureExtractor::new();
        let bars = create_test_bars(50);
        let ema = extractor.calculate_ema(&bars, 12);

        assert_eq!(ema.len(), bars.len());
        for &value in &ema {
            assert!(value.is_finite(), "EMA should be finite");
        }
    }
}