foxhunt/crates/ml-regime/src/transition_probability_features.rs

//! Transition Probability Features (Indices 216-220)
//!
//! Extracts 5 features from regime transition probabilities:
//! - Feature 216: Stability P(i→i) - probability of staying in current regime
//! - Feature 217: Most likely next regime (index) - which regime is most probable next
//! - Feature 218: Shannon entropy H = -Σ P(i→j) log₂ P(i→j) - uncertainty measure
//! - Feature 219: Expected duration - how long regime typically persists
//! - Feature 220: Change probability (1 - stability) - probability of regime change
//!
//! **ARCHITECTURAL DESIGN**:
//! - **REUSE** existing `RegimeTransitionMatrix` for all probability calculations
//! - **REUSE** existing `expected_duration()` method for Feature 219
//! - No duplication of transition tracking logic
//! - Shannon entropy computed with numerical stability (filters p < 1e-10)
//!
//! # Example
//!
//! ```rust
//! use ml::regime::transition_probability_features::TransitionProbabilityFeatures;
//! use ml::ensemble::MarketRegime;
//!
//! let regimes = vec![
//!     MarketRegime::Bull,
//!     MarketRegime::Bear,
//!     MarketRegime::Sideways,
//! ];
//!
//! let mut features = TransitionProbabilityFeatures::new(regimes, 0.1, 10);
//!
//! // Update with observed regime
//! features.update(MarketRegime::Bull);
//! features.update(MarketRegime::Bear);
//!
//! // Compute all 5 features
//! let result = features.compute_features();
//! assert_eq!(result.len(), 5);
//! ```

use crate::MarketRegime;
use crate::transition_matrix::RegimeTransitionMatrix;

/// Transition Probability Feature Extractor
///
/// Maintains a transition matrix and extracts 5 probability-based features:
/// 1. Stability (self-transition probability)
/// 2. Most likely next regime
/// 3. Shannon entropy (uncertainty)
/// 4. Expected duration (persistence)
/// 5. Change probability (1 - stability)
///
/// # Design Principles
///
/// - **REUSE**: Delegates all transition tracking to `RegimeTransitionMatrix`
/// - **PERFORMANCE**: O(N) where N = number of regimes (typically 4-6)
/// - **NUMERICAL STABILITY**: Filters probabilities < 1e-10 before log operations
///
/// # Feature Descriptions
///
/// **Feature 216: Stability P(i→i)**
/// - Probability of staying in current regime
/// - High stability (>0.8) indicates persistent regime
/// - Low stability (<0.3) indicates transitional regime
///
/// **Feature 217: Most Likely Next Regime**
/// - Index of regime with highest transition probability from current regime
/// - Used for predictive regime classification
/// - Value range: [0, N-1] where N = number of regimes
///
/// **Feature 218: Shannon Entropy**
/// - H = -Σ P(i→j) log₂ P(i→j)
/// - Measures uncertainty in regime transitions
/// - High entropy: many possible transitions (uncertain)
/// - Low entropy: few likely transitions (predictable)
/// - Max entropy: log₂(N) for uniform distribution
///
/// **Feature 219: Expected Duration**
/// - `E[T]` = 1 / (1 - `P[i][i]`)
/// - Expected number of periods in current regime
/// - REUSES existing `get_expected_duration()` method
///
/// **Feature 220: Change Probability**
/// - 1 - P(i→i)
/// - Probability of transitioning out of current regime
/// - Complementary to stability (Feature 216)
#[derive(Debug, Clone)]
pub struct TransitionProbabilityFeatures {
    /// Regime transition matrix (REUSED infrastructure)
    matrix: RegimeTransitionMatrix,

    /// Current market regime
    current_regime: MarketRegime,

    /// List of all regimes (for iteration)
    regimes: Vec<MarketRegime>,
}

impl TransitionProbabilityFeatures {
    /// Create a new transition probability feature extractor
    ///
    /// # Arguments
    ///
    /// * `regimes` - List of market regimes to track
    /// * `alpha` - EMA smoothing factor (0 < alpha <= 1)
    /// * `min_obs` - Minimum observations before using empirical probabilities
    ///
    /// # Returns
    ///
    /// New feature extractor initialized with the first regime as current
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_probability_features::TransitionProbabilityFeatures;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![
    ///     MarketRegime::Bull,
    ///     MarketRegime::Bear,
    /// ];
    /// let features = TransitionProbabilityFeatures::new(regimes, 0.1, 10);
    /// ```
    pub fn new(regimes: Vec<MarketRegime>, alpha: f64, min_obs: usize) -> Self {
        let current_regime = regimes.last().copied().unwrap_or(MarketRegime::Unknown);
        let matrix = RegimeTransitionMatrix::new(regimes.clone(), alpha, min_obs);

        Self {
            matrix,
            current_regime,
            regimes,
        }
    }

    /// Update with new regime observation
    ///
    /// If the regime has changed, updates the transition matrix.
    /// If the regime is the same, still updates the matrix to track persistence.
    ///
    /// # Arguments
    ///
    /// * `regime` - Newly observed market regime
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_probability_features::TransitionProbabilityFeatures;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
    /// let mut features = TransitionProbabilityFeatures::new(regimes, 0.1, 1);
    ///
    /// features.update(MarketRegime::Bull);
    /// features.update(MarketRegime::Bear); // Transition recorded
    /// ```
    pub fn update(&mut self, regime: MarketRegime) {
        // Always update the matrix (even for same regime to track persistence)
        self.matrix.update(self.current_regime, regime);
        self.current_regime = regime;
    }

    /// Compute all 5 transition probability features
    ///
    /// Returns array of 5 features:
    /// - `[0]`: Stability P(i→i)
    /// - `[1]`: Most likely next regime (index)
    /// - `[2]`: Shannon entropy
    /// - `[3]`: Expected duration
    /// - `[4]`: Change probability
    ///
    /// # Returns
    ///
    /// Array of 5 f64 values representing the features
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_probability_features::TransitionProbabilityFeatures;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
    /// let mut features = TransitionProbabilityFeatures::new(regimes, 0.1, 1);
    ///
    /// features.update(MarketRegime::Bull);
    /// let result = features.compute_features();
    ///
    /// assert_eq!(result.len(), 5);
    /// ```
    pub fn compute_features(&self) -> [f64; 5] {
        // Feature 216: Stability P(i→i)
        let stability = self
            .matrix
            .get_transition_prob(self.current_regime, self.current_regime);

        // Feature 217: Most likely next regime
        let mut max_prob = 0.0;
        let mut most_likely_idx = 0;
        for (idx, &next_regime) in self.regimes.iter().enumerate() {
            let prob = self
                .matrix
                .get_transition_prob(self.current_regime, next_regime);
            if prob > max_prob {
                max_prob = prob;
                most_likely_idx = idx;
            }
        }

        // Feature 218: Shannon entropy H = -Σ P(i→j) log₂ P(i→j)
        let entropy: f64 = self.regimes.iter()
            .map(|&next| self.matrix.get_transition_prob(self.current_regime, next))
            .filter(|&p| p > 1e-10) // Numerical stability: avoid log(0)
            .map(|p| -p * p.log2())
            .sum();

        // Feature 219: Expected duration (REUSE existing method!)
        let duration = self.matrix.get_expected_duration(self.current_regime);

        // Feature 220: Change probability (1 - stability)
        let change_prob = 1.0 - stability;

        [
            stability,
            most_likely_idx as f64,
            entropy,
            duration,
            change_prob,
        ]
    }

    /// Get current market regime
    ///
    /// # Returns
    ///
    /// Current regime being tracked
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_probability_features::TransitionProbabilityFeatures;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull];
    /// let mut features = TransitionProbabilityFeatures::new(regimes, 0.1, 1);
    ///
    /// features.update(MarketRegime::Bull);
    /// assert_eq!(features.current_regime(), MarketRegime::Bull);
    /// ```
    pub fn current_regime(&self) -> MarketRegime {
        self.current_regime
    }

    /// Get reference to underlying transition matrix (for advanced use)
    ///
    /// Allows direct access to transition probabilities and stationary distribution
    /// when needed for debugging or analysis.
    ///
    /// # Returns
    ///
    /// Reference to the underlying `RegimeTransitionMatrix`
    pub fn transition_matrix(&self) -> &RegimeTransitionMatrix {
        &self.matrix
    }
}

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

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

        let features = TransitionProbabilityFeatures::new(regimes, 0.1, 10);
        assert_eq!(features.current_regime(), MarketRegime::Bear);
    }

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

        let features = TransitionProbabilityFeatures::new(regimes, 0.1, 10);
        let result = features.compute_features();

        assert_eq!(result.len(), 5);
    }

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

        let mut features = TransitionProbabilityFeatures::new(regimes, 0.2, 1);

        for _ in 0..10 {
            features.update(MarketRegime::Sideways);
        }

        let result = features.compute_features();
        let stability = result[0];

        assert!(
            stability >= 0.0 && stability <= 1.0,
            "Stability should be in [0,1], got {}",
            stability
        );
    }

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

        let mut features = TransitionProbabilityFeatures::new(regimes, 0.2, 1);

        features.update(MarketRegime::Bull);
        features.update(MarketRegime::Bear);

        let result = features.compute_features();
        let entropy = result[2];

        assert!(
            entropy >= 0.0,
            "Entropy should be non-negative, got {}",
            entropy
        );
        assert!(
            entropy.is_finite(),
            "Entropy should be finite, got {}",
            entropy
        );
    }

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

        let mut features = TransitionProbabilityFeatures::new(regimes, 0.2, 1);

        features.update(MarketRegime::Bull);
        features.update(MarketRegime::Bear);

        let result = features.compute_features();
        let stability = result[0];
        let change_prob = result[4];

        assert!(
            (stability + change_prob - 1.0).abs() < 1e-10,
            "Stability + change probability should equal 1.0, got {} + {} = {}",
            stability,
            change_prob,
            stability + change_prob
        );
    }
}