foxhunt/ml/tests/transition_matrix_test.rs

//! Regime Transition Matrix Tests
//!
//! TDD tests for regime transition probability matrix implementation.
//! Tests cover:
//! - Transition probability updates
//! - Stationary distribution convergence
//! - Real data regime sequence analysis

use ml::ensemble::MarketRegime;
use ml::regime::transition_matrix::RegimeTransitionMatrix;

#[test]
fn test_transition_matrix_initialization() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
        MarketRegime::Sideways,
        MarketRegime::HighVolatility,
    ];

    let matrix = RegimeTransitionMatrix::new(regimes.clone(), 0.1, 10);

    // Verify all regimes are tracked
    assert_eq!(matrix.regime_count(), 4);

    // Initial transition probabilities should be uniform (1/N for each regime)
    for from in &regimes {
        for to in &regimes {
            let prob = matrix.get_transition_prob(*from, *to);
            assert!((prob - 0.25).abs() < 1e-6,
                "Initial probability should be ~0.25 (uniform), got {}", prob);
        }
    }
}

#[test]
fn test_single_transition_update() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.5, 1);

    // Update: Bull -> Bear
    matrix.update(MarketRegime::Bull, MarketRegime::Bear);

    // After 1 observation with alpha=0.5:
    // P(Bull->Bear) should increase from 0.5 to ~0.75
    // P(Bull->Bull) should decrease from 0.5 to ~0.25
    let p_bull_to_bear = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
    let p_bull_to_bull = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bull);

    assert!(p_bull_to_bear > 0.6, "P(Bull->Bear) should increase, got {}", p_bull_to_bear);
    assert!(p_bull_to_bull < 0.4, "P(Bull->Bull) should decrease, got {}", p_bull_to_bull);

    // Row should sum to 1.0
    let row_sum = p_bull_to_bear + p_bull_to_bull;
    assert!((row_sum - 1.0).abs() < 1e-6, "Row sum should be 1.0, got {}", row_sum);
}

#[test]
fn test_multiple_transitions_same_path() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.2, 1);

    // Repeat Bull -> Bear 10 times
    for _ in 0..10 {
        matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    }

    // P(Bull->Bear) should approach 1.0
    let p_bull_to_bear = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
    assert!(p_bull_to_bear > 0.8, "After 10 observations, P(Bull->Bear) should be >0.8, got {}", p_bull_to_bear);
}

#[test]
fn test_self_transitions() {
    let regimes = vec![
        MarketRegime::Sideways,
        MarketRegime::HighVolatility,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.3, 1);

    // Update: Sideways -> Sideways (persistence)
    for _ in 0..5 {
        matrix.update(MarketRegime::Sideways, MarketRegime::Sideways);
    }

    // P(Sideways->Sideways) should be high (regime persistence)
    let p_sideways_persist = matrix.get_transition_prob(
        MarketRegime::Sideways,
        MarketRegime::Sideways
    );
    assert!(p_sideways_persist > 0.7,
        "Sideways should persist, P(Sideways->Sideways) = {}", p_sideways_persist);
}

#[test]
fn test_row_normalization() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
        MarketRegime::Sideways,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes.clone(), 0.25, 1);

    // Add various transitions
    matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    matrix.update(MarketRegime::Bull, MarketRegime::Sideways);
    matrix.update(MarketRegime::Bear, MarketRegime::Bull);

    // Check that all rows sum to 1.0
    for from in &regimes {
        let row_sum: f64 = regimes.iter()
            .map(|to| matrix.get_transition_prob(*from, *to))
            .sum();

        assert!((row_sum - 1.0).abs() < 1e-6,
            "Row {:?} sum should be 1.0, got {}", from, row_sum);
    }
}

#[test]
fn test_minimum_observations_threshold() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.2, 5); // min_obs = 5

    // Add only 2 observations (below threshold)
    matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    matrix.update(MarketRegime::Bull, MarketRegime::Bear);

    // Should still use uniform priors until min_observations reached
    let p_bull_to_bear = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);

    // With insufficient data, probability should be close to prior (0.5)
    // The exact behavior depends on implementation (Laplace smoothing)
    assert!(p_bull_to_bear >= 0.4 && p_bull_to_bear <= 0.8,
        "With insufficient observations, probability should use smoothing, got {}", p_bull_to_bear);
}

#[test]
fn test_stationary_distribution_uniform() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.2, 1);

    // Create perfectly symmetric transitions: P(Bull->Bear) = P(Bear->Bull) = 0.5
    // This should yield stationary distribution [0.5, 0.5]
    for _ in 0..10 {
        matrix.update(MarketRegime::Bull, MarketRegime::Bear);
        matrix.update(MarketRegime::Bear, MarketRegime::Bull);
    }

    let stationary = matrix.get_stationary_distribution();

    let bull_prob = stationary.get(&MarketRegime::Bull).unwrap_or(&0.0);
    let bear_prob = stationary.get(&MarketRegime::Bear).unwrap_or(&0.0);

    // Should be approximately equal
    assert!((bull_prob - 0.5).abs() < 0.15,
        "Bull stationary probability should be ~0.5, got {}", bull_prob);
    assert!((bear_prob - 0.5).abs() < 0.15,
        "Bear stationary probability should be ~0.5, got {}", bear_prob);

    // Should sum to 1.0
    let total: f64 = stationary.values().sum();
    assert!((total - 1.0).abs() < 1e-6, "Stationary distribution should sum to 1.0, got {}", total);
}

#[test]
fn test_stationary_distribution_absorbing() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.3, 1);

    // Create Bull as absorbing state: P(Bull->Bull) = 1.0
    for _ in 0..20 {
        matrix.update(MarketRegime::Bull, MarketRegime::Bull);
        matrix.update(MarketRegime::Bear, MarketRegime::Bull);
    }

    let stationary = matrix.get_stationary_distribution();

    let bull_prob = stationary.get(&MarketRegime::Bull).unwrap_or(&0.0);

    // Bull should dominate stationary distribution
    assert!(*bull_prob > 0.7,
        "Bull should dominate as absorbing state, got {}", bull_prob);
}

#[test]
fn test_expected_duration_high_persistence() {
    let regimes = vec![
        MarketRegime::Sideways,
        MarketRegime::HighVolatility,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.2, 1);

    // Make Sideways highly persistent: P(Sideways->Sideways) = 0.9
    for _ in 0..20 {
        matrix.update(MarketRegime::Sideways, MarketRegime::Sideways);
        matrix.update(MarketRegime::Sideways, MarketRegime::Sideways);
        matrix.update(MarketRegime::Sideways, MarketRegime::HighVolatility);
    }

    // Expected duration = 1 / (1 - P(i->i))
    // If P(Sideways->Sideways) = 0.9, duration = 1 / 0.1 = 10
    let duration = matrix.get_expected_duration(MarketRegime::Sideways);

    assert!(duration > 3.0,
        "High persistence should yield long duration, got {}", duration);
    assert!(duration < 50.0,
        "Duration should be finite, got {}", duration);
}

#[test]
fn test_expected_duration_low_persistence() {
    let regimes = vec![
        MarketRegime::HighVolatility,
        MarketRegime::Sideways,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes, 0.3, 1);

    // Make HighVolatility transient: P(HV->HV) = 0.2
    for _ in 0..20 {
        matrix.update(MarketRegime::HighVolatility, MarketRegime::Sideways);
        matrix.update(MarketRegime::HighVolatility, MarketRegime::Sideways);
        matrix.update(MarketRegime::HighVolatility, MarketRegime::Sideways);
        matrix.update(MarketRegime::HighVolatility, MarketRegime::HighVolatility);
    }

    // Low persistence -> short duration
    let duration = matrix.get_expected_duration(MarketRegime::HighVolatility);

    assert!(duration >= 1.0 && duration < 3.0,
        "Low persistence should yield short duration, got {}", duration);
}

#[test]
fn test_four_regime_matrix() {
    let regimes = vec![
        MarketRegime::Bull,
        MarketRegime::Bear,
        MarketRegime::Sideways,
        MarketRegime::HighVolatility,
    ];

    let mut matrix = RegimeTransitionMatrix::new(regimes.clone(), 0.15, 1);

    // Simulate realistic regime transitions
    let transitions = vec![
        (MarketRegime::Sideways, MarketRegime::Bull),       // Breakout to bull
        (MarketRegime::Bull, MarketRegime::Bull),           // Bull persistence
        (MarketRegime::Bull, MarketRegime::HighVolatility), // Volatility spike
        (MarketRegime::HighVolatility, MarketRegime::Bear), // Crash
        (MarketRegime::Bear, MarketRegime::Bear),           // Bear persistence
        (MarketRegime::Bear, MarketRegime::Sideways),       // Stabilization
    ];

    for (from, to) in transitions {
        matrix.update(from, to);
    }

    // Verify all rows still sum to 1.0
    for from in &regimes {
        let row_sum: f64 = regimes.iter()
            .map(|to| matrix.get_transition_prob(*from, *to))
            .sum();

        assert!((row_sum - 1.0).abs() < 1e-6,
            "Row {:?} sum should be 1.0, got {}", from, row_sum);
    }

    // Verify stationary distribution sums to 1.0
    let stationary = matrix.get_stationary_distribution();
    let total: f64 = stationary.values().sum();
    assert!((total - 1.0).abs() < 1e-6,
        "Stationary distribution should sum to 1.0, got {}", total);
}