foxhunt/crates/ml-validation/src/temporal_guard.rs

//! Temporal guard for preventing forward-looking data leakage.
//!
//! Wraps a [`TimeSeriesData`] reference with a cutoff index that partitions
//! the series into a training region `[0, cutoff)` and a test region
//! `[cutoff, len)`.  Any attempt to slice across the boundary is rejected,
//! ensuring that training code never accidentally sees future data.

use chrono::{DateTime, Utc};

use crate::types::TimeSeriesData;
use crate::MLError;

// ---------------------------------------------------------------------------
// LeakageAuditReport
// ---------------------------------------------------------------------------

/// Result of a forward-leakage audit on a [`TemporalGuard`].
#[derive(Debug, Clone)]
pub struct LeakageAuditReport {
    /// `true` if any training timestamp is >= any test timestamp.
    pub has_future_timestamps: bool,
    /// Number of bars in the training partition.
    pub training_bars: usize,
    /// The timestamp at the cutoff index, if available.
    pub cutoff_timestamp: Option<DateTime<Utc>>,
}

// ---------------------------------------------------------------------------
// NormalizationStats
// ---------------------------------------------------------------------------

/// Feature-wise normalization statistics computed from training data only.
#[derive(Debug, Clone)]
pub struct NormalizationStats {
    /// Per-feature means.
    pub means: Vec<f64>,
    /// Per-feature standard deviations.
    pub stds: Vec<f64>,
    /// Number of samples (bars) used.
    pub sample_count: usize,
}

// ---------------------------------------------------------------------------
// TemporalGuard
// ---------------------------------------------------------------------------

/// Prevents forward-looking data leakage by enforcing a temporal cutoff.
///
/// The guard partitions the underlying [`TimeSeriesData`] into:
/// - **Training region**: `[0, cutoff_idx)`
/// - **Test region**: `[cutoff_idx, data.len())`
///
/// Slicing across the boundary is an error.
#[derive(Debug)]
pub struct TemporalGuard<'a> {
    data: &'a TimeSeriesData,
    cutoff_idx: usize,
}

impl<'a> TemporalGuard<'a> {
    /// Create a new `TemporalGuard`.
    ///
    /// # Errors
    ///
    /// Returns [`MLError::InvalidInput`] if `cutoff_idx` is 0, less than 2
    /// (training needs at least 2 bars), or exceeds `data.len()`.
    pub fn new(data: &'a TimeSeriesData, cutoff_idx: usize) -> Result<Self, MLError> {
        if cutoff_idx < 2 {
            return Err(MLError::InvalidInput(format!(
                "TemporalGuard cutoff must be >= 2 for valid training slice, got {}",
                cutoff_idx,
            )));
        }
        if cutoff_idx > data.len() {
            return Err(MLError::InvalidInput(format!(
                "TemporalGuard cutoff ({}) exceeds data length ({})",
                cutoff_idx,
                data.len(),
            )));
        }
        Ok(Self { data, cutoff_idx })
    }

    /// Return the training partition `[0, cutoff)` as a new `TimeSeriesData`.
    ///
    /// # Errors
    ///
    /// Propagates errors from [`TimeSeriesData::slice`].
    pub fn training_slice(&self) -> Result<TimeSeriesData, MLError> {
        self.data.slice(0, self.cutoff_idx)
    }

    /// Return a test sub-range `[start, end)`.
    ///
    /// Both `start` and `end` must be at or after the cutoff index.
    ///
    /// # Errors
    ///
    /// Returns [`MLError::InvalidInput`] if `start < cutoff_idx`, or
    /// propagates errors from [`TimeSeriesData::slice`].
    pub fn test_slice(&self, start: usize, end: usize) -> Result<TimeSeriesData, MLError> {
        if start < self.cutoff_idx {
            return Err(MLError::InvalidInput(format!(
                "test_slice start ({}) is before cutoff ({}); this would leak training data",
                start, self.cutoff_idx,
            )));
        }
        self.data.slice(start, end)
    }

    /// Slice the data within a single partition.
    ///
    /// The range `[start, end)` must lie entirely within `[0, cutoff)` or
    /// entirely within `[cutoff, len)`.  Crossing the boundary is an error.
    ///
    /// # Errors
    ///
    /// Returns [`MLError::InvalidInput`] if the range crosses the cutoff
    /// boundary, or propagates errors from [`TimeSeriesData::slice`].
    pub fn slice(&self, start: usize, end: usize) -> Result<TimeSeriesData, MLError> {
        if start >= end {
            return Err(MLError::InvalidInput(format!(
                "Invalid slice range: start ({}) must be less than end ({})",
                start, end,
            )));
        }
        // Entirely in training region
        let in_train = end <= self.cutoff_idx;
        // Entirely in test region
        let in_test = start >= self.cutoff_idx;

        if !in_train && !in_test {
            return Err(MLError::InvalidInput(format!(
                "Slice [{}, {}) crosses the temporal cutoff at index {}; \
                 this would leak future data into the training partition",
                start, end, self.cutoff_idx,
            )));
        }
        self.data.slice(start, end)
    }

    /// Audit for timestamp-based forward leakage.
    ///
    /// Checks whether any training timestamp is >= the earliest test
    /// timestamp, which would indicate temporal ordering issues.
    pub fn audit_leakage(&self) -> LeakageAuditReport {
        let cutoff_timestamp = self.data.timestamps.get(self.cutoff_idx).copied();

        let has_future_timestamps = if self.cutoff_idx < self.data.len() {
            // Find the minimum test timestamp
            let min_test_ts = self
                .data
                .timestamps
                .get(self.cutoff_idx..)
                .unwrap_or_default()
                .iter()
                .min()
                .copied();

            // Find the maximum training timestamp
            let max_train_ts = self
                .data
                .timestamps
                .get(..self.cutoff_idx)
                .unwrap_or_default()
                .iter()
                .max()
                .copied();

            match (max_train_ts, min_test_ts) {
                (Some(train_max), Some(test_min)) => train_max >= test_min,
                _ => false,
            }
        } else {
            false
        };

        LeakageAuditReport {
            has_future_timestamps,
            training_bars: self.cutoff_idx,
            cutoff_timestamp,
        }
    }

    /// Compute per-feature normalization statistics from training data only.
    ///
    /// Returns means and standard deviations for each feature dimension,
    /// computed exclusively from `[0, cutoff)`.  These can then be applied
    /// to the test partition without leaking future information.
    pub fn compute_normalization_stats(&self) -> NormalizationStats {
        let train_features = self
            .data
            .features
            .get(..self.cutoff_idx)
            .unwrap_or_default();

        let sample_count = train_features.len();
        if sample_count == 0 {
            return NormalizationStats {
                means: Vec::new(),
                stds: Vec::new(),
                sample_count: 0,
            };
        }

        // Determine feature dimension from first row
        let dim = train_features
            .first()
            .map(|r| r.len())
            .unwrap_or_default();

        let mut means = vec![0.0_f64; dim];
        let mut sq_sums = vec![0.0_f64; dim];

        for row in train_features {
            for (j, val) in row.iter().enumerate() {
                if let Some(m) = means.get_mut(j) {
                    *m += *val as f64;
                }
                if let Some(s) = sq_sums.get_mut(j) {
                    *s += (*val as f64) * (*val as f64);
                }
            }
        }

        let n = sample_count as f64;
        for m in &mut means {
            *m /= n;
        }

        let stds: Vec<f64> = means
            .iter()
            .zip(sq_sums.iter())
            .map(|(&mean, &sq_sum)| {
                let variance = (sq_sum / n) - mean * mean;
                if variance > 0.0 {
                    variance.sqrt()
                } else {
                    0.0
                }
            })
            .collect();

        NormalizationStats {
            means,
            stds,
            sample_count,
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Helper: create N UTC timestamps starting from 2024-01-01.
    fn make_timestamps(n: usize) -> Vec<DateTime<Utc>> {
        (0..n)
            .map(|i| {
                Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0)
                    .single()
                    .unwrap_or_else(Utc::now)
                    + chrono::Duration::days(i as i64)
            })
            .collect()
    }

    /// Helper: create N feature rows with a deterministic pattern.
    fn make_features(n: usize, dim: usize) -> Vec<Vec<f32>> {
        (0..n)
            .map(|i| (0..dim).map(|j| (i * dim + j) as f32).collect())
            .collect()
    }

    /// Helper: create a simple ascending-price time series.
    fn make_test_data(n: usize) -> Result<TimeSeriesData, Box<dyn std::error::Error>> {
        let prices: Vec<f64> = (0..n).map(|i| 100.0 + i as f64).collect();
        Ok(TimeSeriesData::new(make_timestamps(n), make_features(n, 3), prices)?)
    }

    #[test]
    fn test_training_slice_returns_correct_range() -> Result<(), Box<dyn std::error::Error>> {
        let data = make_test_data(10)?;
        let guard = TemporalGuard::new(&data, 5)?;
        let train = guard.training_slice()?;

        assert_eq!(train.len(), 5);
        // First price should be 100.0, last should be 104.0
        let first = train.prices.first().copied().unwrap_or(0.0);
        let last = train.prices.last().copied().unwrap_or(0.0);
        assert!((first - 100.0).abs() < 1e-12);
        assert!((last - 104.0).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_test_slice_rejects_before_cutoff() -> Result<(), Box<dyn std::error::Error>> {
        let data = make_test_data(10)?;
        let guard = TemporalGuard::new(&data, 5)?;

        // start=3 is before cutoff=5 → should fail
        let result = guard.test_slice(3, 8);
        assert!(result.is_err());
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("leak"),
            "Expected 'leak' in error, got: {err_msg}",
        );
        Ok(())
    }

    #[test]
    fn test_slice_rejects_cross_boundary() -> Result<(), Box<dyn std::error::Error>> {
        let data = make_test_data(10)?;
        let guard = TemporalGuard::new(&data, 5)?;

        // [3, 7) crosses cutoff=5
        let result = guard.slice(3, 7);
        assert!(result.is_err());
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("crosses"),
            "Expected 'crosses' in error, got: {err_msg}",
        );

        // [0, 4) is entirely in training → should succeed
        let train_ok = guard.slice(0, 4);
        assert!(train_ok.is_ok());

        // [5, 8) is entirely in test → should succeed
        let test_ok = guard.slice(5, 8);
        assert!(test_ok.is_ok());
        Ok(())
    }

    #[test]
    fn test_audit_detects_no_leakage_in_sorted_data() -> Result<(), Box<dyn std::error::Error>> {
        let data = make_test_data(10)?;
        let guard = TemporalGuard::new(&data, 5)?;

        let report = guard.audit_leakage();
        assert!(!report.has_future_timestamps);
        assert_eq!(report.training_bars, 5);
        assert!(report.cutoff_timestamp.is_some());
        Ok(())
    }

    #[test]
    fn test_normalization_stats_from_training_only() -> Result<(), Box<dyn std::error::Error>> {
        let data = make_test_data(10)?;
        let guard = TemporalGuard::new(&data, 5)?;

        let stats = guard.compute_normalization_stats();
        assert_eq!(stats.sample_count, 5);
        assert_eq!(stats.means.len(), 3);
        assert_eq!(stats.stds.len(), 3);

        // Verify means are computed from training features only (rows 0..5)
        // Feature column 0: values 0, 3, 6, 9, 12 → mean = 6.0
        let expected_mean_0 = (0.0 + 3.0 + 6.0 + 9.0 + 12.0) / 5.0;
        let got_mean_0 = stats.means.first().copied().unwrap_or(f64::NAN);
        assert!(
            (got_mean_0 - expected_mean_0).abs() < 1e-10,
            "Expected mean_0={expected_mean_0}, got {got_mean_0}",
        );
        Ok(())
    }
}