//! Regime Transition Feature Extractor
//!
//! Wraps the RegimeTransitionMatrix to extract transition probability features
//! for ML models. Tracks current regime and provides features based on transition
//! dynamics between market regimes.
//!
//! # Feature Indices (216-220)
//! - 216: Current regime persistence (self-transition probability)
//! - 217: Most likely next regime (highest transition probability)
//! - 218: Transition entropy (uncertainty in next regime)
//! - 219: Regime stability score (stationary distribution weight)
//! - 220: Expected regime duration (bars until transition)

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

/// Regime Transition Feature Extractor
///
/// Maintains a transition matrix and current regime state to compute
/// features related to regime transition dynamics.
///
/// # Features
/// This extractor computes 5 features (indices 216-220):
/// 1. Persistence: P(current_regime | current_regime)
/// 2. Next regime ID: argmax_j P(j | current_regime)
/// 3. Entropy: -Σ P(j | current_regime) * log(P(j | current_regime))
/// 4. Stability: `π`[current]`` from stationary distribution
/// 5. Duration: Expected bars in current regime
///
/// # Example
/// ```rust,ignore
/// use ml::features::regime_transition::RegimeTransitionFeatures;
/// use ml::ensemble::adaptive_ml_integration::MarketRegime;
///
/// let mut features = RegimeTransitionFeatures::new(4, 0.1);
/// features.update(MarketRegime::Sideways);
/// features.update(MarketRegime::Bull);
/// let feature_vec = features.update(MarketRegime::Bull);
/// ```
#[derive(Debug)]
pub struct RegimeTransitionFeatures {
    /// Underlying transition matrix tracking regime changes
    matrix: RegimeTransitionMatrix,

    /// Current market regime
    current_regime: MarketRegime,
}

impl RegimeTransitionFeatures {
    /// Creates a new RegimeTransitionFeatures extractor.
    ///
    /// # Arguments
    /// * `num_regimes` - Number of distinct regimes to track (typically 4-6)
    /// * `ema_alpha` - EMA smoothing factor for transition updates (0 < alpha <= 1)
    ///
    /// # Returns
    /// New feature extractor with uniform initial transition probabilities.
    ///
    /// # Example
    /// ```rust,ignore
    /// let features = RegimeTransitionFeatures::new(4, 0.1);
    /// ```
    pub fn new(num_regimes: usize, ema_alpha: f64) -> Self {
        // Define standard regime set based on num_regimes
        // Using MarketRegime from adaptive_ml_integration: Normal, Trending, Bull, Bear, Sideways, HighVolatility, Crisis, Unknown
        let regimes = match num_regimes {
            3 => vec![
                MarketRegime::Bull,
                MarketRegime::Bear,
                MarketRegime::Sideways,
            ],
            4 => vec![
                MarketRegime::Bull,
                MarketRegime::Bear,
                MarketRegime::Sideways,
                MarketRegime::HighVolatility,
            ],
            5 => vec![
                MarketRegime::Bull,
                MarketRegime::Bear,
                MarketRegime::Sideways,
                MarketRegime::HighVolatility,
                MarketRegime::Unknown,
            ],
            6 => vec![
                MarketRegime::Normal,
                MarketRegime::Trending,
                MarketRegime::Bull,
                MarketRegime::Bear,
                MarketRegime::Sideways,
                MarketRegime::HighVolatility,
            ],
            _ => vec![
                MarketRegime::Bull,
                MarketRegime::Bear,
                MarketRegime::Sideways,
                MarketRegime::HighVolatility,
            ], // Default to 4 regimes
        };

        Self {
            matrix: RegimeTransitionMatrix::new(regimes, ema_alpha, 10),
            current_regime: MarketRegime::Sideways, // Start with neutral regime
        }
    }

    /// Updates the feature extractor with a new regime observation.
    ///
    /// Records the transition from current_regime to the new regime,
    /// updates the transition matrix, and computes 5 transition features.
    ///
    /// # Arguments
    /// * `regime` - The newly observed market regime
    ///
    /// # Returns
    /// Array of 5 features:
    /// - `[0]`: Persistence probability (how stable is current regime)
    /// - `[1]`: Most likely next regime (encoded as f64: 0-5)
    /// - `[2]`: Transition entropy (uncertainty about next regime)
    /// - `[3]`: Regime stability from stationary distribution
    /// - `[4]`: Expected duration in current regime (bars)
    ///
    /// # Note
    /// Full feature calculation logic will be implemented in Agent D15.2.
    /// This stub ensures the struct compiles and is ready for integration.
    ///
    /// # Example
    /// ```rust,ignore
    /// let mut features = RegimeTransitionFeatures::new(4, 0.1);
    /// let result = features.update(MarketRegime::Trending);
    /// println!("Persistence: {:.4}", result[0]);
    /// println!("Expected duration: {:.2} bars", result[4]);
    /// ```
    pub fn update(&mut self, regime: MarketRegime) -> [f64; 5] {
        // Update transition matrix with observed transition (from current to new regime)
        self.matrix.update(self.current_regime, regime);

        // Update current regime for next iteration
        self.current_regime = regime;

        // Extract all 5 transition probability features (indices 216-220)
        self.compute_features()
    }

    /// Compute all 5 transition probability features
    ///
    /// Returns array of 5 features:
    /// - `[0]`: Feature 216 - Persistence P(i→i)
    /// - `[1]`: Feature 217 - Most likely next regime (index)
    /// - `[2]`: Feature 218 - Shannon entropy
    /// - `[3]`: Feature 219 - Expected duration
    /// - `[4]`: Feature 220 - Change probability
    ///
    /// # Implementation
    ///
    /// Delegates all probability calculations to the underlying RegimeTransitionMatrix
    /// to maintain architectural principle: REUSE existing infrastructure.
    pub fn compute_features(&self) -> [f64; 5] {
        // Feature 216: Persistence P(i→i) - probability of staying in current regime
        let persistence = self
            .matrix
            .get_transition_prob(self.current_regime, self.current_regime);

        // Feature 217: Most likely next regime (index)
        // Find regime with highest transition probability from current regime
        let regimes = self.matrix.get_regimes();
        let mut max_prob = 0.0;
        let mut most_likely_idx = 0;
        for (idx, &next_regime) in 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)
        // Measures uncertainty in regime transitions
        let entropy: f64 = 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!)
        // E[T] = 1 / (1 - P[i][i])
        let duration = self.matrix.get_expected_duration(self.current_regime);

        // Feature 220: Change probability (1 - persistence)
        // Probability of transitioning out of current regime
        let change_prob = 1.0 - persistence;

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

    /// Get current market regime
    ///
    /// Returns the current regime being tracked
    pub fn current_regime(&self) -> MarketRegime {
        self.current_regime
    }

    /// Get reference to underlying transition matrix
    ///
    /// Allows direct access to transition probabilities and stationary distribution
    /// for advanced use cases
    pub fn transition_matrix(&self) -> &RegimeTransitionMatrix {
        &self.matrix
    }
}

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

    #[test]
    fn test_regime_transition_features_new() {
        let features = RegimeTransitionFeatures::new(4, 0.1);
        assert_eq!(features.current_regime, MarketRegime::Sideways);
        assert_eq!(features.matrix.regime_count(), 4);
    }

    #[test]
    fn test_regime_transition_features_new_5_regimes() {
        let features = RegimeTransitionFeatures::new(5, 0.1);
        assert_eq!(features.matrix.regime_count(), 5);
    }

    #[test]
    fn test_regime_transition_features_new_6_regimes() {
        let features = RegimeTransitionFeatures::new(6, 0.1);
        assert_eq!(features.matrix.regime_count(), 6);
    }

    #[test]
    fn test_regime_transition_features_update() {
        let mut features = RegimeTransitionFeatures::new(4, 0.1);
        let result = features.update(MarketRegime::Bull);

        // Verify return type is correct
        assert_eq!(result.len(), 5);

        // Verify current regime is updated
        assert_eq!(features.current_regime, MarketRegime::Bull);

        // Verify all features are finite and within valid bounds
        assert!(
            result.iter().all(|&x| x.is_finite()),
            "All features should be finite"
        );

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

        // Feature 217 (most likely next regime index) should be in [0, 3] for 4 regimes
        assert!(
            result[1] >= 0.0 && result[1] <= 3.0,
            "Most likely index should be in [0, 3], got {}",
            result[1]
        );

        // Feature 218 (entropy) should be non-negative
        assert!(
            result[2] >= 0.0,
            "Entropy should be non-negative, got {}",
            result[2]
        );

        // Feature 219 (expected duration) should be >= 1.0
        assert!(
            result[3] >= 1.0,
            "Expected duration should be >= 1.0, got {}",
            result[3]
        );

        // Feature 220 (change probability) should be in [0, 1]
        assert!(
            result[4] >= 0.0 && result[4] <= 1.0,
            "Change probability should be in [0, 1], got {}",
            result[4]
        );

        // Features 216 and 220 should be complementary (stability + change_prob = 1.0)
        assert!(
            (result[0] + result[4] - 1.0).abs() < 1e-9,
            "Stability + change_prob should equal 1.0, got {} + {} = {}",
            result[0],
            result[4],
            result[0] + result[4]
        );
    }

    #[test]
    fn test_regime_transition_features_multiple_updates() {
        let mut features = RegimeTransitionFeatures::new(4, 0.1);

        let regimes = vec![
            MarketRegime::Sideways,
            MarketRegime::Bull,
            MarketRegime::Bull,
            MarketRegime::Bear,
            MarketRegime::HighVolatility,
        ];

        for regime in regimes {
            let result = features.update(regime);
            assert_eq!(result.len(), 5);
        }

        // Verify final regime state
        assert_eq!(features.current_regime, MarketRegime::HighVolatility);
    }

    #[test]
    fn test_regime_transition_features_default_num_regimes() {
        // Test that invalid num_regimes defaults to 4
        let features = RegimeTransitionFeatures::new(10, 0.1);
        assert_eq!(features.matrix.regime_count(), 4);
    }
}
