//! Regime Transition Matrix
//!
//! Implements a Markov chain transition matrix for modeling regime change probabilities.
//! Tracks regime transitions and calculates:
//! - Transition probabilities P(regime_t | regime_{t-1})
//! - Stationary distribution (long-run regime probabilities)
//! - Expected regime duration (1 / (1 - P(i->i)))
//!
//! Uses exponential moving average for online updates and Laplace smoothing
//! to handle sparse transitions.

use crate::MarketRegime;
use std::collections::HashMap;

/// Regime Transition Matrix
///
/// N×N matrix tracking transition probabilities between N market regimes.
/// Updates use exponential moving average (EMA) for online learning.
///
/// # Mathematical Foundation
///
/// For a Markov chain with transition matrix P:
/// - `P[i][j]` = P(state_t = j | state_{t-1} = i)
/// - Row sums: Σ_j `P[i][j]` = 1.0 for all i
/// - Stationary distribution π: π = πP (eigenvector with eigenvalue 1)
/// - Expected duration in state i: `E[T_i]` = 1 / (1 - `P[i][i]`)
///
/// # Example
///
/// ```rust
/// use ml::regime::transition_matrix::RegimeTransitionMatrix;
/// use ml::ensemble::MarketRegime;
///
/// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
/// let mut matrix = RegimeTransitionMatrix::new(regimes, 0.1, 10);
///
/// // Update with observed transitions
/// matrix.update(MarketRegime::Bull, MarketRegime::Bear);
/// matrix.update(MarketRegime::Bear, MarketRegime::Bull);
///
/// // Query transition probability
/// let prob = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
/// println!("P(Bull->Bear) = {:.4}", prob);
///
/// // Get stationary distribution
/// let stationary = matrix.get_stationary_distribution();
/// println!("Stationary: {:?}", stationary);
/// ```
#[derive(Debug, Clone)]
pub struct RegimeTransitionMatrix {
    /// Ordered list of regimes (defines matrix indices)
    regimes: Vec<MarketRegime>,

    /// N×N transition probability matrix
    /// ``P`[i]``[j]` = P(regime_t = j | regime_{t-1} = i)
    transition_matrix: Vec<Vec<f64>>,

    /// N×N transition count matrix (raw observations)
    transition_counts: Vec<Vec<usize>>,

    /// EMA smoothing factor (0 < alpha <= 1)
    /// Higher values give more weight to recent observations
    smoothing_alpha: f64,

    /// Minimum observations before using empirical probabilities
    /// Below this threshold, uses Laplace smoothing
    min_observations: usize,

    /// Regime index lookup (for O(1) access)
    regime_to_index: HashMap<MarketRegime, usize>,
}

impl RegimeTransitionMatrix {
    /// Create a new regime transition matrix
    ///
    /// # 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 transition matrix with uniform initial probabilities (1/N for each transition)
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_matrix::RegimeTransitionMatrix;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![
    ///     MarketRegime::Bull,
    ///     MarketRegime::Bear,
    ///     MarketRegime::Sideways,
    /// ];
    /// let matrix = RegimeTransitionMatrix::new(regimes, 0.1, 10);
    /// ```
    pub fn new(regimes: Vec<MarketRegime>, alpha: f64, min_obs: usize) -> Self {
        let n = regimes.len();
        let uniform_prob = 1.0 / n as f64;

        // Initialize transition matrix with uniform probabilities
        let transition_matrix = vec![vec![uniform_prob; n]; n];

        // Initialize counts to zero
        let transition_counts = vec![vec![0; n]; n];

        // Build regime index lookup
        let regime_to_index: HashMap<MarketRegime, usize> = regimes
            .iter()
            .enumerate()
            .map(|(i, &regime)| (regime, i))
            .collect();

        Self {
            regimes,
            transition_matrix,
            transition_counts,
            smoothing_alpha: alpha.clamp(0.01, 1.0), // Ensure valid range
            min_observations: min_obs,
            regime_to_index,
        }
    }

    /// Update transition matrix with observed regime transition
    ///
    /// Uses exponential moving average (EMA) for online updates:
    /// `P_new[i][j]` = (1 - alpha) * `P_old[i][j]` + alpha * `delta[i][j]`
    ///
    /// where `delta[i][j]` = 1 if transition i->j observed, else 0
    ///
    /// # Arguments
    ///
    /// * `from` - Previous regime
    /// * `to` - Current regime
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_matrix::RegimeTransitionMatrix;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
    /// let mut matrix = RegimeTransitionMatrix::new(regimes, 0.1, 1);
    ///
    /// // Observe transition: Bull -> Bear
    /// matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    /// ```
    pub fn update(&mut self, from: MarketRegime, to: MarketRegime) {
        let from_idx = match self.regime_to_index.get(&from) {
            Some(&idx) => idx,
            None => return, // Unknown regime, skip
        };

        let to_idx = match self.regime_to_index.get(&to) {
            Some(&idx) => idx,
            None => return, // Unknown regime, skip
        };

        let n = self.regimes.len();
        let alpha = self.smoothing_alpha;

        // Increment count
        self.transition_counts[from_idx][to_idx] += 1;

        // Update probabilities using EMA
        // P_new[i][j] = (1 - alpha) * P_old[i][j] + alpha * delta[i][j]
        for j in 0..n {
            if j == to_idx {
                // Observed transition: increase probability
                self.transition_matrix[from_idx][j] =
                    (1.0 - alpha) * self.transition_matrix[from_idx][j] + alpha;
            } else {
                // Other transitions: decrease probability
                self.transition_matrix[from_idx][j] *= 1.0 - alpha;
            }
        }

        // Normalize row to ensure probabilities sum to 1.0
        self.normalize_row(from_idx);
    }

    /// Get transition probability P(to | from)
    ///
    /// Returns the probability of transitioning from `from` regime to `to` regime.
    /// If insufficient observations, uses Laplace smoothing.
    ///
    /// # Arguments
    ///
    /// * `from` - Source regime
    /// * `to` - Target regime
    ///
    /// # Returns
    ///
    /// Transition probability in [0, 1]
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_matrix::RegimeTransitionMatrix;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
    /// let mut matrix = RegimeTransitionMatrix::new(regimes, 0.1, 1);
    ///
    /// matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    ///
    /// let prob = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
    /// println!("P(Bull->Bear) = {:.4}", prob);
    /// ```
    pub fn get_transition_prob(&self, from: MarketRegime, to: MarketRegime) -> f64 {
        let from_idx = match self.regime_to_index.get(&from) {
            Some(&idx) => idx,
            None => return 0.0, // Unknown regime
        };

        let to_idx = match self.regime_to_index.get(&to) {
            Some(&idx) => idx,
            None => return 0.0, // Unknown regime
        };

        // Check if we have sufficient observations
        let total_from_count: usize = self.transition_counts[from_idx].iter().sum();

        if total_from_count < self.min_observations {
            // Insufficient data: use Laplace smoothing
            // P[i][j] = (count[i][j] + 1) / (total_count[i] + N)
            let n = self.regimes.len();
            let count = self.transition_counts[from_idx][to_idx] as f64;
            (count + 1.0) / (total_from_count as f64 + n as f64)
        } else {
            // Sufficient data: use EMA-updated probabilities
            self.transition_matrix[from_idx][to_idx]
        }
    }

    /// Get stationary distribution of the Markov chain
    ///
    /// Computes the long-run probability distribution π where π = πP.
    /// Uses iterative power method: π^(k+1) = π^(k) * P until convergence.
    ///
    /// # Returns
    ///
    /// HashMap mapping each regime to its stationary probability.
    /// Probabilities sum to 1.0.
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_matrix::RegimeTransitionMatrix;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
    /// let mut matrix = RegimeTransitionMatrix::new(regimes, 0.1, 1);
    ///
    /// // Add symmetric transitions
    /// for _ in 0..10 {
    ///     matrix.update(MarketRegime::Bull, MarketRegime::Bear);
    ///     matrix.update(MarketRegime::Bear, MarketRegime::Bull);
    /// }
    ///
    /// let stationary = matrix.get_stationary_distribution();
    /// println!("Stationary: {:?}", stationary);
    /// ```
    pub fn get_stationary_distribution(&self) -> HashMap<MarketRegime, f64> {
        let n = self.regimes.len();

        // Start with uniform distribution
        let mut pi = vec![1.0 / n as f64; n];

        // Power iteration: π^(k+1) = π^(k) * P
        // Converge when ||π^(k+1) - π^(k)|| < epsilon
        let max_iterations = 1000;
        let epsilon = 1e-8;

        for _ in 0..max_iterations {
            let mut pi_new = vec![0.0; n];

            // Matrix multiplication: π_new = π * P
            for j in 0..n {
                for i in 0..n {
                    pi_new[j] += pi[i] * self.transition_matrix[i][j];
                }
            }

            // Check convergence
            let delta: f64 = pi_new
                .iter()
                .zip(pi.iter())
                .map(|(new, old)| (new - old).abs())
                .sum();

            pi = pi_new;

            if delta < epsilon {
                break;
            }
        }

        // Normalize to ensure sum = 1.0 (numerical stability)
        let sum: f64 = pi.iter().sum();
        if sum > 0.0 {
            pi.iter_mut().for_each(|p| *p /= sum);
        }

        // Convert to HashMap
        self.regimes
            .iter()
            .zip(pi.iter())
            .map(|(&regime, &prob)| (regime, prob))
            .collect()
    }

    /// Get expected duration in a regime
    ///
    /// Calculates `E[T_i]` = 1 / (1 - `P[i][i]`), the expected number of time steps
    /// spent in regime i before transitioning out.
    ///
    /// # Arguments
    ///
    /// * `regime` - Regime to analyze
    ///
    /// # Returns
    ///
    /// Expected duration (number of periods). Returns 1.0 if regime has zero persistence.
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::regime::transition_matrix::RegimeTransitionMatrix;
    /// use ml::ensemble::MarketRegime;
    ///
    /// let regimes = vec![MarketRegime::Sideways];
    /// let mut matrix = RegimeTransitionMatrix::new(regimes, 0.1, 1);
    ///
    /// // Make Sideways persistent
    /// for _ in 0..10 {
    ///     matrix.update(MarketRegime::Sideways, MarketRegime::Sideways);
    /// }
    ///
    /// let duration = matrix.get_expected_duration(MarketRegime::Sideways);
    /// println!("Expected duration: {:.2} periods", duration);
    /// ```
    pub fn get_expected_duration(&self, regime: MarketRegime) -> f64 {
        let idx = match self.regime_to_index.get(&regime) {
            Some(&i) => i,
            None => return 1.0, // Unknown regime
        };

        let self_prob = self.transition_matrix[idx][idx];

        // E[T] = 1 / (1 - P[i][i])
        // Clamp to avoid division by zero or negative values
        let exit_prob = (1.0 - self_prob).max(0.001);
        1.0 / exit_prob
    }

    /// Get number of regimes tracked
    pub fn regime_count(&self) -> usize {
        self.regimes.len()
    }

    /// Get the list of regimes tracked by this matrix
    ///
    /// # Returns
    ///
    /// Reference to the vector of regimes in index order
    pub fn get_regimes(&self) -> &Vec<MarketRegime> {
        &self.regimes
    }

    /// Normalize a row of the transition matrix to sum to 1.0
    ///
    /// Ensures row sums equal 1.0 (probability distribution property).
    fn normalize_row(&mut self, row_idx: usize) {
        let row_sum: f64 = self.transition_matrix[row_idx].iter().sum();

        if row_sum > 0.0 {
            let n = self.regimes.len();
            for j in 0..n {
                self.transition_matrix[row_idx][j] /= row_sum;
            }
        }
    }
}

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

    #[test]
    fn test_new_initialization() {
        let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
        let matrix = RegimeTransitionMatrix::new(regimes, 0.1, 10);

        assert_eq!(matrix.regime_count(), 2);

        // Check uniform initialization
        for from in [MarketRegime::Bull, MarketRegime::Bear] {
            for to in [MarketRegime::Bull, MarketRegime::Bear] {
                let prob = matrix.get_transition_prob(from, to);
                assert!((prob - 0.5).abs() < 1e-6);
            }
        }
    }

    #[test]
    fn test_update_and_normalization() {
        let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
        let mut matrix = RegimeTransitionMatrix::new(regimes, 0.5, 1);

        matrix.update(MarketRegime::Bull, MarketRegime::Bear);

        // Row should sum to 1.0
        let p_bull_bear = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
        let p_bull_bull = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bull);

        let sum = p_bull_bear + p_bull_bull;
        assert!(
            (sum - 1.0).abs() < 1e-6,
            "Row should sum to 1.0, got {}",
            sum
        );
    }

    #[test]
    fn test_laplace_smoothing() {
        let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
        let matrix = RegimeTransitionMatrix::new(regimes, 0.1, 10);

        // With zero observations and min_obs=10, should use Laplace smoothing
        // P[i][j] = (0 + 1) / (0 + 2) = 0.5
        let prob = matrix.get_transition_prob(MarketRegime::Bull, MarketRegime::Bear);
        assert!((prob - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_stationary_convergence() {
        let regimes = vec![MarketRegime::Bull, MarketRegime::Bear];
        let matrix = RegimeTransitionMatrix::new(regimes, 0.1, 1);

        let stationary = matrix.get_stationary_distribution();

        // Uniform transition matrix -> uniform stationary distribution
        let bull_prob = stationary.get(&MarketRegime::Bull).unwrap();
        let bear_prob = stationary.get(&MarketRegime::Bear).unwrap();

        assert!((bull_prob - 0.5).abs() < 0.01);
        assert!((bear_prob - 0.5).abs() < 0.01);

        // Sum to 1.0
        let sum: f64 = stationary.values().sum();
        assert!((sum - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_expected_duration() {
        let regimes = vec![MarketRegime::Sideways];
        let mut matrix = RegimeTransitionMatrix::new(regimes, 0.2, 1);

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

        let duration = matrix.get_expected_duration(MarketRegime::Sideways);

        // Should be > 1 (high persistence)
        assert!(duration > 1.0);
    }
}
