foxhunt/crates/ml/src/validation/regime_analysis.rs

//! Per-regime performance analysis.
//!
//! Groups daily returns by market regime and computes performance metrics
//! (Sharpe ratio, win rate, average return) for each regime bucket.
//!
//! Provides both a CPU fallback ([`per_regime_breakdown`]) and a GPU-accelerated
//! path ([`per_regime_breakdown_gpu`]) that classifies regimes via tensor ops on
//! device, reading back only 3 mask vectors for the final grouping.

use std::collections::HashMap;

use ml_core::cuda_autograd::GpuTensor;
use ml_core::device::MlDevice;
use ml_core::native_types::NativeDevice;
use serde::{Deserialize, Serialize};

use crate::dqn::{RegimeType, RegimeClassConfig};
use ml_validation::statistical::sharpe_ratio;

/// Per-regime performance metrics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegimeMetrics {
    /// The market regime this bucket corresponds to.
    pub regime: RegimeType,
    /// Sharpe ratio of returns in this regime.
    pub sharpe: f64,
    /// Number of bars observed in this regime.
    pub num_bars: usize,
    /// Fraction of bars with positive returns (0.0-1.0).
    pub win_rate: f64,
    /// Arithmetic mean of returns in this regime.
    pub avg_return: f64,
}

/// Break down daily returns by market regime and compute per-regime metrics.
///
/// Iterates over `min(daily_returns.len(), features.len())` bars. Each bar is
/// classified into a regime via [`RegimeType::classify_from_features`], and
/// returns are grouped accordingly. For each regime bucket the function
/// computes the Sharpe ratio, win rate, average return, and bar count.
///
/// Returns an empty [`HashMap`] when both inputs are empty.
pub fn per_regime_breakdown(
    daily_returns: &[f64],
    features: &[Vec<f32>],
) -> HashMap<RegimeType, RegimeMetrics> {
    let n = daily_returns.len().min(features.len());

    // Group returns by regime using safe indexing via iterators.
    let mut regime_returns: HashMap<RegimeType, Vec<f64>> = HashMap::new();
    let regime_cfg = RegimeClassConfig::default();

    for (ret, feat) in daily_returns.iter().zip(features.iter()).take(n) {
        let regime = RegimeType::classify_from_features(feat, &regime_cfg);
        regime_returns.entry(regime).or_default().push(*ret);
    }

    // Compute metrics for each regime bucket.
    compute_regime_metrics(regime_returns)
}

/// Minimum number of features required for GPU regime classification.
///
/// The GPU path needs ADX at index 40 and CUSUM direction at index 41 (via
/// [`RegimeClassConfig`] defaults), so each feature vector must have at least
/// 42 elements.
const MIN_FEATURES_FOR_GPU: usize = 42;

/// GPU-accelerated per-regime breakdown using tensor-based classification.
///
/// Builds a `[n, feature_dim]` tensor on `device`, calls
/// [`RegimeType::classify_regime_masks_gpu`] to produce binary masks entirely
/// on device, then reads back the 3 mask vectors (single GPU-to-CPU transfer)
/// and groups returns on the CPU side.
///
/// Falls back to the CPU path ([`per_regime_breakdown`]) when:
/// - Any feature vector has fewer than [`MIN_FEATURES_FOR_GPU`] elements
/// - Tensor construction or GPU classification fails
///
/// # Errors
///
/// Returns [`MLError`] only on truly unrecoverable failures. Soft errors
/// (short features, GPU OOM) are caught and redirected to the CPU fallback.
pub fn per_regime_breakdown_gpu(
    daily_returns: &[f64],
    features: &[Vec<f32>],
    device: &NativeDevice,
) -> Result<HashMap<RegimeType, RegimeMetrics>, crate::MLError> {
    let n = daily_returns.len().min(features.len());
    if n == 0 {
        return Ok(HashMap::new());
    }

    // All feature vectors must be long enough for the GPU regime classifier.
    let feature_dim = features.iter().take(n).map(Vec::len).min().unwrap_or(0);
    if feature_dim < MIN_FEATURES_FOR_GPU {
        return Ok(per_regime_breakdown(daily_returns, features));
    }

    // Get CUDA ordinal from device; fall back to CPU if not CUDA.
    let ordinal = match device.cuda_ordinal() {
        Some(id) => id,
        None => return Ok(per_regime_breakdown(daily_returns, features)),
    };

    // Create MlDevice to obtain a CUDA stream.
    let ml_device = match MlDevice::cuda(ordinal) {
        Ok(d) => d,
        Err(_) => return Ok(per_regime_breakdown(daily_returns, features)),
    };
    let stream = match ml_device.cuda_stream() {
        Ok(s) => s,
        Err(_) => return Ok(per_regime_breakdown(daily_returns, features)),
    };

    // Build a flat f32 buffer, truncating each row to `feature_dim`.
    let mut flat: Vec<f32> = Vec::with_capacity(n * feature_dim);
    for feat in features.iter().take(n) {
        flat.extend_from_slice(feat.get(..feature_dim).unwrap_or(feat));
    }

    // Construct the [n, feature_dim] GpuTensor on device.
    let states = match GpuTensor::from_host(&flat, vec![n, feature_dim], stream) {
        Ok(t) => t,
        Err(_) => return Ok(per_regime_breakdown(daily_returns, features)),
    };

    // Run GPU regime classification.
    let regime_cfg = RegimeClassConfig::default();
    let (trending_mask, ranging_mask, volatile_mask) =
        match RegimeType::classify_regime_masks_gpu(&states, &regime_cfg, stream) {
            Ok(masks) => masks,
            Err(_) => return Ok(per_regime_breakdown(daily_returns, features)),
        };

    // Read masks back to CPU — single transfer of 3 x n f32 values.
    let trending: Vec<f32> = trending_mask.to_host(stream).map_err(|e| {
        crate::MLError::ModelError(format!("trending mask readback: {e}"))
    })?;
    let ranging: Vec<f32> = ranging_mask.to_host(stream).map_err(|e| {
        crate::MLError::ModelError(format!("ranging mask readback: {e}"))
    })?;
    let volatile: Vec<f32> = volatile_mask.to_host(stream).map_err(|e| {
        crate::MLError::ModelError(format!("volatile mask readback: {e}"))
    })?;

    // Group returns by regime using the binary masks.
    let mut regime_returns: HashMap<RegimeType, Vec<f64>> = HashMap::new();
    for i in 0..n {
        let ret = daily_returns.get(i).copied().unwrap_or(0.0);
        let regime = if trending.get(i).copied().unwrap_or(0.0) > 0.5 {
            RegimeType::Trending
        } else if volatile.get(i).copied().unwrap_or(0.0) > 0.5 {
            RegimeType::Volatile
        } else if ranging.get(i).copied().unwrap_or(0.0) > 0.5 {
            RegimeType::Ranging
        } else {
            // Fallback — should not happen with well-formed masks.
            RegimeType::Ranging
        };
        regime_returns.entry(regime).or_default().push(ret);
    }

    // Compute per-regime metrics (same logic as the CPU path).
    Ok(compute_regime_metrics(regime_returns))
}

/// Compute [`RegimeMetrics`] for each regime bucket from grouped returns.
///
/// Shared helper used by both the CPU and GPU paths.
fn compute_regime_metrics(
    regime_returns: HashMap<RegimeType, Vec<f64>>,
) -> HashMap<RegimeType, RegimeMetrics> {
    regime_returns
        .into_iter()
        .map(|(regime, returns)| {
            let num_bars = returns.len();
            let sharpe_val = sharpe_ratio(&returns);

            let avg_return = if num_bars > 0 {
                returns.iter().sum::<f64>() / num_bars as f64
            } else {
                0.0
            };

            let win_rate = if num_bars > 0 {
                let wins = returns.iter().filter(|r| **r > 0.0).count();
                wins as f64 / num_bars as f64
            } else {
                0.0
            };

            let metrics = RegimeMetrics {
                regime,
                sharpe: sharpe_val,
                num_bars,
                win_rate,
                avg_return,
            };

            (regime, metrics)
        })
        .collect()
}

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

    #[test]
    fn test_per_regime_breakdown_groups_correctly() {
        // With short feature vectors (< 211 elements), all bars classify as Ranging.
        let daily_returns = vec![0.01, -0.005, 0.02, -0.01, 0.015];
        let features: Vec<Vec<f32>> = (0..5).map(|_| vec![0.0_f32; 50]).collect();

        let result = per_regime_breakdown(&daily_returns, &features);

        // All should fall into the Ranging bucket.
        assert_eq!(result.len(), 1, "Expected exactly one regime bucket");
        assert!(
            result.contains_key(&RegimeType::Ranging),
            "Expected Ranging regime bucket"
        );

        let ranging = result.get(&RegimeType::Ranging);
        assert!(ranging.is_some(), "Ranging bucket should exist");
        if let Some(m) = ranging {
            assert_eq!(m.num_bars, 5, "Expected 5 bars in Ranging bucket");
            assert_eq!(m.regime, RegimeType::Ranging);
        }
    }

    #[test]
    fn test_per_regime_metrics_computation() {
        // All returns land in Ranging (short features), so we test metrics directly.
        // Returns: [0.10, -0.05, 0.20, -0.10, 0.15]
        //   positive: 3 out of 5 -> win_rate = 0.6
        //   sum = 0.30, avg_return = 0.06
        let daily_returns = vec![0.10, -0.05, 0.20, -0.10, 0.15];
        let features: Vec<Vec<f32>> = (0..5).map(|_| vec![0.0_f32; 10]).collect();

        let result = per_regime_breakdown(&daily_returns, &features);
        let metrics = result
            .get(&RegimeType::Ranging)
            .expect("Ranging bucket must exist");

        // win_rate: 3 positive out of 5
        assert!(
            (metrics.win_rate - 0.6).abs() < 1e-10,
            "Expected win_rate=0.6, got {}",
            metrics.win_rate
        );

        // avg_return: (0.10 + (-0.05) + 0.20 + (-0.10) + 0.15) / 5 = 0.06
        assert!(
            (metrics.avg_return - 0.06).abs() < 1e-10,
            "Expected avg_return=0.06, got {}",
            metrics.avg_return
        );

        // num_bars
        assert_eq!(metrics.num_bars, 5);

        // sharpe should be positive (mean > 0 with reasonable variance)
        assert!(
            metrics.sharpe > 0.0,
            "Expected positive Sharpe, got {}",
            metrics.sharpe
        );
    }

    #[test]
    fn test_empty_returns() {
        let result = per_regime_breakdown(&[], &[]);
        assert!(
            result.is_empty(),
            "Expected empty HashMap for empty inputs, got {} entries",
            result.len()
        );
    }
}