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

//! Walk-forward cross-validation splitter with embargo periods.
//!
//! Provides a time-series-aware splitting strategy that slides a
//! `[train | embargo | test]` window across bar indices. The embargo gap
//! between train and test prevents information leakage from lookback
//! features (e.g., moving averages, rolling volatility).
//!
//! # Example
//!
//! ```
//! use ml_validation::walk_forward::{WalkForwardConfig, walk_forward_split};
//!
//! let cfg = WalkForwardConfig {
//!     train_bars: 100,
//!     test_bars: 20,
//!     embargo_bars: 5,
//!     step_bars: 20,
//!     min_train_samples: 50,
//! };
//! let folds = walk_forward_split(200, &cfg);
//! assert!(!folds.is_empty());
//! ```

use serde::{Deserialize, Serialize};
use std::ops::Range;

/// Configuration for walk-forward cross-validation.
///
/// Controls the sizes of the training window, embargo gap, test window,
/// and the step size between successive folds.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalkForwardConfig {
    /// Number of bars in each training window.
    pub train_bars: usize,
    /// Number of bars in each test window.
    pub test_bars: usize,
    /// Number of bars between train end and test start (prevents leakage).
    pub embargo_bars: usize,
    /// How far to advance the window origin each fold.
    pub step_bars: usize,
    /// Minimum number of training bars required to produce a fold.
    pub min_train_samples: usize,
}

impl Default for WalkForwardConfig {
    fn default() -> Self {
        Self {
            train_bars: 1260,  // ~5 years of daily bars
            test_bars: 252,    // ~1 year of daily bars
            embargo_bars: 20,  // ~1 month gap
            step_bars: 126,    // ~6 months advance per fold
            min_train_samples: 252, // ~1 year minimum training data
        }
    }
}

/// A single fold produced by the walk-forward splitter.
///
/// Each fold defines three contiguous, non-overlapping ranges:
/// `[train_range | embargo_range | test_range]`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Fold {
    /// Zero-based index of this fold in the sequence.
    pub fold_index: usize,
    /// Bar range used for training (exclusive end).
    pub train_range: Range<usize>,
    /// Bar range of the embargo gap (exclusive end).
    pub embargo_range: Range<usize>,
    /// Bar range used for testing (exclusive end).
    pub test_range: Range<usize>,
}

/// Generate walk-forward cross-validation folds with embargo.
///
/// Uses a **rolling window** design: each fold has a fixed-size training window
/// that slides forward. This is the standard approach for evaluating strategies
/// on recent data (avoids stale early data diluting training signal).
///
/// An expanding window variant (where training grows from an anchor point) may
/// be added later if needed for strategies that benefit from larger datasets.
///
/// Slides a `[train | embargo | test]` window across `num_bars` bar indices,
/// advancing by `config.step_bars` each iteration. A fold is only emitted when:
/// - The training window has at least `config.min_train_samples` bars
/// - The test window fits entirely within `[0, num_bars)`
///
/// Returns an empty `Vec` if the data is too short for even one fold.
pub fn walk_forward_split(num_bars: usize, config: &WalkForwardConfig) -> Vec<Fold> {
    let mut folds = Vec::new();

    // Total bars consumed by one complete window
    let window_size = config
        .train_bars
        .saturating_add(config.embargo_bars)
        .saturating_add(config.test_bars);

    // Early exit: not enough data for even one fold
    if num_bars < window_size {
        return folds;
    }

    // Also bail if train window is smaller than the minimum
    if config.train_bars < config.min_train_samples {
        return folds;
    }

    let mut train_start: usize = 0;
    let mut fold_index: usize = 0;

    loop {
        let train_end = train_start.saturating_add(config.train_bars);
        let embargo_end = train_end.saturating_add(config.embargo_bars);
        let test_end = embargo_end.saturating_add(config.test_bars);

        // Stop if test window exceeds available data
        if test_end > num_bars {
            break;
        }

        folds.push(Fold {
            fold_index,
            train_range: train_start..train_end,
            embargo_range: train_end..embargo_end,
            test_range: embargo_end..test_end,
        });

        fold_index = fold_index.saturating_add(1);

        // Advance origin by step_bars; use saturating_add to avoid overflow
        train_start = train_start.saturating_add(config.step_bars);
    }

    folds
}

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

    #[test]
    fn test_basic_walk_forward_split() {
        let cfg = WalkForwardConfig {
            train_bars: 100,
            test_bars: 20,
            embargo_bars: 5,
            step_bars: 20,
            min_train_samples: 50,
        };
        let folds = walk_forward_split(300, &cfg);

        // Should produce at least 2 folds
        assert!(
            folds.len() >= 2,
            "Expected at least 2 folds, got {}",
            folds.len()
        );

        // First fold: train [0..100), embargo [100..105), test [105..125)
        let f0 = folds.first().expect("first fold must exist");
        assert_eq!(f0.train_range, 0..100);
        assert_eq!(f0.embargo_range, 100..105);
        assert_eq!(f0.test_range, 105..125);

        // Second fold: train [20..120), embargo [120..125), test [125..145)
        let f1 = folds.get(1).expect("second fold must exist");
        assert_eq!(f1.train_range, 20..120);
        assert_eq!(f1.embargo_range, 120..125);
        assert_eq!(f1.test_range, 125..145);
    }

    #[test]
    fn test_no_overlap_between_train_and_test() {
        let cfg = WalkForwardConfig {
            train_bars: 80,
            test_bars: 30,
            embargo_bars: 10,
            step_bars: 25,
            min_train_samples: 40,
        };
        let folds = walk_forward_split(500, &cfg);
        assert!(!folds.is_empty(), "Expected at least one fold");

        for fold in &folds {
            // Train end == embargo start
            assert_eq!(
                fold.train_range.end, fold.embargo_range.start,
                "Fold {}: train.end ({}) != embargo.start ({})",
                fold.fold_index, fold.train_range.end, fold.embargo_range.start
            );
            // Embargo end == test start
            assert_eq!(
                fold.embargo_range.end, fold.test_range.start,
                "Fold {}: embargo.end ({}) != test.start ({})",
                fold.fold_index, fold.embargo_range.end, fold.test_range.start
            );
            // No overlap: train end <= test start (with embargo in between)
            assert!(
                fold.train_range.end <= fold.test_range.start,
                "Fold {}: train overlaps test ({} > {})",
                fold.fold_index,
                fold.train_range.end,
                fold.test_range.start,
            );
        }
    }

    #[test]
    fn test_temporal_ordering_across_folds() {
        let cfg = WalkForwardConfig {
            train_bars: 60,
            test_bars: 15,
            embargo_bars: 5,
            step_bars: 15,
            min_train_samples: 30,
        };
        let folds = walk_forward_split(300, &cfg);
        assert!(
            folds.len() >= 2,
            "Need at least 2 folds for ordering test, got {}",
            folds.len()
        );

        for pair in folds.windows(2) {
            if let [prev, curr] = pair {
                assert!(
                    curr.test_range.start >= prev.test_range.start,
                    "Fold {} test start ({}) < fold {} test start ({})",
                    curr.fold_index,
                    curr.test_range.start,
                    prev.fold_index,
                    prev.test_range.start,
                );
            }
        }
    }

    #[test]
    fn test_data_too_short_returns_empty() {
        let cfg = WalkForwardConfig {
            train_bars: 100,
            test_bars: 20,
            embargo_bars: 5,
            step_bars: 20,
            min_train_samples: 50,
        };
        // Total window = 100 + 5 + 20 = 125, but only 100 bars available
        let folds = walk_forward_split(100, &cfg);
        assert!(
            folds.is_empty(),
            "Expected empty folds for insufficient data, got {} folds",
            folds.len()
        );
    }

    #[test]
    fn test_fold_indices_are_sequential() {
        let cfg = WalkForwardConfig {
            train_bars: 50,
            test_bars: 10,
            embargo_bars: 3,
            step_bars: 10,
            min_train_samples: 20,
        };
        let folds = walk_forward_split(200, &cfg);
        assert!(!folds.is_empty(), "Expected at least one fold");

        for (i, fold) in folds.iter().enumerate() {
            assert_eq!(
                fold.fold_index, i,
                "Fold index mismatch: expected {}, got {}",
                i, fold.fold_index
            );
        }
    }

    #[test]
    fn test_test_range_within_bounds() {
        let num_bars = 250;
        let cfg = WalkForwardConfig {
            train_bars: 80,
            test_bars: 20,
            embargo_bars: 5,
            step_bars: 15,
            min_train_samples: 40,
        };
        let folds = walk_forward_split(num_bars, &cfg);
        assert!(!folds.is_empty(), "Expected at least one fold");

        for fold in &folds {
            assert!(
                fold.test_range.end <= num_bars,
                "Fold {} test_range.end ({}) exceeds num_bars ({})",
                fold.fold_index,
                fold.test_range.end,
                num_bars,
            );
        }
    }

    #[test]
    fn test_default_config() {
        let cfg = WalkForwardConfig::default();
        assert_eq!(cfg.train_bars, 1260, "Default train_bars should be ~5yr daily");
        assert_eq!(cfg.test_bars, 252, "Default test_bars should be ~1yr daily");
        assert_eq!(cfg.embargo_bars, 20, "Default embargo_bars should be ~1mo");
        assert_eq!(cfg.step_bars, 126, "Default step_bars should be ~6mo");
        assert_eq!(cfg.min_train_samples, 252, "Default min_train_samples should be ~1yr");
    }
}