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
    );
}