foxhunt/crates/ml-data-validation/src/corrector.rs

//! # Data Corrector
//!
//! Automatic correction of data quality issues including:
//! - Price spike interpolation
//! - Outlier removal/capping
//! - Missing bar interpolation
//!
//! ## Safety
//!
//! All corrections are conservative and preserve data integrity.
//! Original data is never modified in-place.

use crate::OHLCVBar;
use anyhow::Result;

/// Data corrector for automatic fixes
pub struct DataCorrector {
    /// Correction statistics
    corrections_applied: std::sync::atomic::AtomicUsize,
}

impl std::fmt::Debug for DataCorrector {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DataCorrector")
            .field(
                "corrections_applied",
                &self
                    .corrections_applied
                    .load(std::sync::atomic::Ordering::Relaxed),
            )
            .finish()
    }
}

impl DataCorrector {
    /// Create new data corrector
    pub const fn new() -> Self {
        Self {
            corrections_applied: std::sync::atomic::AtomicUsize::new(0),
        }
    }

    /// Correct price spikes using linear interpolation
    ///
    /// When a price spike >threshold is detected, replaces the spiked bar
    /// with interpolated values from surrounding bars.
    ///
    /// # Arguments
    ///
    /// * `bars` - Original OHLCV bars
    /// * `threshold` - Spike threshold (0.20 = 20%)
    ///
    /// # Returns
    ///
    /// Corrected bars with spikes interpolated
    pub fn correct_price_spikes(&self, bars: &[OHLCVBar], threshold: f64) -> Result<Vec<OHLCVBar>> {
        if bars.len() < 3 {
            return Ok(bars.to_vec());
        }

        let mut corrected = bars.to_vec();
        let mut corrections = 0;

        for i in 1..(bars.len() - 1) {
            let prev_close = bars[i - 1].close;
            let curr_open = bars[i].open;
            let next_open = bars[i + 1].open;

            // Calculate percentage changes
            let pct_change_prev = ((curr_open - prev_close) / prev_close).abs();
            let pct_change_next = ((next_open - curr_open) / curr_open).abs();

            // Detect spike: large change in, large change out
            if pct_change_prev > threshold && pct_change_next > threshold {
                // Interpolate the spiked bar
                let interpolated_close = (prev_close + next_open) / 2.0;
                let interpolated_open = interpolated_close * 0.999; // Slightly below close
                let interpolated_high = interpolated_close * 1.002; // Slightly above
                let interpolated_low = interpolated_close * 0.998; // Slightly below

                corrected[i].open = interpolated_open;
                corrected[i].high = interpolated_high;
                corrected[i].low = interpolated_low;
                corrected[i].close = interpolated_close;

                corrections += 1;
            }
        }

        // Update correction counter
        self.corrections_applied
            .fetch_add(corrections, std::sync::atomic::Ordering::Relaxed);

        Ok(corrected)
    }

    /// Remove outliers using z-score method
    ///
    /// Identifies outliers in volume and prices using standard deviations,
    /// then caps them at reasonable values.
    ///
    /// # Arguments
    ///
    /// * `bars` - Original OHLCV bars
    /// * `z_threshold` - Z-score threshold (3.0 = 3 standard deviations)
    ///
    /// # Returns
    ///
    /// Corrected bars with outliers capped
    pub fn remove_outliers(&self, bars: &[OHLCVBar], z_threshold: f64) -> Result<Vec<OHLCVBar>> {
        if bars.is_empty() {
            return Ok(bars.to_vec());
        }

        let mut corrected = bars.to_vec();
        let mut corrections = 0;

        // Calculate volume statistics
        let volumes: Vec<f64> = bars.iter().map(|b| b.volume).collect();
        // Use median and MAD for robust outlier detection (resistant to outliers)
        let vol_median = calculate_median(&volumes);
        let vol_mad = calculate_mad(&volumes, vol_median);

        // Correct volume outliers
        for bar in corrected.iter_mut() {
            // Use modified z-score with MAD: z = 0.6745 * (x - median) / MAD
            // This is more robust to outliers than standard z-score
            let modified_z = if vol_mad > 0.0 {
                0.6745 * (bar.volume - vol_median).abs() / vol_mad
            } else {
                0.0
            };

            if modified_z > z_threshold {
                // Cap volume at median + threshold * MAD (robust capping)
                let max_volume = vol_median + (z_threshold * vol_mad / 0.6745);
                bar.volume = max_volume;
                corrections += 1;
            }
        }

        // Update correction counter
        self.corrections_applied
            .fetch_add(corrections, std::sync::atomic::Ordering::Relaxed);

        Ok(corrected)
    }

    /// Fill missing bars with interpolated values
    ///
    /// Detects gaps in time series and fills them with interpolated OHLCV values.
    ///
    /// # Arguments
    ///
    /// * `bars` - Original OHLCV bars (may have gaps)
    /// * `expected_interval_secs` - Expected interval between bars
    ///
    /// # Returns
    ///
    /// Complete time series with interpolated bars
    pub fn fill_missing_bars(
        &self,
        bars: &[OHLCVBar],
        expected_interval_secs: i64,
    ) -> Result<Vec<OHLCVBar>> {
        if bars.len() < 2 {
            return Ok(bars.to_vec());
        }

        let mut filled = Vec::new();
        filled.push(bars[0]);

        for i in 1..bars.len() {
            let prev = &bars[i - 1];
            let curr = &bars[i];

            let gap_secs = (curr.timestamp - prev.timestamp).num_seconds();
            let missing_bars = (gap_secs / expected_interval_secs) - 1;

            // If gap detected, interpolate missing bars
            if missing_bars > 0 && missing_bars < 10 {
                // Only fill small gaps
                for j in 1..=missing_bars {
                    let ratio = j as f64 / (missing_bars + 1) as f64;
                    let interpolated = interpolate_bar(prev, curr, ratio);
                    filled.push(interpolated);
                }

                // Update correction counter
                self.corrections_applied
                    .fetch_add(missing_bars as usize, std::sync::atomic::Ordering::Relaxed);
            }

            filled.push(*curr);
        }

        Ok(filled)
    }

    /// Get total corrections applied
    pub fn corrections_count(&self) -> usize {
        self.corrections_applied
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Reset correction counter
    pub fn reset_counter(&self) {
        self.corrections_applied
            .store(0, std::sync::atomic::Ordering::Relaxed);
    }
}

impl Default for DataCorrector {
    fn default() -> Self {
        Self::new()
    }
}

// Helper functions

/// Calculate mean and standard deviation
#[cfg_attr(not(test), allow(dead_code))]
fn calculate_mean_std(values: &[f64]) -> (f64, f64) {
    if values.is_empty() {
        return (0.0, 0.0);
    }

    let mean = values.iter().sum::<f64>() / values.len() as f64;

    let variance = values
        .iter()
        .map(|&v| {
            let diff = v - mean;
            diff * diff
        })
        .sum::<f64>()
        / values.len() as f64;

    let std = variance.sqrt();

    (mean, std)
}

/// Calculate median of a set of values
fn calculate_median(values: &[f64]) -> f64 {
    if values.is_empty() {
        return 0.0;
    }

    let mut sorted = values.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let len = sorted.len();
    if len % 2 == 0 {
        (sorted[len / 2 - 1] + sorted[len / 2]) / 2.0
    } else {
        sorted[len / 2]
    }
}

/// Calculate Median Absolute Deviation (MAD)
fn calculate_mad(values: &[f64], median: f64) -> f64 {
    if values.is_empty() {
        return 0.0;
    }

    let deviations: Vec<f64> = values.iter().map(|&v| (v - median).abs()).collect();
    calculate_median(&deviations)
}

/// Interpolate bar between two bars
fn interpolate_bar(prev: &OHLCVBar, next: &OHLCVBar, ratio: f64) -> OHLCVBar {
    // Linear interpolation for timestamp
    let duration = next.timestamp - prev.timestamp;
    let interpolated_duration =
        chrono::Duration::milliseconds((duration.num_milliseconds() as f64 * ratio) as i64);
    let interpolated_timestamp = prev.timestamp + interpolated_duration;

    // Linear interpolation for prices
    let interpolated_close = prev.close + (next.close - prev.close) * ratio;
    let interpolated_open = prev.open + (next.open - prev.open) * ratio;
    let interpolated_high = interpolated_close * 1.001; // Slightly above close
    let interpolated_low = interpolated_close * 0.999; // Slightly below close

    // Average volume
    let interpolated_volume = (prev.volume + next.volume) / 2.0;

    OHLCVBar {
        timestamp: interpolated_timestamp,
        open: interpolated_open,
        high: interpolated_high,
        low: interpolated_low,
        close: interpolated_close,
        volume: interpolated_volume,
    }
}

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

    fn create_test_bar(close: f64, volume: f64) -> OHLCVBar {
        OHLCVBar {
            timestamp: Utc::now(),
            open: close * 0.999,
            high: close * 1.001,
            low: close * 0.999,
            close,
            volume,
        }
    }

    #[test]
    fn test_spike_correction() {
        let corrector = DataCorrector::new();

        let bars = vec![
            create_test_bar(100.0, 1000.0),
            create_test_bar(200.0, 1000.0), // Spike
            create_test_bar(102.0, 1000.0),
        ];

        let corrected = corrector.correct_price_spikes(&bars, 0.20).unwrap();

        // Middle bar should be interpolated
        assert!(corrected[1].close > 100.0);
        assert!(corrected[1].close < 110.0);
        assert_ne!(corrected[1].close, bars[1].close);
    }

    #[test]
    fn test_outlier_removal() {
        let corrector = DataCorrector::new();

        let bars = vec![
            create_test_bar(100.0, 1000.0),
            create_test_bar(101.0, 1100.0),
            create_test_bar(102.0, 50000.0), // Outlier
            create_test_bar(103.0, 1050.0),
        ];

        let corrected = corrector.remove_outliers(&bars, 3.0).unwrap();

        // Outlier volume should be capped
        assert!(corrected[2].volume < 10000.0);
        assert_ne!(corrected[2].volume, bars[2].volume);
    }

    #[test]
    fn test_mean_std_calculation() {
        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let (mean, std) = calculate_mean_std(&values);

        assert!((mean - 3.0).abs() < 0.01);
        assert!(std > 0.0);
    }

    #[test]
    fn test_correction_counter() {
        let corrector = DataCorrector::new();
        assert_eq!(corrector.corrections_count(), 0);

        let bars = vec![
            create_test_bar(100.0, 1000.0),
            create_test_bar(200.0, 1000.0),
            create_test_bar(102.0, 1000.0),
        ];

        let _ = corrector.correct_price_spikes(&bars, 0.20).unwrap();
        assert!(corrector.corrections_count() > 0);
    }
}