foxhunt/crates/ml/src/walk_forward.rs

//! Walk-forward evaluation framework for time-series ML models.
//!
//! Implements expanding-window walk-forward cross-validation with configurable
//! train/val/test splits and feature normalization. This prevents lookahead bias
//! by ensuring models are always evaluated on unseen future data.
//!
//! # Walk-Forward Split Diagram
//!
//! ```text
//! Fold 0: |------- Train (12mo) -------|-- Val (3mo) --|-- Test (3mo) --|
//! Fold 1: |---------- Train (15mo) ----------|-- Val --|-- Test --|
//! Fold 2: |------------- Train (18mo) --------------|-- Val --|-- Test --|
//! ```
//!
//! # Usage
//!
//! ```rust,no_run
//! use ml::walk_forward::{WalkForwardConfig, generate_walk_forward_windows, NormStats};
//!
//! let config = WalkForwardConfig::default();
//! let windows = generate_walk_forward_windows(&bars, &config);
//!
//! for window in &windows {
//!     let features = extract_ml_features(&window.train)?;
//!     let stats = NormStats::from_features(&features);
//!     let normalized = stats.normalize_batch(&features);
//!     // Train model on normalized features...
//! }
//! ```

use crate::types::OHLCVBar;
use chrono::{Months, NaiveDate};

/// Configuration for walk-forward evaluation window generation.
#[derive(Debug, Clone)]
pub struct WalkForwardConfig {
    /// Number of months in the initial training window (default: 12).
    pub initial_train_months: u32,
    /// Number of months in the validation window (default: 3).
    pub val_months: u32,
    /// Number of months in the test window (default: 3).
    pub test_months: u32,
    /// Number of months to step forward between folds (default: 3).
    pub step_months: u32,
}

impl Default for WalkForwardConfig {
    fn default() -> Self {
        Self {
            initial_train_months: 12,
            val_months: 3,
            test_months: 3,
            step_months: 3,
        }
    }
}

/// A single walk-forward evaluation window containing train/val/test splits.
#[derive(Debug, Clone)]
pub struct WalkForwardWindow {
    /// Zero-based fold index.
    pub fold: usize,
    /// Training bars (expanding window).
    pub train: Vec<OHLCVBar>,
    /// Validation bars.
    pub val: Vec<OHLCVBar>,
    /// Test bars.
    pub test: Vec<OHLCVBar>,
    /// End date of the training period (exclusive boundary).
    pub train_end: NaiveDate,
    /// End date of the validation period (exclusive boundary).
    pub val_end: NaiveDate,
    /// End date of the test period (exclusive boundary).
    pub test_end: NaiveDate,
    /// Difficulty score for curriculum learning (mean ADX over training bars).
    /// Low ADX (<20) = easy trending market, high ADX fluctuation = harder.
    /// Computed via [`compute_difficulty`] after window generation.
    pub difficulty_score: f64,
}

/// Curriculum phase for difficulty-based walk-forward window filtering.
///
/// Controls which windows are included in training based on their ADX-derived
/// difficulty score. Progression: Easy -> Mixed -> Full mirrors the standard
/// curriculum learning approach of starting with simple examples.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DifficultyPhase {
    /// Only windows with mean ADX > 30 (clear trending markets, easier to learn).
    Easy,
    /// All windows included, but trending windows (ADX > 30) are weighted 2x
    /// via duplication in the returned list.
    Mixed,
    /// All windows with equal weight (full training, no filtering).
    Full,
}

impl DifficultyPhase {
    /// Create a `DifficultyPhase` from a numeric index.
    ///
    /// * `0` => `Easy`
    /// * `1` => `Mixed`
    /// * `2` (or any other value) => `Full`
    pub fn from_index(idx: usize) -> Self {
        match idx {
            0 => Self::Easy,
            1 => Self::Mixed,
            _ => Self::Full,
        }
    }
}

/// ADX threshold used to classify a window as "trending" (easy) for curriculum learning.
const ADX_TRENDING_THRESHOLD: f64 = 30.0;

/// Compute a difficulty score for a slice of OHLCV bars based on mean ADX(14).
///
/// ADX (Average Directional Index) measures trend strength on a 0-100 scale.
/// High ADX (>30) indicates a clear trend which is easier for RL agents to exploit.
/// Low or fluctuating ADX indicates choppy, range-bound markets that are harder.
///
/// Returns the mean ADX value over the bar slice. If the slice has fewer than
/// 15 bars (minimum for ADX(14) calculation), returns 0.0.
///
/// # Algorithm
///
/// 1. Compute True Range (TR) for each bar
/// 2. Compute +DM / -DM directional movement
/// 3. Smooth TR, +DM, -DM with 14-period Wilder EMA
/// 4. Compute +DI / -DI and DX
/// 5. Smooth DX with 14-period Wilder EMA to get ADX
/// 6. Return mean ADX over valid bars
pub fn compute_difficulty(bars: &[OHLCVBar]) -> f64 {
    let period: usize = 14;
    // Need at least period + 1 bars for one ADX value
    if bars.len() < period.saturating_add(1) {
        return 0.0;
    }

    // Step 1-2: Compute True Range and Directional Movement
    let mut tr_values = Vec::with_capacity(bars.len().saturating_sub(1));
    let mut plus_dm_values = Vec::with_capacity(bars.len().saturating_sub(1));
    let mut minus_dm_values = Vec::with_capacity(bars.len().saturating_sub(1));

    for i in 1..bars.len() {
        let prev = match bars.get(i.wrapping_sub(1)) {
            Some(b) => b,
            None => continue,
        };
        let curr = match bars.get(i) {
            Some(b) => b,
            None => continue,
        };

        // True Range = max(H-L, |H-prev_close|, |L-prev_close|)
        let hl = curr.high - curr.low;
        let hpc = (curr.high - prev.close).abs();
        let lpc = (curr.low - prev.close).abs();
        let tr = hl.max(hpc).max(lpc);
        tr_values.push(tr);

        // +DM = max(H - prev_H, 0) if > max(prev_L - L, 0), else 0
        let up_move = curr.high - prev.high;
        let down_move = prev.low - curr.low;

        let plus_dm = if up_move > down_move && up_move > 0.0 {
            up_move
        } else {
            0.0
        };
        let minus_dm = if down_move > up_move && down_move > 0.0 {
            down_move
        } else {
            0.0
        };

        plus_dm_values.push(plus_dm);
        minus_dm_values.push(minus_dm);
    }

    if tr_values.len() < period {
        return 0.0;
    }

    // Step 3: Wilder smoothing (initial = sum of first `period` values, then EMA)
    let mut smoothed_tr: f64 = tr_values.iter().take(period).sum();
    let mut smoothed_plus_dm: f64 = plus_dm_values.iter().take(period).sum();
    let mut smoothed_minus_dm: f64 = minus_dm_values.iter().take(period).sum();

    let period_f64 = period as f64;
    let mut dx_values = Vec::with_capacity(tr_values.len().saturating_sub(period));

    // First DX from initial smoothed values
    if smoothed_tr.abs() > 1e-12 {
        let plus_di = 100.0 * smoothed_plus_dm / smoothed_tr;
        let minus_di = 100.0 * smoothed_minus_dm / smoothed_tr;
        let di_sum = plus_di + minus_di;
        if di_sum.abs() > 1e-12 {
            let dx = 100.0 * (plus_di - minus_di).abs() / di_sum;
            dx_values.push(dx);
        }
    }

    // Subsequent smoothed values using Wilder's method
    for i in period..tr_values.len() {
        let tr_val = match tr_values.get(i) {
            Some(&v) => v,
            None => continue,
        };
        let pdm_val = match plus_dm_values.get(i) {
            Some(&v) => v,
            None => continue,
        };
        let mdm_val = match minus_dm_values.get(i) {
            Some(&v) => v,
            None => continue,
        };

        smoothed_tr = smoothed_tr - (smoothed_tr / period_f64) + tr_val;
        smoothed_plus_dm = smoothed_plus_dm - (smoothed_plus_dm / period_f64) + pdm_val;
        smoothed_minus_dm = smoothed_minus_dm - (smoothed_minus_dm / period_f64) + mdm_val;

        if smoothed_tr.abs() > 1e-12 {
            let plus_di = 100.0 * smoothed_plus_dm / smoothed_tr;
            let minus_di = 100.0 * smoothed_minus_dm / smoothed_tr;
            let di_sum = plus_di + minus_di;
            if di_sum.abs() > 1e-12 {
                let dx = 100.0 * (plus_di - minus_di).abs() / di_sum;
                dx_values.push(dx);
            }
        }
    }

    if dx_values.len() < period {
        return 0.0;
    }

    // Step 5: Smooth DX with Wilder EMA to get ADX
    let mut adx: f64 = dx_values.iter().take(period).sum::<f64>() / period_f64;
    let mut adx_values = vec![adx];

    for i in period..dx_values.len() {
        let dx_val = match dx_values.get(i) {
            Some(&v) => v,
            None => continue,
        };
        adx = (adx * (period_f64 - 1.0) + dx_val) / period_f64;
        adx_values.push(adx);
    }

    // Step 6: Return mean ADX
    if adx_values.is_empty() {
        return 0.0;
    }
    let sum: f64 = adx_values.iter().sum();
    sum / adx_values.len() as f64
}

/// Filter walk-forward windows by curriculum difficulty phase.
///
/// This function applies difficulty-based filtering to a set of walk-forward
/// windows, enabling curriculum learning by controlling which market conditions
/// the agent trains on.
///
/// # Phases
///
/// * [`DifficultyPhase::Easy`] — Only windows with ADX > 30 (trending, easier markets).
/// * [`DifficultyPhase::Mixed`] — All windows, but trending windows are duplicated for 2x weight.
/// * [`DifficultyPhase::Full`] — All windows with equal weight (no filtering).
///
/// # Arguments
///
/// * `windows` — Walk-forward windows with pre-computed `difficulty_score`.
/// * `phase` — The curriculum phase to apply.
pub fn filter_by_difficulty(
    windows: &[WalkForwardWindow],
    phase: DifficultyPhase,
) -> Vec<WalkForwardWindow> {
    match phase {
        DifficultyPhase::Easy => {
            // Only trending windows (ADX > 30)
            windows
                .iter()
                .filter(|w| w.difficulty_score > ADX_TRENDING_THRESHOLD)
                .cloned()
                .collect()
        }
        DifficultyPhase::Mixed => {
            // All windows, trending ones duplicated for 2x weight
            let mut result = Vec::with_capacity(windows.len().saturating_mul(2));
            for w in windows {
                result.push(w.clone());
                if w.difficulty_score > ADX_TRENDING_THRESHOLD {
                    result.push(w.clone());
                }
            }
            result
        }
        DifficultyPhase::Full => {
            // All windows, equal weight
            windows.to_vec()
        }
    }
}

/// Minimum number of bars required in each split to form a valid window.
const MIN_BARS_PER_SPLIT: usize = 50;

/// Generate expanding walk-forward windows from sorted OHLCV bars.
///
/// Each successive fold extends the training window while keeping
/// val/test windows the same size, stepping forward by `step_months`.
///
/// Returns an empty `Vec` if the data is insufficient for even one fold.
///
/// # Arguments
///
/// * `bars` - Chronologically sorted OHLCV bars
/// * `config` - Walk-forward configuration
pub fn generate_walk_forward_windows(
    bars: &[OHLCVBar],
    config: &WalkForwardConfig,
) -> Vec<WalkForwardWindow> {
    if bars.is_empty() {
        return Vec::new();
    }

    // Determine date range from the data
    let first_bar = match bars.first() {
        Some(b) => b,
        None => return Vec::new(),
    };
    let last_bar = match bars.last() {
        Some(b) => b,
        None => return Vec::new(),
    };

    let data_start = first_bar.timestamp.date_naive();
    let data_end = last_bar.timestamp.date_naive();

    let mut windows = Vec::new();
    let mut fold = 0_usize;

    loop {
        // Train end: initial_train_months + fold * step_months from data_start
        let total_train_months = config
            .initial_train_months
            .saturating_add(fold as u32 * config.step_months);
        let train_end = match data_start.checked_add_months(Months::new(total_train_months)) {
            Some(d) => d,
            None => break,
        };

        // Val end: train_end + val_months
        let val_end = match train_end.checked_add_months(Months::new(config.val_months)) {
            Some(d) => d,
            None => break,
        };

        // Test end: val_end + test_months
        let test_end = match val_end.checked_add_months(Months::new(config.test_months)) {
            Some(d) => d,
            None => break,
        };

        // If test_end exceeds data range, we cannot form this fold.
        // Use `>` with +1 day margin: bars with `date_naive() < test_end` captures
        // all bars ON the last calendar day of the test window. Without the margin,
        // quarterly data ending on e.g. Jun-30 fails when test_end = Jul-01.
        if test_end > data_end + chrono::Duration::days(1) {
            break;
        }

        // Partition bars into train/val/test using date boundaries
        let train_bars: Vec<OHLCVBar> = bars
            .iter()
            .filter(|b| b.timestamp.date_naive() < train_end)
            .copied()
            .collect();

        let val_bars: Vec<OHLCVBar> = bars
            .iter()
            .filter(|b| {
                let d = b.timestamp.date_naive();
                d >= train_end && d < val_end
            })
            .copied()
            .collect();

        let test_bars: Vec<OHLCVBar> = bars
            .iter()
            .filter(|b| {
                let d = b.timestamp.date_naive();
                d >= val_end && d < test_end
            })
            .copied()
            .collect();

        // Skip folds where any split has fewer than MIN_BARS_PER_SPLIT bars
        if train_bars.len() < MIN_BARS_PER_SPLIT
            || val_bars.len() < MIN_BARS_PER_SPLIT
            || test_bars.len() < MIN_BARS_PER_SPLIT
        {
            // If training set is too small, later folds might still work
            // But if val/test are too small at the end, no point continuing
            if val_bars.len() < MIN_BARS_PER_SPLIT || test_bars.len() < MIN_BARS_PER_SPLIT {
                // Check if we have already generated at least one fold or if
                // the data simply does not support this fold size
                if fold > 0 {
                    break;
                }
                // For fold 0 with insufficient data, skip but try next fold
                fold = fold.saturating_add(1);
                continue;
            }
            fold = fold.saturating_add(1);
            continue;
        }

        let difficulty_score = compute_difficulty(&train_bars);
        windows.push(WalkForwardWindow {
            fold,
            train: train_bars,
            val: val_bars,
            test: test_bars,
            train_end,
            val_end,
            test_end,
            difficulty_score,
        });

        fold = fold.saturating_add(1);
    }

    windows
}

/// Walk-forward fold as index ranges into pre-loaded arrays.
/// Zero data copying — just start/end indices into bars/features arrays.
#[derive(Debug, Clone)]
pub struct FoldRange {
    pub fold: usize,
    pub train_start: usize,
    pub train_end: usize,
    pub val_start: usize,
    pub val_end: usize,
    pub difficulty_score: f64,
}

/// Generate walk-forward folds as index ranges (zero-copy).
///
/// Same date-boundary logic as `generate_walk_forward_windows` but returns
/// `(start, end)` indices instead of cloning bars. Uses `partition_point`
/// for O(log n) index lookup.
pub fn generate_walk_forward_indices(
    bars: &[OHLCVBar],
    config: &WalkForwardConfig,
) -> Vec<FoldRange> {
    if bars.is_empty() {
        return Vec::new();
    }

    let data_start = bars[0].timestamp.date_naive();
    let data_end = bars[bars.len() - 1].timestamp.date_naive();

    let mut ranges = Vec::new();
    let mut fold = 0_usize;

    loop {
        let total_train_months = config
            .initial_train_months
            .saturating_add(fold as u32 * config.step_months);
        let train_end = match data_start.checked_add_months(Months::new(total_train_months)) {
            Some(d) => d,
            None => break,
        };
        let val_end = match train_end.checked_add_months(Months::new(config.val_months)) {
            Some(d) => d,
            None => break,
        };
        let test_end = match val_end.checked_add_months(Months::new(config.test_months)) {
            Some(d) => d,
            None => break,
        };

        if test_end > data_end + chrono::Duration::days(1) {
            break;
        }

        // O(log n) index lookups via partition_point (bars are sorted by timestamp)
        let train_end_idx = bars.partition_point(|b| b.timestamp.date_naive() < train_end);
        let val_end_idx = bars.partition_point(|b| b.timestamp.date_naive() < val_end);
        let test_end_idx = bars.partition_point(|b| b.timestamp.date_naive() < test_end);

        let train_start = 0; // expanding window: always starts at beginning
        let val_start = train_end_idx;
        let val_end_bound = val_end_idx;

        // Skip folds where any split is too small
        let train_len = train_end_idx - train_start;
        let val_len = val_end_bound - val_start;
        let test_len = test_end_idx - val_end_bound;

        if train_len < MIN_BARS_PER_SPLIT || val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT {
            if val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT {
                if fold > 0 { break; }
            }
            fold = fold.saturating_add(1);
            continue;
        }

        let difficulty_score = compute_difficulty(&bars[train_start..train_end_idx]);

        ranges.push(FoldRange {
            fold,
            train_start,
            train_end: train_end_idx,
            val_start,
            val_end: val_end_bound,
            difficulty_score,
        });

        fold = fold.saturating_add(1);
    }

    ranges
}

/// Generate walk-forward folds as index ranges from nanosecond timestamps.
///
/// Same date-boundary logic as [`generate_walk_forward_indices`] but operates
/// on raw `i64` timestamps (nanoseconds since Unix epoch) from fxcache data
/// instead of `OHLCVBar`. Difficulty score is always 0.0 because we don't
/// have OHLC data to compute ADX.
pub fn generate_walk_forward_indices_from_timestamps(
    timestamps: &[i64],
    config: &WalkForwardConfig,
) -> Vec<FoldRange> {
    if timestamps.is_empty() {
        return Vec::new();
    }

    let first_ts = timestamps[0];
    let last_ts = timestamps[timestamps.len() - 1];

    let data_start = chrono::DateTime::from_timestamp_nanos(first_ts).date_naive();
    let data_end = chrono::DateTime::from_timestamp_nanos(last_ts).date_naive();

    let mut ranges = Vec::new();
    let mut fold = 0_usize;

    loop {
        let total_train_months = config
            .initial_train_months
            .saturating_add(fold as u32 * config.step_months);
        let train_end = match data_start.checked_add_months(Months::new(total_train_months)) {
            Some(d) => d,
            None => break,
        };
        let val_end = match train_end.checked_add_months(Months::new(config.val_months)) {
            Some(d) => d,
            None => break,
        };
        let test_end = match val_end.checked_add_months(Months::new(config.test_months)) {
            Some(d) => d,
            None => break,
        };

        if test_end > data_end + chrono::Duration::days(1) {
            break;
        }

        // O(log n) index lookups via partition_point (timestamps are sorted)
        let train_end_idx = timestamps.partition_point(|&ts| {
            chrono::DateTime::from_timestamp_nanos(ts).date_naive() < train_end
        });
        let val_end_idx = timestamps.partition_point(|&ts| {
            chrono::DateTime::from_timestamp_nanos(ts).date_naive() < val_end
        });
        let test_end_idx = timestamps.partition_point(|&ts| {
            chrono::DateTime::from_timestamp_nanos(ts).date_naive() < test_end
        });

        let train_start = 0; // expanding window: always starts at beginning
        let val_start = train_end_idx;
        let val_end_bound = val_end_idx;

        // Skip folds where any split is too small
        let train_len = train_end_idx - train_start;
        let val_len = val_end_bound - val_start;
        let test_len = test_end_idx - val_end_bound;

        if train_len < MIN_BARS_PER_SPLIT || val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT {
            if val_len < MIN_BARS_PER_SPLIT || test_len < MIN_BARS_PER_SPLIT {
                if fold > 0 { break; }
            }
            fold = fold.saturating_add(1);
            continue;
        }

        // No ADX computation — fxcache doesn't have OHLC bars for difficulty
        ranges.push(FoldRange {
            fold,
            train_start,
            train_end: train_end_idx,
            val_start,
            val_end: val_end_bound,
            difficulty_score: 0.0,
        });

        fold = fold.saturating_add(1);
    }

    ranges
}

/// Per-feature normalization statistics (mean and standard deviation).
///
/// Computed from training data and applied to val/test data to prevent
/// information leakage. Uses z-score normalization: `(x - mean) / std`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct NormStats {
    /// Per-feature mean values (length 51).
    pub mean: Vec<f64>,
    /// Per-feature standard deviation values (length 51, clamped >= 1e-8).
    pub std: Vec<f64>,
}

/// Number of features in a standard feature vector (42 market features: 40 base + 2 regime).
const FEATURE_DIM: usize = 42;

/// Minimum standard deviation to prevent division by zero.
const MIN_STD: f64 = 1e-8;

/// Maximum legitimate magnitude of a z-score-normalized feature.
///
/// Z-scores in real-market 42-dim feature vectors live well within ±5 even on
/// extreme bars (10σ events are vanishingly rare); ±20 is a generous ceiling
/// that traps every observed corruption pattern (raw-price leak, corrupt bar
/// with log_return = -30 / std=1e-3 = -30000, etc.) without rejecting any
/// legitimate market signal. Used as the shared invariant in
/// [`validate_normalized_features`].
pub const NORMALIZED_FEATURE_BOUND: f64 = 20.0;

impl NormStats {
    /// Compute normalization statistics from a batch of feature vectors (FEATURE_DIM dimensions).
    ///
    /// If the input is empty, returns mean = 0.0 and std = 1.0 for all features.
    pub fn from_features(features: &[[f64; FEATURE_DIM]]) -> Self {
        if features.is_empty() {
            return Self {
                mean: vec![0.0; FEATURE_DIM],
                std: vec![1.0; FEATURE_DIM],
            };
        }

        let n = features.len() as f64;
        let mut mean = vec![0.0_f64; FEATURE_DIM];
        let mut variance = vec![0.0_f64; FEATURE_DIM];

        // Accumulate sums for mean
        for feature_vec in features {
            for (m, val) in mean.iter_mut().zip(feature_vec.iter()) {
                *m += val;
            }
        }

        // Compute mean
        for m in &mut mean {
            *m /= n;
        }

        // Accumulate variance
        for feature_vec in features {
            for (v, (val, m)) in variance
                .iter_mut()
                .zip(feature_vec.iter().zip(mean.iter()))
            {
                let diff = val - m;
                *v += diff * diff;
            }
        }

        // Compute std with minimum clamp
        let std_vec: Vec<f64> = variance
            .iter()
            .map(|v| (v / n).sqrt().max(MIN_STD))
            .collect();

        Self {
            mean,
            std: std_vec,
        }
    }

    /// Z-score normalize a single 42-dimensional feature vector.
    ///
    /// Returns `clamp((feature - mean) / std, ±NORMALIZED_FEATURE_BOUND)`
    /// element-wise. The clamp is the structural guarantee that downstream GPU
    /// consumers (state_gather → state[0] in `experience_state_gather`, aux
    /// label_scale_ema, vol_normalizer) see bounded inputs regardless of what
    /// corruption survives upstream. Without it, a single outlier from a
    /// corrupt DBN bar inflates `std` enough to amplify itself by sqrt(N) post-
    /// normalization and flow through training as a ±30000 magnitude state[0]
    /// outlier.
    pub fn normalize(&self, features: &[f64; FEATURE_DIM]) -> [f64; FEATURE_DIM] {
        let mut result = [0.0_f64; FEATURE_DIM];
        for ((r, val), (m, s)) in result
            .iter_mut()
            .zip(features.iter())
            .zip(self.mean.iter().zip(self.std.iter()))
        {
            *r = ((val - m) / s).clamp(-NORMALIZED_FEATURE_BOUND, NORMALIZED_FEATURE_BOUND);
        }
        result
    }

    /// Z-score normalize a batch of 42-dimensional feature vectors.
    pub fn normalize_batch(&self, features: &[[f64; FEATURE_DIM]]) -> Vec<[f64; FEATURE_DIM]> {
        features.iter().map(|f| self.normalize(f)).collect()
    }
}

/// Validate that a batch of normalized features is bounded and free of NaN/Inf.
///
/// Single source of truth for the post-normalization invariant — used by both
/// the fxcache fast path (after `discover_and_load`) and the DBN fallback
/// (after `normalize_batch`). Either path producing data that fails this gate
/// must error out rather than upload poisoned values to the GPU.
///
/// Returns `Ok(())` if every element of every vector is finite and within
/// `±NORMALIZED_FEATURE_BOUND`. On the first violation, returns an `Err` with
/// the bar index, feature index, and value so the operator can locate the
/// corrupt source row.
pub fn validate_normalized_features(
    features: &[[f64; FEATURE_DIM]],
    source: &str,
) -> anyhow::Result<()> {
    for (bar_idx, feat_vec) in features.iter().enumerate() {
        for (feat_idx, &val) in feat_vec.iter().enumerate() {
            if !val.is_finite() {
                anyhow::bail!(
                    "{source}: non-finite value at bar={bar_idx}, feat={feat_idx}: {val} \
                     — data corruption between source and post-normalization."
                );
            }
            if val.abs() > NORMALIZED_FEATURE_BOUND {
                anyhow::bail!(
                    "{source}: |feature[{bar_idx}][{feat_idx}]| = {} exceeds normalized bound \
                     (±{NORMALIZED_FEATURE_BOUND}) — corrupt bar (e.g. bar.open ≈ 0 producing \
                     extreme log return) or unit-mismatch (raw price written instead of \
                     z-normalized log return). Reject the dataset; do not upload to GPU.",
                    val
                );
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{Datelike, NaiveDate, NaiveTime, TimeZone, Utc, Weekday};

    /// Create a single OHLCV bar at the given date.
    fn make_bar(year: i32, month: u32, day: u32) -> OHLCVBar {
        let date = NaiveDate::from_ymd_opt(year, month, day).unwrap_or_default();
        let time = NaiveTime::from_hms_opt(10, 0, 0).unwrap_or_default();
        let dt = date.and_time(time);
        let timestamp = Utc.from_utc_datetime(&dt);
        OHLCVBar {
            timestamp,
            open: 100.0,
            high: 101.0,
            low: 99.0,
            close: 100.5,
            volume: 1000.0,
        }
    }

    /// Create bars for every weekday in [start, end) with one bar per day.
    /// For testing purposes, this creates enough density to exceed the minimum
    /// bars threshold when covering multi-month ranges.
    fn make_bars_range(start: NaiveDate, end: NaiveDate) -> Vec<OHLCVBar> {
        let mut bars = Vec::new();
        let mut current = start;
        let time = NaiveTime::from_hms_opt(10, 0, 0).unwrap_or_default();
        while current < end {
            let weekday = current.weekday();
            if weekday != Weekday::Sat && weekday != Weekday::Sun {
                // Generate multiple bars per trading day to simulate 1-min bars
                // We need enough density: ~390 bars per day for 6.5h of trading
                // For test efficiency, generate 20 bars per day (enough for density)
                for minute_offset in 0..20 {
                    let bar_time =
                        NaiveTime::from_hms_opt(9, 30_u32.saturating_add(minute_offset), 0)
                            .unwrap_or(time);
                    let dt = current.and_time(bar_time);
                    let timestamp = Utc.from_utc_datetime(&dt);
                    bars.push(OHLCVBar {
                        timestamp,
                        open: 100.0 + (bars.len() as f64 * 0.001),
                        high: 101.0 + (bars.len() as f64 * 0.001),
                        low: 99.0 + (bars.len() as f64 * 0.001),
                        close: 100.5 + (bars.len() as f64 * 0.001),
                        volume: 1000.0,
                    });
                }
            }
            current = current
                .succ_opt()
                .unwrap_or(current);
        }
        bars
    }

    #[test]
    fn test_generate_walk_forward_empty_bars() {
        let config = WalkForwardConfig::default();
        let windows = generate_walk_forward_windows(&[], &config);
        assert!(windows.is_empty());
    }

    #[test]
    fn test_walk_forward_windows_24_months() {
        // Generate 24 months of synthetic bars
        let start = NaiveDate::from_ymd_opt(2023, 1, 1).unwrap_or_default();
        let end = NaiveDate::from_ymd_opt(2025, 1, 1).unwrap_or_default();
        let bars = make_bars_range(start, end);

        assert!(
            !bars.is_empty(),
            "Should have generated bars for 24-month range"
        );

        let config = WalkForwardConfig::default(); // 12/3/3/3
        let windows = generate_walk_forward_windows(&bars, &config);

        // With 24 months of data and 12+3+3=18 months for first fold, stepping by 3:
        // Fold 0: train 12mo, val 3mo, test 3mo = 18 months (fits in 24)
        // Fold 1: train 15mo, val 3mo, test 3mo = 21 months (fits in 24)
        // Fold 2: train 18mo, val 3mo, test 3mo = 24 months (might barely fit)
        assert!(
            windows.len() >= 2,
            "Expected at least 2 windows, got {}",
            windows.len()
        );

        // Check fold 0 has non-empty splits
        if let Some(w0) = windows.first() {
            assert!(!w0.train.is_empty(), "Fold 0 train should be non-empty");
            assert!(!w0.val.is_empty(), "Fold 0 val should be non-empty");
            assert!(!w0.test.is_empty(), "Fold 0 test should be non-empty");

            // Train must end before val
            assert!(
                w0.train_end <= w0.val_end,
                "Train end ({}) should be <= val end ({})",
                w0.train_end,
                w0.val_end
            );
        }
    }

    #[test]
    fn test_walk_forward_expanding_windows() {
        // Generate 30 months to get multiple folds
        let start = NaiveDate::from_ymd_opt(2023, 1, 1).unwrap_or_default();
        let end = NaiveDate::from_ymd_opt(2025, 7, 1).unwrap_or_default();
        let bars = make_bars_range(start, end);

        let config = WalkForwardConfig::default();
        let windows = generate_walk_forward_windows(&bars, &config);

        assert!(
            windows.len() >= 2,
            "Need at least 2 windows to test expanding, got {}",
            windows.len()
        );

        // Each successive fold should have strictly more training data
        for pair in windows.windows(2) {
            let (prev, next) = match (pair.first(), pair.get(1)) {
                (Some(p), Some(n)) => (p, n),
                _ => continue,
            };
            assert!(
                next.train.len() > prev.train.len(),
                "Fold {} train ({} bars) should have more data than fold {} train ({} bars)",
                next.fold,
                next.train.len(),
                prev.fold,
                prev.train.len()
            );
        }
    }

    #[test]
    fn test_norm_stats_roundtrip() {
        // Features: [[1,1,...], [2,2,...], [3,3,...]]
        let f1 = [1.0_f64; FEATURE_DIM];
        let f2 = [2.0_f64; FEATURE_DIM];
        let f3 = [3.0_f64; FEATURE_DIM];
        let features = [f1, f2, f3];

        let stats = NormStats::from_features(&features);

        // Mean should be 2.0 for all features
        for m in &stats.mean {
            assert!(
                (*m - 2.0).abs() < 1e-10,
                "Expected mean ~2.0, got {}",
                m
            );
        }

        // Normalizing the mean vector should give ~0 for all features
        let normalized = stats.normalize(&f2);
        for val in &normalized {
            assert!(
                val.abs() < 1e-10,
                "Expected normalized mean ~0.0, got {}",
                val
            );
        }
    }

    #[test]
    fn test_norm_stats_empty() {
        let features: &[[f64; FEATURE_DIM]] = &[];
        let stats = NormStats::from_features(features);

        // Should return safe defaults
        assert_eq!(stats.mean.len(), FEATURE_DIM);
        assert_eq!(stats.std.len(), FEATURE_DIM);

        // Mean should be 0.0
        for m in &stats.mean {
            assert!((*m).abs() < 1e-10, "Expected mean 0.0, got {}", m);
        }

        // Std should be 1.0
        for s in &stats.std {
            assert!(
                (*s - 1.0).abs() < 1e-10,
                "Expected std 1.0, got {}",
                s
            );
        }
    }
}