//! Hidden Markov Model for regime detection.
//!
//! Gaussian-emission HMM with:
//! - Baum-Welch (EM) training — O(T * N^2) per iteration
//! - Viterbi decoding — O(T * N^2) most-likely state sequence
//! - Online forward filter — O(N^2) per observation (real-time)
//! - One-step-ahead prediction — O(N^2)
//! - Forward-backward smoothing — O(T * N^2)
//!
//! N = number of states (typically 2-3), T = observations.
//! At N=3 this is ~9 multiplies per tick — effectively free.

use pyo3::prelude::*;
use std::sync::Mutex;

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const MAX_STATES: usize = 8;
const MAX_EM_ITERS: usize = 500;
const EM_TOL: f64 = 1e-8;
const LOG_2PI: f64 = 1.8378770664093453; // ln(2*pi)
const MIN_VARIANCE: f64 = 1e-12;

// ---------------------------------------------------------------------------
// MarkovRegimeModel
// ---------------------------------------------------------------------------

/// Hidden Markov Model with Gaussian emissions for regime detection.
///
/// Each state has its own (mean, variance) emission distribution.
/// Designed for financial return series: state 0 = low vol, state 1 = high vol, etc.
#[pyclass]
pub struct MarkovRegimeModel {
    inner: Mutex<HmmInner>,
}

struct HmmInner {
    n_states: usize,
    /// Transition matrix: trans[i][j] = P(state_j | state_i)
    trans: Vec<Vec<f64>>,
    /// Initial state probabilities
    initial: Vec<f64>,
    /// Emission means per state
    means: Vec<f64>,
    /// Emission variances per state
    variances: Vec<f64>,
    /// Current filtered state probabilities (for online use)
    filtered: Vec<f64>,
    /// Whether the model has been trained
    trained: bool,
}

impl HmmInner {
    fn new(n_states: usize) -> Result<Self, String> {
        if n_states < 2 || n_states > MAX_STATES {
            return Err(format!(
                "n_states must be between 2 and {}, got {}",
                MAX_STATES, n_states
            ));
        }
        let uniform = 1.0 / n_states as f64;

        // Default: uniform transitions, uniform initial
        let trans = vec![vec![uniform; n_states]; n_states];
        let initial = vec![uniform; n_states];

        // Default emissions: spread means around 0, increasing variance
        let mut means = vec![0.0; n_states];
        let mut variances = vec![0.0; n_states];
        for i in 0..n_states {
            means[i] = 0.0;
            variances[i] = 0.001 * (1.0 + i as f64 * 2.0); // 0.001, 0.003, 0.005, ...
        }

        let filtered = vec![uniform; n_states];

        Ok(Self {
            n_states,
            trans,
            initial,
            means,
            variances,
            filtered,
            trained: false,
        })
    }

    /// Gaussian log-likelihood: ln N(x | mu, sigma^2)
    #[inline(always)]
    fn log_emission(&self, state: usize, x: f64) -> f64 {
        let mu = self.means[state];
        let var = self.variances[state].max(MIN_VARIANCE);
        -0.5 * (LOG_2PI + var.ln() + (x - mu).powi(2) / var)
    }

    /// Gaussian pdf: N(x | mu, sigma^2)
    #[inline(always)]
    fn emission(&self, state: usize, x: f64) -> f64 {
        self.log_emission(state, x).exp()
    }

    /// Forward algorithm. Returns (alpha, scaling_factors, log_likelihood).
    /// alpha[t][i] = scaled P(state_i, obs_0..t)
    fn forward(&self, obs: &[f64]) -> (Vec<Vec<f64>>, Vec<f64>, f64) {
        let n = self.n_states;
        let t_len = obs.len();
        let mut alpha = vec![vec![0.0; n]; t_len];
        let mut scales = vec![0.0; t_len];

        // t=0
        for i in 0..n {
            alpha[0][i] = self.initial[i] * self.emission(i, obs[0]);
        }
        let s0: f64 = alpha[0].iter().sum();
        scales[0] = if s0 > 0.0 { 1.0 / s0 } else { 1.0 };
        for i in 0..n {
            alpha[0][i] *= scales[0];
        }

        // t=1..T-1
        for t in 1..t_len {
            for j in 0..n {
                let mut sum = 0.0;
                for i in 0..n {
                    sum += alpha[t - 1][i] * self.trans[i][j];
                }
                alpha[t][j] = sum * self.emission(j, obs[t]);
            }
            let st: f64 = alpha[t].iter().sum();
            scales[t] = if st > 0.0 { 1.0 / st } else { 1.0 };
            for j in 0..n {
                alpha[t][j] *= scales[t];
            }
        }

        let log_lik: f64 = -scales.iter().map(|s| s.ln()).sum::<f64>();
        (alpha, scales, log_lik)
    }

    /// Backward algorithm (using scaling factors from forward pass).
    fn backward(&self, obs: &[f64], scales: &[f64]) -> Vec<Vec<f64>> {
        let n = self.n_states;
        let t_len = obs.len();
        let mut beta = vec![vec![0.0; n]; t_len];

        // t=T-1
        for i in 0..n {
            beta[t_len - 1][i] = scales[t_len - 1];
        }

        // t=T-2..0
        for t in (0..t_len - 1).rev() {
            for i in 0..n {
                let mut sum = 0.0;
                for j in 0..n {
                    sum += self.trans[i][j] * self.emission(j, obs[t + 1]) * beta[t + 1][j];
                }
                beta[t][i] = sum * scales[t];
            }
        }

        beta
    }

    /// Baum-Welch EM training.
    fn fit(&mut self, obs: &[f64], max_iters: usize, tol: f64) -> f64 {
        let n = self.n_states;
        let t_len = obs.len();
        if t_len < 2 {
            return f64::NEG_INFINITY;
        }

        // Initialize means/variances from data using quantile-based heuristic
        self.init_from_data(obs);

        let mut prev_ll = f64::NEG_INFINITY;

        for _iter in 0..max_iters {
            // E-step: forward-backward
            let (alpha, scales, log_lik) = self.forward(obs);
            let beta = self.backward(obs, &scales);

            // Check convergence
            if (log_lik - prev_ll).abs() < tol {
                prev_ll = log_lik;
                break;
            }
            prev_ll = log_lik;

            // Compute gamma[t][i] = P(state_i at time t | obs)
            let mut gamma = vec![vec![0.0; n]; t_len];
            for t in 0..t_len {
                let mut denom = 0.0;
                for i in 0..n {
                    gamma[t][i] = alpha[t][i] * beta[t][i];
                    denom += gamma[t][i];
                }
                if denom > 0.0 {
                    for i in 0..n {
                        gamma[t][i] /= denom;
                    }
                }
            }

            // Compute xi[t][i][j] = P(state_i at t, state_j at t+1 | obs)
            // We accumulate sums directly to save memory
            let mut xi_sum = vec![vec![0.0; n]; n]; // [i][j]
            let mut xi_t = vec![vec![0.0; n]; n]; // allocated once, reused per timestep
            for t in 0..t_len - 1 {
                let mut denom = 0.0;
                // Compute all xi values for this t (reuse xi_t buffer)
                for i in 0..n {
                    for j in 0..n {
                        xi_t[i][j] = alpha[t][i]
                            * self.trans[i][j]
                            * self.emission(j, obs[t + 1])
                            * beta[t + 1][j];
                        denom += xi_t[i][j];
                    }
                }
                if denom > 0.0 {
                    for i in 0..n {
                        for j in 0..n {
                            xi_sum[i][j] += xi_t[i][j] / denom;
                        }
                    }
                }
            }

            // M-step: update parameters
            // Initial probabilities
            for i in 0..n {
                self.initial[i] = gamma[0][i].max(1e-300);
            }
            let init_sum: f64 = self.initial.iter().sum();
            for i in 0..n {
                self.initial[i] /= init_sum;
            }

            // Transition matrix
            for i in 0..n {
                let gamma_sum: f64 = (0..t_len - 1).map(|t| gamma[t][i]).sum();
                if gamma_sum > 0.0 {
                    for j in 0..n {
                        self.trans[i][j] = (xi_sum[i][j] / gamma_sum).max(1e-300);
                    }
                }
                // Normalize row
                let row_sum: f64 = self.trans[i].iter().sum();
                if row_sum > 0.0 {
                    for j in 0..n {
                        self.trans[i][j] /= row_sum;
                    }
                }
            }

            // Emission means and variances
            for i in 0..n {
                let gamma_sum: f64 = gamma.iter().map(|g| g[i]).sum();
                if gamma_sum > 0.0 {
                    // Mean
                    let weighted_sum: f64 = (0..t_len).map(|t| gamma[t][i] * obs[t]).sum();
                    self.means[i] = weighted_sum / gamma_sum;

                    // Variance
                    let var_sum: f64 = (0..t_len)
                        .map(|t| gamma[t][i] * (obs[t] - self.means[i]).powi(2))
                        .sum();
                    self.variances[i] = (var_sum / gamma_sum).max(MIN_VARIANCE);
                }
            }
        }

        self.trained = true;
        // Reset filtered state to initial
        self.filtered = self.initial.clone();
        prev_ll
    }

    /// Initialize emission params from data using sorted quantiles.
    fn init_from_data(&mut self, obs: &[f64]) {
        let n = self.n_states;
        let t_len = obs.len();

        let mut sorted = obs.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        // Split data into n_states quantile groups
        let chunk = (t_len + n - 1) / n;
        for i in 0..n {
            let start = i * chunk;
            let end = ((i + 1) * chunk).min(t_len);
            if start >= t_len {
                break;
            }
            let slice = &sorted[start..end];
            let mean: f64 = slice.iter().sum::<f64>() / slice.len() as f64;
            let var: f64 = slice.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
                / slice.len() as f64;
            self.means[i] = mean;
            self.variances[i] = var.max(MIN_VARIANCE);
        }

        // Sort states by variance (state 0 = lowest vol)
        let mut indices: Vec<usize> = (0..n).collect();
        indices.sort_by(|&a, &b| {
            self.variances[a]
                .partial_cmp(&self.variances[b])
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let old_means = self.means.clone();
        let old_vars = self.variances.clone();
        for (new_i, &old_i) in indices.iter().enumerate() {
            self.means[new_i] = old_means[old_i];
            self.variances[new_i] = old_vars[old_i];
        }
    }

    /// Viterbi decoding: most likely state sequence.
    fn viterbi(&self, obs: &[f64]) -> Vec<usize> {
        let n = self.n_states;
        let t_len = obs.len();
        if t_len == 0 {
            return vec![];
        }

        let mut v = vec![vec![f64::NEG_INFINITY; n]; t_len];
        let mut path = vec![vec![0usize; n]; t_len];

        // t=0
        for i in 0..n {
            v[0][i] = self.initial[i].max(1e-300).ln() + self.log_emission(i, obs[0]);
        }

        // t=1..T-1
        for t in 1..t_len {
            for j in 0..n {
                let mut best_val = f64::NEG_INFINITY;
                let mut best_i = 0;
                for i in 0..n {
                    let val = v[t - 1][i] + self.trans[i][j].max(1e-300).ln();
                    if val > best_val {
                        best_val = val;
                        best_i = i;
                    }
                }
                v[t][j] = best_val + self.log_emission(j, obs[t]);
                path[t][j] = best_i;
            }
        }

        // Backtrack
        let mut states = vec![0usize; t_len];
        states[t_len - 1] = v[t_len - 1]
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);

        for t in (0..t_len - 1).rev() {
            states[t] = path[t + 1][states[t + 1]];
        }

        states
    }

    /// Online forward filter: process one observation, update state probabilities.
    /// Returns the new filtered state probabilities. O(N^2).
    fn filter_step(&mut self, obs: f64) -> Vec<f64> {
        let n = self.n_states;
        let mut new_filtered = vec![0.0; n];

        // Predict: multiply by transition matrix
        for j in 0..n {
            let mut sum = 0.0;
            for i in 0..n {
                sum += self.filtered[i] * self.trans[i][j];
            }
            // Update: multiply by emission
            new_filtered[j] = sum * self.emission(j, obs);
        }

        // Normalize
        let total: f64 = new_filtered.iter().sum();
        if total > 0.0 {
            for j in 0..n {
                new_filtered[j] /= total;
            }
        } else {
            // Fallback to uniform
            for j in 0..n {
                new_filtered[j] = 1.0 / n as f64;
            }
        }

        self.filtered = new_filtered.clone();
        new_filtered
    }

    /// One-step-ahead prediction: P(state at t+1 | obs_0..t).
    fn predict(&self) -> Vec<f64> {
        let n = self.n_states;
        let mut pred = vec![0.0; n];
        for j in 0..n {
            for i in 0..n {
                pred[j] += self.filtered[i] * self.trans[i][j];
            }
        }
        pred
    }

    /// Forward-backward smoothing: P(state_i at t | all observations).
    fn smooth(&self, obs: &[f64]) -> Vec<Vec<f64>> {
        let n = self.n_states;
        let t_len = obs.len();
        if t_len == 0 {
            return vec![];
        }

        let (alpha, scales, _) = self.forward(obs);
        let beta = self.backward(obs, &scales);

        let mut gamma = vec![vec![0.0; n]; t_len];
        for t in 0..t_len {
            let mut denom = 0.0;
            for i in 0..n {
                gamma[t][i] = alpha[t][i] * beta[t][i];
                denom += gamma[t][i];
            }
            if denom > 0.0 {
                for i in 0..n {
                    gamma[t][i] /= denom;
                }
            }
        }

        gamma
    }
}

// ---------------------------------------------------------------------------
// PyO3 interface
// ---------------------------------------------------------------------------

#[pymethods]
impl MarkovRegimeModel {
    /// Create a new HMM with n_states Gaussian emission states.
    ///
    /// Args:
    ///     n_states: Number of hidden states (2-8). Default 2 (calm/volatile).
    #[new]
    #[pyo3(signature = (n_states=2))]
    fn new(n_states: usize) -> PyResult<Self> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if n_states < 2 || n_states > MAX_STATES {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "n_states must be between 2 and {}, got {}",
                MAX_STATES, n_states
            )));
        }
        let hmm = HmmInner::new(n_states).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(e)
        })?;
        Ok(Self {
            inner: Mutex::new(hmm),
        })
    }

    /// Train the model on a sequence of observations using Baum-Welch EM.
    ///
    /// Args:
    ///     observations: List of float values (e.g., log returns).
    ///     max_iters: Maximum EM iterations (default 100).
    ///     tol: Convergence tolerance on log-likelihood (default 1e-6).
    ///
    /// Returns:
    ///     Final log-likelihood.
    #[pyo3(signature = (observations, max_iters=100, tol=1e-6))]
    fn fit(&self, observations: Vec<f64>, max_iters: usize, tol: f64) -> PyResult<f64> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if observations.len() < 2 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "need at least 2 observations to train",
            ));
        }
        // Validate no NaN/Inf
        for (i, &v) in observations.iter().enumerate() {
            if !v.is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "observation[{}] is not finite: {}",
                    i, v
                )));
            }
        }
        let max_iters = max_iters.min(MAX_EM_ITERS);
        let mut inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        Ok(inner.fit(&observations, max_iters, tol))
    }

    /// Viterbi decoding: find the most likely state sequence.
    ///
    /// Args:
    ///     observations: List of float values.
    ///
    /// Returns:
    ///     List of state indices (0 = lowest variance state).
    fn decode(&self, observations: Vec<f64>) -> PyResult<Vec<usize>> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if !inner.trained {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "model must be trained before decoding (call fit() first)",
            ));
        }
        Ok(inner.viterbi(&observations))
    }

    /// Online filter: process one observation and return updated state probabilities.
    ///
    /// This is the real-time method — call once per tick with the latest return.
    /// O(N^2) where N = number of states. At N=3, this is 9 multiplies.
    ///
    /// Args:
    ///     observation: Single float value (e.g., latest log return).
    ///
    /// Returns:
    ///     List of state probabilities [P(state_0), P(state_1), ...].
    fn filter_step(&self, observation: f64) -> PyResult<Vec<f64>> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if !observation.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "observation must be finite",
            ));
        }
        let mut inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if !inner.trained {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "model must be trained before filtering (call fit() first)",
            ));
        }
        Ok(inner.filter_step(observation))
    }

    /// One-step-ahead regime prediction from current filtered state.
    ///
    /// Returns:
    ///     List of predicted state probabilities for the next time step.
    fn predict(&self) -> PyResult<Vec<f64>> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if !inner.trained {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "model must be trained before predicting (call fit() first)",
            ));
        }
        Ok(inner.predict())
    }

    /// Forward-backward smoothing: compute P(state | all observations) for each timestep.
    ///
    /// Args:
    ///     observations: Full observation sequence.
    ///
    /// Returns:
    ///     List of lists: smoothed[t][state] = P(state at t | all obs).
    fn smooth(&self, observations: Vec<f64>) -> PyResult<Vec<Vec<f64>>> {
        crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        if !inner.trained {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "model must be trained before smoothing (call fit() first)",
            ));
        }
        Ok(inner.smooth(&observations))
    }

    /// Get the current filtered state probabilities.
    fn filtered_probs(&self) -> Vec<f64> {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner.filtered.clone()
    }

    /// Reset the filter to initial state probabilities.
    fn reset_filter(&self) {
        let mut inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner.filtered = inner.initial.clone();
    }

    /// Get the number of hidden states.
    fn n_states(&self) -> usize {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner.n_states
    }

    /// Get the transition matrix as a list of lists.
    fn transition_matrix(&self) -> Vec<Vec<f64>> {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner.trans.clone()
    }

    /// Get emission parameters: list of (mean, variance) per state.
    fn emission_params(&self) -> Vec<(f64, f64)> {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        (0..inner.n_states)
            .map(|i| (inner.means[i], inner.variances[i]))
            .collect()
    }

    /// Get the most likely current regime (argmax of filtered probabilities).
    fn current_regime(&self) -> usize {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner
            .filtered
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0)
    }

    /// Check if the model has been trained.
    fn is_trained(&self) -> bool {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        inner.trained
    }

    fn __repr__(&self) -> String {
        let inner = self
            .inner
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        let state_info: Vec<String> = (0..inner.n_states)
            .map(|i| format!("s{}(mu={:.4}, var={:.6})", i, inner.means[i], inner.variances[i]))
            .collect();
        format!(
            "MarkovRegimeModel(n_states={}, trained={}, states=[{}])",
            inner.n_states,
            inner.trained,
            state_info.join(", ")
        )
    }
}

/// Convenience: compute log returns from a price series.
///
/// Returns vec of length len(prices)-1.
#[pyfunction]
pub fn prices_to_returns(prices: Vec<f64>) -> PyResult<Vec<f64>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if prices.len() < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 2 prices",
        ));
    }
    let mut returns = Vec::with_capacity(prices.len() - 1);
    for i in 1..prices.len() {
        if prices[i - 1] <= 0.0 {
            returns.push(0.0);
        } else {
            returns.push((prices[i] / prices[i - 1]).ln());
        }
    }
    Ok(returns)
}

/// Register Markov types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<MarkovRegimeModel>()?;
    m.add_function(wrap_pyfunction!(prices_to_returns, m)?)?;
    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_two_regime_data() -> Vec<f64> {
        // 200 observations: first 100 from N(0, 0.01), next 100 from N(0, 0.05)
        let mut data = Vec::with_capacity(200);
        // Simple LCG for determinism
        let mut rng: u64 = 42;
        for i in 0..200 {
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u = (rng >> 33) as f64 / (1u64 << 31) as f64; // [0, 1)
            let noise = (u - 0.5) * 2.0; // [-1, 1)
            let sigma = if i < 100 { 0.01 } else { 0.05 };
            data.push(noise * sigma);
        }
        data
    }

    #[test]
    fn test_create_model() {
        let m = HmmInner::new(2).unwrap();
        assert_eq!(m.n_states, 2);
        assert_eq!(m.trans.len(), 2);
        assert_eq!(m.means.len(), 2);
    }

    #[test]
    fn test_fit_two_states() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(2).unwrap();
        let ll = m.fit(&data, 50, 1e-6);
        assert!(ll.is_finite());
        assert!(m.trained);
        // State 0 should have lower variance than state 1
        assert!(m.variances[0] < m.variances[1]);
    }

    #[test]
    fn test_viterbi_separation() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(2).unwrap();
        m.fit(&data, 50, 1e-6);
        let states = m.viterbi(&data);
        assert_eq!(states.len(), 200);
        // First half should be mostly state 0, second half mostly state 1
        let first_half_mean: f64 =
            states[0..100].iter().map(|&s| s as f64).sum::<f64>() / 100.0;
        let second_half_mean: f64 =
            states[100..200].iter().map(|&s| s as f64).sum::<f64>() / 100.0;
        assert!(second_half_mean > first_half_mean);
    }

    #[test]
    fn test_filter_step() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(2).unwrap();
        m.fit(&data, 50, 1e-6);

        // Filter through low-vol data
        m.filtered = m.initial.clone();
        for &x in &data[0..20] {
            m.filter_step(x);
        }
        // After low-vol observations, state 0 (low vol) should dominate
        assert!(m.filtered[0] > m.filtered[1]);
    }

    #[test]
    fn test_predict() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(2).unwrap();
        m.fit(&data, 50, 1e-6);
        let pred = m.predict();
        assert_eq!(pred.len(), 2);
        let sum: f64 = pred.iter().sum();
        assert!((sum - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_smooth() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(2).unwrap();
        m.fit(&data, 50, 1e-6);
        let smoothed = m.smooth(&data);
        assert_eq!(smoothed.len(), 200);
        for row in &smoothed {
            let sum: f64 = row.iter().sum();
            assert!((sum - 1.0).abs() < 1e-10);
        }
    }

    #[test]
    fn test_three_states() {
        let data = make_two_regime_data();
        let mut m = HmmInner::new(3).unwrap();
        let ll = m.fit(&data, 50, 1e-6);
        assert!(ll.is_finite());
        assert_eq!(m.n_states, 3);
    }

    #[test]
    fn test_prices_to_returns() {
        let prices = vec![100.0, 101.0, 99.0, 102.0];
        let rets = prices_to_returns(prices).unwrap();
        assert_eq!(rets.len(), 3);
        assert!((rets[0] - (101.0f64 / 100.0).ln()).abs() < 1e-10);
    }

    #[test]
    fn test_empty_data() {
        let rets = prices_to_returns(vec![1.0]);
        assert!(rets.is_err());
    }

    #[test]
    fn test_emission_gaussian() {
        let m = HmmInner::new(2).unwrap();
        // Emission at mean should be highest
        let at_mean = m.emission(0, m.means[0]);
        let away = m.emission(0, m.means[0] + 1.0);
        assert!(at_mean > away);
    }
}
