//! Simple 3-state Gaussian Hidden Markov Model for offline regime detection.
//!
//! This is a basic HMM implementation with:
//! - Gaussian emissions (single-dimensional observations)
//! - Baum-Welch EM fitting with K-means initialization
//! - Viterbi-like single-observation prediction
//!
//! Intended for offline validation and regime transition estimation,
//! not real-time inference.

use serde::{Deserialize, Serialize};

use ml_core::MLError;

/// Configuration for the Hidden Markov Model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HMMConfig {
    /// Number of hidden states (default: 3)
    pub num_states: usize,
    /// Maximum EM iterations (default: 100)
    pub max_iterations: usize,
    /// Log-likelihood convergence threshold (default: 1e-6)
    pub convergence_threshold: f64,
}

impl Default for HMMConfig {
    fn default() -> Self {
        Self {
            num_states: 3,
            max_iterations: 100,
            convergence_threshold: 1e-6,
        }
    }
}

/// A simple Gaussian-emission Hidden Markov Model.
///
/// Each hidden state emits observations from a univariate Gaussian distribution
/// N(mean, variance). Transition probabilities are stored in a row-stochastic matrix.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HiddenMarkovModel {
    config: HMMConfig,
    /// Row-stochastic transition matrix: `transition_matrix[i][j]` = P(state j | state i)
    transition_matrix: Vec<Vec<f64>>,
    /// Per-state emission means
    emission_means: Vec<f64>,
    /// Per-state emission variances
    emission_vars: Vec<f64>,
    /// Initial state probability distribution
    initial_probs: Vec<f64>,
    /// Whether the model has been fitted
    fitted: bool,
}

impl HiddenMarkovModel {
    /// Create a new HMM with the given configuration.
    ///
    /// The model is initialized with uniform priors and must be fitted before use.
    pub fn new(config: HMMConfig) -> Self {
        let n = config.num_states;
        let uniform = 1.0 / n as f64;

        let transition_matrix = vec![vec![uniform; n]; n];
        let initial_probs = vec![uniform; n];
        let emission_means = vec![0.0; n];
        let emission_vars = vec![1.0; n];

        Self {
            config,
            transition_matrix,
            emission_means,
            emission_vars,
            initial_probs,
            fitted: false,
        }
    }

    /// Fit the HMM to observation data using simplified Baum-Welch EM.
    ///
    /// Initialization is done via K-means to place emission means, then EM
    /// iterates until convergence or `max_iterations` is reached.
    ///
    /// # Arguments
    ///
    /// * `observations` - Slice of f64 observations (e.g. returns, volatility)
    ///
    /// # Errors
    ///
    /// Returns error if observations are empty or contain fewer points than states.
    pub fn fit(&mut self, observations: &[f64]) -> Result<(), MLError> {
        let n = self.config.num_states;
        let t = observations.len();

        if t < n {
            return Err(MLError::ConfigError(format!(
                    "Need at least {} observations for {} states, got {}",
                    n, n, t
                )));
        }

        // -- K-means initialization for emission parameters --
        self.kmeans_init(observations);

        let mut prev_log_likelihood = f64::NEG_INFINITY;

        for _iter in 0..self.config.max_iterations {
            // -- E-step: forward-backward --
            let (alpha, scale_factors) = self.forward(observations);
            let beta = self.backward(observations, &scale_factors);

            // Compute log-likelihood from scaling factors
            let log_likelihood: f64 = scale_factors
                .iter()
                .filter(|&&s| s > 0.0)
                .map(|&s| s.ln())
                .sum();

            // Check convergence
            if (log_likelihood - prev_log_likelihood).abs() < self.config.convergence_threshold {
                break;
            }
            prev_log_likelihood = log_likelihood;

            // -- Compute gamma (state posteriors) and xi (transition posteriors) --
            let gamma = self.compute_gamma(&alpha, &beta, t, n);
            let xi = self.compute_xi(&alpha, &beta, observations, &scale_factors, t, n);

            // -- M-step: re-estimate parameters --
            self.m_step(observations, &gamma, &xi, t, n);
        }

        self.fitted = true;
        Ok(())
    }

    /// Predict the most likely hidden state for a single observation.
    ///
    /// Uses emission likelihoods weighted by steady-state probabilities
    /// (or initial probs if not converged to steady state).
    ///
    /// # Arguments
    ///
    /// * `observation` - A single observation value
    ///
    /// # Returns
    ///
    /// State index (0..num_states-1)
    pub fn predict(&self, observation: f64) -> usize {
        let n = self.config.num_states;
        let mut best_state = 0;
        let mut best_score = f64::NEG_INFINITY;

        for i in 0..n {
            let emission = self.gaussian_pdf(observation, i);
            let score = self.initial_probs.get(i).copied().unwrap_or(0.0) * emission;
            if score > best_score {
                best_score = score;
                best_state = i;
            }
        }

        best_state
    }

    /// Get a reference to the transition matrix.
    ///
    /// `transition_matrix[i][j]` = P(state j at t+1 | state i at t)
    pub fn transition_matrix(&self) -> &[Vec<f64>] {
        &self.transition_matrix
    }

    /// Get emission means for each state.
    pub fn emission_means(&self) -> &[f64] {
        &self.emission_means
    }

    /// Get emission variances for each state.
    pub fn emission_vars(&self) -> &[f64] {
        &self.emission_vars
    }

    /// Get initial state probabilities.
    pub fn initial_probs(&self) -> &[f64] {
        &self.initial_probs
    }

    /// Whether the model has been fitted.
    pub fn is_fitted(&self) -> bool {
        self.fitted
    }

    // ---- Private helpers ----

    /// Initialize emission means via K-means clustering.
    fn kmeans_init(&mut self, observations: &[f64]) {
        let n = self.config.num_states;
        let t = observations.len();

        // Sort observations and pick evenly spaced quantile centers
        let mut sorted: Vec<f64> = observations.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        for i in 0..n {
            let idx = (i * t) / n + t / (2 * n);
            let clamped = idx.min(t.saturating_sub(1));
            self.emission_means[i] = sorted.get(clamped).copied().unwrap_or(0.0);
        }

        // Run a few K-means iterations
        let mut assignments = vec![0_usize; t];
        for _km_iter in 0..20 {
            // Assign each observation to nearest center
            for (obs_idx, &obs) in observations.iter().enumerate() {
                let mut best = 0;
                let mut best_dist = f64::INFINITY;
                for (state_idx, &mean) in self.emission_means.iter().enumerate() {
                    let d = (obs - mean).abs();
                    if d < best_dist {
                        best_dist = d;
                        best = state_idx;
                    }
                }
                assignments[obs_idx] = best;
            }

            // Recompute centers
            let mut sums = vec![0.0; n];
            let mut counts = vec![0_usize; n];
            for (obs_idx, &obs) in observations.iter().enumerate() {
                let s = assignments.get(obs_idx).copied().unwrap_or(0);
                sums[s] += obs;
                counts[s] += 1;
            }
            for i in 0..n {
                if counts.get(i).copied().unwrap_or(0) > 0 {
                    self.emission_means[i] = sums.get(i).copied().unwrap_or(0.0)
                        / counts.get(i).copied().unwrap_or(1) as f64;
                }
            }
        }

        // Compute per-cluster variance
        let mut var_sums = vec![0.0; n];
        let mut var_counts = vec![0_usize; n];
        for (obs_idx, &obs) in observations.iter().enumerate() {
            let s = assignments.get(obs_idx).copied().unwrap_or(0);
            let mean = self.emission_means.get(s).copied().unwrap_or(0.0);
            var_sums[s] += (obs - mean).powi(2);
            var_counts[s] += 1;
        }
        for i in 0..n {
            let count = var_counts.get(i).copied().unwrap_or(0);
            if count > 1 {
                self.emission_vars[i] =
                    (var_sums.get(i).copied().unwrap_or(1.0) / count as f64).max(1e-8);
            } else {
                // Single or no points -- use global variance as fallback
                let global_mean: f64 =
                    observations.iter().sum::<f64>() / observations.len() as f64;
                let global_var: f64 = observations
                    .iter()
                    .map(|x| (x - global_mean).powi(2))
                    .sum::<f64>()
                    / observations.len() as f64;
                self.emission_vars[i] = global_var.max(1e-8);
            }
        }

    }

    /// Gaussian PDF for state `state` evaluated at `x`.
    fn gaussian_pdf(&self, x: f64, state: usize) -> f64 {
        let mean = self.emission_means.get(state).copied().unwrap_or(0.0);
        let var = self.emission_vars.get(state).copied().unwrap_or(1.0).max(1e-12);
        let diff = x - mean;
        let exponent = -0.5 * diff * diff / var;
        let norm = (2.0 * std::f64::consts::PI * var).sqrt();
        exponent.exp() / norm
    }

    /// Forward pass with scaling. Returns (alpha, scale_factors).
    fn forward(&self, observations: &[f64]) -> (Vec<Vec<f64>>, Vec<f64>) {
        let n = self.config.num_states;
        let t = observations.len();

        let mut alpha = vec![vec![0.0; n]; t];
        let mut scale_factors = vec![0.0; t];

        // Initialization
        let first_obs = observations.first().copied().unwrap_or(0.0);
        for i in 0..n {
            let init_p = self.initial_probs.get(i).copied().unwrap_or(0.0);
            alpha[0][i] = init_p * self.gaussian_pdf(first_obs, i);
        }
        let sum: f64 = alpha[0].iter().sum();
        scale_factors[0] = sum.max(f64::MIN_POSITIVE);
        for i in 0..n {
            alpha[0][i] /= scale_factors[0];
        }

        // Recursion
        for step in 1..t {
            let obs = observations.get(step).copied().unwrap_or(0.0);
            for j in 0..n {
                let mut s = 0.0;
                for i in 0..n {
                    let a_ij = self
                        .transition_matrix
                        .get(i)
                        .and_then(|row| row.get(j))
                        .copied()
                        .unwrap_or(0.0);
                    let prev = alpha
                        .get(step - 1)
                        .and_then(|row| row.get(i))
                        .copied()
                        .unwrap_or(0.0);
                    s += prev * a_ij;
                }
                alpha[step][j] = s * self.gaussian_pdf(obs, j);
            }
            let sum: f64 = alpha[step].iter().sum();
            scale_factors[step] = sum.max(f64::MIN_POSITIVE);
            for j in 0..n {
                alpha[step][j] /= scale_factors[step];
            }
        }

        (alpha, scale_factors)
    }

    /// Backward pass with scaling.
    fn backward(&self, observations: &[f64], scale_factors: &[f64]) -> Vec<Vec<f64>> {
        let n = self.config.num_states;
        let t = observations.len();

        let mut beta = vec![vec![0.0; n]; t];

        // Initialization: beta[T-1] = 1 / scale[T-1]
        let last_scale = scale_factors
            .last()
            .copied()
            .unwrap_or(1.0)
            .max(f64::MIN_POSITIVE);
        for i in 0..n {
            beta[t - 1][i] = 1.0 / last_scale;
        }

        // Recursion (backwards)
        for step in (0..t.saturating_sub(1)).rev() {
            let next_obs = observations.get(step + 1).copied().unwrap_or(0.0);
            for i in 0..n {
                let mut s = 0.0;
                for j in 0..n {
                    let a_ij = self
                        .transition_matrix
                        .get(i)
                        .and_then(|row| row.get(j))
                        .copied()
                        .unwrap_or(0.0);
                    let b_next = beta
                        .get(step + 1)
                        .and_then(|row| row.get(j))
                        .copied()
                        .unwrap_or(0.0);
                    s += a_ij * self.gaussian_pdf(next_obs, j) * b_next;
                }
                let scale = scale_factors
                    .get(step)
                    .copied()
                    .unwrap_or(1.0)
                    .max(f64::MIN_POSITIVE);
                beta[step][i] = s / scale;
            }
        }

        beta
    }

    /// Compute gamma (state posterior probabilities) for each time step.
    fn compute_gamma(
        &self,
        alpha: &[Vec<f64>],
        beta: &[Vec<f64>],
        t: usize,
        n: usize,
    ) -> Vec<Vec<f64>> {
        let mut gamma = vec![vec![0.0; n]; t];

        for (step, gamma_row) in gamma.iter_mut().enumerate().take(t) {
            let mut denom = 0.0;
            for (i, g_val) in gamma_row.iter_mut().enumerate().take(n) {
                let a = alpha.get(step).and_then(|r| r.get(i)).copied().unwrap_or(0.0);
                let b = beta.get(step).and_then(|r| r.get(i)).copied().unwrap_or(0.0);
                *g_val = a * b;
                denom += *g_val;
            }
            if denom > f64::MIN_POSITIVE {
                for val in gamma_row.iter_mut().take(n) {
                    *val /= denom;
                }
            }
        }

        gamma
    }

    /// Compute xi (transition posterior probabilities) for consecutive time steps.
    fn compute_xi(
        &self,
        alpha: &[Vec<f64>],
        beta: &[Vec<f64>],
        observations: &[f64],
        scale_factors: &[f64],
        t: usize,
        n: usize,
    ) -> Vec<Vec<Vec<f64>>> {
        let mut xi = vec![vec![vec![0.0; n]; n]; t.saturating_sub(1)];

        for (step, xi_row) in xi.iter_mut().enumerate().take(t.saturating_sub(1)) {
            let next_obs = observations.get(step + 1).copied().unwrap_or(0.0);
            let mut denom = 0.0;

            for (i, xi_i) in xi_row.iter_mut().enumerate().take(n) {
                for (j, xi_ij) in xi_i.iter_mut().enumerate().take(n) {
                    let a_prev = alpha.get(step).and_then(|r| r.get(i)).copied().unwrap_or(0.0);
                    let a_ij = self
                        .transition_matrix
                        .get(i)
                        .and_then(|row| row.get(j))
                        .copied()
                        .unwrap_or(0.0);
                    let b_next = beta
                        .get(step + 1)
                        .and_then(|r| r.get(j))
                        .copied()
                        .unwrap_or(0.0);
                    let emission = self.gaussian_pdf(next_obs, j);

                    *xi_ij = a_prev * a_ij * emission * b_next;
                    denom += *xi_ij;
                }
            }

            let _ = scale_factors; // used implicitly through scaled alpha/beta

            if denom > f64::MIN_POSITIVE {
                for row in xi_row.iter_mut().take(n) {
                    for val in row.iter_mut().take(n) {
                        *val /= denom;
                    }
                }
            }
        }

        xi
    }

    /// M-step: re-estimate HMM parameters from gamma and xi posteriors.
    fn m_step(
        &mut self,
        observations: &[f64],
        gamma: &[Vec<f64>],
        xi: &[Vec<Vec<f64>>],
        t: usize,
        n: usize,
    ) {
        // Re-estimate initial probs
        for i in 0..n {
            self.initial_probs[i] = gamma
                .first()
                .and_then(|r| r.get(i))
                .copied()
                .unwrap_or(1.0 / n as f64);
        }

        // Re-estimate transition matrix
        for i in 0..n {
            let gamma_sum: f64 = (0..t.saturating_sub(1))
                .map(|step| {
                    gamma
                        .get(step)
                        .and_then(|r| r.get(i))
                        .copied()
                        .unwrap_or(0.0)
                })
                .sum();

            for j in 0..n {
                let xi_sum: f64 = (0..t.saturating_sub(1))
                    .map(|step| {
                        xi.get(step)
                            .and_then(|r| r.get(i))
                            .and_then(|r| r.get(j))
                            .copied()
                            .unwrap_or(0.0)
                    })
                    .sum();

                self.transition_matrix[i][j] = if gamma_sum > f64::MIN_POSITIVE {
                    (xi_sum / gamma_sum).max(1e-10)
                } else {
                    1.0 / n as f64
                };
            }

            // Normalize row
            let row_sum: f64 = self.transition_matrix[i].iter().sum();
            if row_sum > f64::MIN_POSITIVE {
                for j in 0..n {
                    self.transition_matrix[i][j] /= row_sum;
                }
            }
        }

        // Re-estimate emission means
        for i in 0..n {
            let gamma_sum: f64 = (0..t)
                .map(|step| {
                    gamma
                        .get(step)
                        .and_then(|r| r.get(i))
                        .copied()
                        .unwrap_or(0.0)
                })
                .sum();

            if gamma_sum > f64::MIN_POSITIVE {
                let weighted_sum: f64 = (0..t)
                    .map(|step| {
                        let g = gamma
                            .get(step)
                            .and_then(|r| r.get(i))
                            .copied()
                            .unwrap_or(0.0);
                        let obs = observations.get(step).copied().unwrap_or(0.0);
                        g * obs
                    })
                    .sum();

                self.emission_means[i] = weighted_sum / gamma_sum;

                // Re-estimate emission variance
                let weighted_var: f64 = (0..t)
                    .map(|step| {
                        let g = gamma
                            .get(step)
                            .and_then(|r| r.get(i))
                            .copied()
                            .unwrap_or(0.0);
                        let obs = observations.get(step).copied().unwrap_or(0.0);
                        let diff = obs - self.emission_means[i];
                        g * diff * diff
                    })
                    .sum();

                self.emission_vars[i] = (weighted_var / gamma_sum).max(1e-8);
            }
        }
    }
}

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

    /// Generate synthetic data from 3 regimes: low-vol, medium-vol, high-vol.
    fn synthetic_regime_data() -> Vec<f64> {
        let mut data = Vec::with_capacity(300);
        // Regime 0: low volatility, mean ~ 0.0
        for i in 0..100 {
            data.push(0.001 * (i as f64 % 10.0 - 5.0));
        }
        // Regime 1: medium volatility, mean ~ 0.5
        for i in 0..100 {
            data.push(0.5 + 0.05 * (i as f64 % 10.0 - 5.0));
        }
        // Regime 2: high volatility, mean ~ -0.5
        for i in 0..100 {
            data.push(-0.5 + 0.2 * (i as f64 % 10.0 - 5.0));
        }
        data
    }

    #[test]
    fn test_fit_on_synthetic_data() {
        let data = synthetic_regime_data();
        let mut hmm = HiddenMarkovModel::new(HMMConfig::default());

        let result = hmm.fit(&data);
        assert!(result.is_ok(), "HMM fit should succeed on synthetic data");
        assert!(hmm.is_fitted(), "Model should be marked as fitted");

        // Emission means should be separated (3 distinct clusters)
        let means = hmm.emission_means();
        assert_eq!(means.len(), 3);
        // Means should not all be equal
        let all_equal = means.windows(2).all(|w| (w[0] - w[1]).abs() < 0.01);
        assert!(!all_equal, "Emission means should be distinct after fitting");
    }

    #[test]
    fn test_predict_returns_valid_state() {
        let data = synthetic_regime_data();
        let mut hmm = HiddenMarkovModel::new(HMMConfig::default());
        hmm.fit(&data).ok();

        // Predict on values from each regime
        let state_low = hmm.predict(0.0);
        let state_high = hmm.predict(-0.5);

        assert!(state_low < 3, "Predicted state should be 0, 1, or 2");
        assert!(state_high < 3, "Predicted state should be 0, 1, or 2");
    }

    #[test]
    fn test_transition_matrix_rows_sum_to_one() {
        let data = synthetic_regime_data();
        let mut hmm = HiddenMarkovModel::new(HMMConfig::default());
        hmm.fit(&data).ok();

        for (i, row) in hmm.transition_matrix().iter().enumerate() {
            let sum: f64 = row.iter().sum();
            assert!(
                (sum - 1.0).abs() < 1e-6,
                "Row {} sums to {} instead of 1.0",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_fit_too_few_observations() {
        let mut hmm = HiddenMarkovModel::new(HMMConfig::default());
        let result = hmm.fit(&[1.0, 2.0]);
        assert!(result.is_err(), "Should fail with fewer observations than states");
    }

    #[test]
    fn test_predict_unfitted_model() {
        let hmm = HiddenMarkovModel::new(HMMConfig::default());
        // Should still return a valid state (uses initial uniform priors)
        let state = hmm.predict(0.5);
        assert!(state < 3, "Unfitted model should still return valid state index");
    }

    #[test]
    fn test_initial_probs_sum_to_one() {
        let data = synthetic_regime_data();
        let mut hmm = HiddenMarkovModel::new(HMMConfig::default());
        hmm.fit(&data).ok();

        let sum: f64 = hmm.initial_probs().iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-6,
            "Initial probs should sum to 1.0, got {}",
            sum
        );
    }
}
