"""
Hidden Markov Model implementation
"""

import warnings
from typing import Optional
import numpy as np

# Try to import Rust backend
try:
    from optimizr._core import fit_hmm as _rust_fit_hmm, viterbi_decode as _rust_viterbi
    RUST_AVAILABLE = True
except ImportError:
    RUST_AVAILABLE = False


class HMM:
    """
    Hidden Markov Model for regime detection and sequence analysis.
    
    Uses the Baum-Welch algorithm (EM) for parameter estimation and
    Viterbi algorithm for finding the most likely state sequence.
    
    Parameters
    ----------
    n_states : int
        Number of hidden states
    
    Attributes
    ----------
    transition_matrix_ : np.ndarray or None
        Learned transition probabilities (n_states × n_states)
    emission_means_ : np.ndarray or None
        Mean of Gaussian emission for each state
    emission_stds_ : np.ndarray or None
        Std dev of Gaussian emission for each state
    
    Examples
    --------
    >>> import numpy as np
    >>> from optimizr import HMM
    >>> 
    >>> # Generate data with regime changes
    >>> returns = np.concatenate([
    ...     np.random.normal(0.01, 0.02, 500),  # Bull market
    ...     np.random.normal(-0.01, 0.03, 500),  # Bear market
    ... ])
    >>> 
    >>> # Fit HMM
    >>> hmm = HMM(n_states=2)
    >>> hmm.fit(returns, n_iterations=100)
    >>> 
    >>> # Decode states
    >>> states = hmm.predict(returns)
    >>> print(f"Detected states: {np.unique(states)}")
    """
    
    def __init__(self, n_states: int = 2):
        if n_states < 2:
            raise ValueError("n_states must be at least 2")
        
        self.n_states = n_states
        self.transition_matrix_: np.ndarray = np.zeros((n_states, n_states))
        self.emission_means_: np.ndarray = np.zeros(n_states)
        self.emission_stds_: np.ndarray = np.ones(n_states)
        self._params = None
    
    def fit(self, X: np.ndarray, n_iterations: int = 100, tolerance: float = 1e-6) -> 'HMM':
        """
        Fit HMM parameters using Baum-Welch algorithm.
        
        Parameters
        ----------
        X : np.ndarray
            Time series observations (1D array)
        n_iterations : int, default=100
            Maximum number of EM iterations
        tolerance : float, default=1e-6
            Convergence threshold for log-likelihood change
            
        Returns
        -------
        self : HMM
            Fitted model
        """
        X = np.asarray(X).flatten()
        
        if len(X) == 0:
            raise ValueError("X cannot be empty")
        
        if RUST_AVAILABLE:
            # Use Rust implementation
            self._params = _rust_fit_hmm(
                observations=X.tolist(),
                n_states=self.n_states,
                n_iterations=n_iterations,
                tolerance=tolerance
            )
            
            self.transition_matrix_ = np.array(self._params.transition_matrix)
            self.emission_means_ = np.array(self._params.emission_means)
            self.emission_stds_ = np.array(self._params.emission_stds)
        else:
            # Pure Python fallback
            warnings.warn(
                "Rust backend not available. Using slower Python implementation.",
                RuntimeWarning
            )
            self._fit_python(X, n_iterations, tolerance)
        
        return self
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict most likely state sequence using Viterbi algorithm.
        
        Parameters
        ----------
        X : np.ndarray
            Time series observations (1D array)
            
        Returns
        -------
        states : np.ndarray
            Most likely state at each time step
        """
        if self.transition_matrix_ is None:
            raise ValueError("Model must be fitted before prediction")
        
        X = np.asarray(X).flatten()
        
        if RUST_AVAILABLE and self._params is not None:
            states = _rust_viterbi(X.tolist(), self._params)
            return np.array(states)
        else:
            return self._viterbi_python(X)
    
    def score(self, X: np.ndarray) -> float:
        """
        Compute log-likelihood of observations.
        
        Parameters
        ----------
        X : np.ndarray
            Time series observations
            
        Returns
        -------
        log_likelihood : float
            Log P(X | model)
        """
        if self.transition_matrix_ is None:
            raise ValueError("Model must be fitted before scoring")
        
        X = np.asarray(X).flatten()
        alpha = self._forward_python(X)
        return np.log(np.sum(alpha[-1]))
    
    def _fit_python(self, X: np.ndarray, n_iterations: int, tolerance: float):
        """Pure Python implementation of Baum-Welch"""
        n_obs = len(X)
        
        # Initialize parameters
        self.transition_matrix_ = np.ones((self.n_states, self.n_states)) / self.n_states
        
        # Initialize emissions based on quantiles
        quantiles = np.linspace(0, 1, self.n_states + 1)
        self.emission_means_ = np.zeros(self.n_states)
        self.emission_stds_ = np.ones(self.n_states)
        
        for i in range(self.n_states):
            mask = (X >= np.quantile(X, quantiles[i])) & (X < np.quantile(X, quantiles[i+1]))
            if np.any(mask):
                self.emission_means_[i] = np.mean(X[mask])
                self.emission_stds_[i] = max(np.std(X[mask]), 1e-6)
        
        prev_ll = float('-inf')
        
        # EM iterations
        for _ in range(n_iterations):
            # E-step
            alpha = self._forward_python(X)
            beta = self._backward_python(X)
            gamma = self._compute_gamma_python(alpha, beta)
            xi = self._compute_xi_python(X, alpha, beta)
            
            # M-step
            self._update_parameters_python(X, gamma, xi)
            
            # Check convergence
            ll = np.log(np.sum(alpha[-1]))
            if abs(ll - prev_ll) < tolerance:
                break
            prev_ll = ll
    
    def _forward_python(self, X: np.ndarray) -> np.ndarray:
        """Forward algorithm"""
        n_obs = len(X)
        alpha = np.zeros((n_obs, self.n_states))
        
        # Initialize
        for s in range(self.n_states):
            alpha[0, s] = (1.0 / self.n_states) * self._emission_prob(X[0], s)
        
        alpha[0] /= np.sum(alpha[0])
        
        # Recursion
        for t in range(1, n_obs):
            for s in range(self.n_states):
                alpha[t, s] = np.sum(alpha[t-1] * self.transition_matrix_[:, s]) * self._emission_prob(X[t], s)
            alpha[t] /= max(np.sum(alpha[t]), 1e-10)
        
        return alpha
    
    def _backward_python(self, X: np.ndarray) -> np.ndarray:
        """Backward algorithm"""
        n_obs = len(X)
        beta = np.zeros((n_obs, self.n_states))
        beta[-1] = 1.0
        
        for t in range(n_obs - 2, -1, -1):
            for s in range(self.n_states):
                beta[t, s] = np.sum(
                    self.transition_matrix_[s] * 
                    np.array([self._emission_prob(X[t+1], s2) for s2 in range(self.n_states)]) *
                    beta[t+1]
                )
            beta[t] /= max(np.sum(beta[t]), 1e-10)
        
        return beta
    
    def _compute_gamma_python(self, alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
        """Compute state occupation probabilities"""
        gamma = alpha * beta
        gamma /= np.sum(gamma, axis=1, keepdims=True)
        return gamma
    
    def _compute_xi_python(self, X: np.ndarray, alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
        """Compute transition probabilities"""
        n_obs = len(X)
        xi = np.zeros((n_obs - 1, self.n_states, self.n_states))
        
        for t in range(n_obs - 1):
            for i in range(self.n_states):
                for j in range(self.n_states):
                    xi[t, i, j] = (alpha[t, i] * self.transition_matrix_[i, j] *
                                  self._emission_prob(X[t+1], j) * beta[t+1, j])
            xi[t] /= max(np.sum(xi[t]), 1e-10)
        
        return xi
    
    def _update_parameters_python(self, X: np.ndarray, gamma: np.ndarray, xi: np.ndarray):
        """M-step: update parameters"""
        # Update transitions
        for i in range(self.n_states):
            denom = np.sum(gamma[:-1, i])
            for j in range(self.n_states):
                numer = np.sum(xi[:, i, j])
                self.transition_matrix_[i, j] = numer / max(denom, 1e-10)
        
        # Update emissions
        for s in range(self.n_states):
            weights = gamma[:, s]
            sum_weights = np.sum(weights)
            
            if sum_weights > 1e-10:
                self.emission_means_[s] = np.sum(weights * X) / sum_weights
                self.emission_stds_[s] = max(
                    np.sqrt(np.sum(weights * (X - self.emission_means_[s])**2) / sum_weights),
                    1e-6
                )
    
    def _emission_prob(self, obs: float, state: int) -> float:
        """Gaussian emission probability"""
        mean = self.emission_means_[state]
        std = self.emission_stds_[state]
        z = (obs - mean) / std
        return max(np.exp(-0.5 * z**2) / (std * np.sqrt(2 * np.pi)), 1e-10)
    
    def _viterbi_python(self, X: np.ndarray) -> np.ndarray:
        """Viterbi decoding"""
        n_obs = len(X)
        delta = np.full((n_obs, self.n_states), float('-inf'))
        psi = np.zeros((n_obs, self.n_states), dtype=int)
        
        # Initialize
        for s in range(self.n_states):
            delta[0, s] = np.log(1.0 / self.n_states) + np.log(self._emission_prob(X[0], s))
        
        # Recursion
        for t in range(1, n_obs):
            for s in range(self.n_states):
                trans_probs = delta[t-1] + np.log(self.transition_matrix_[:, s] + 1e-10)
                psi[t, s] = np.argmax(trans_probs)
                delta[t, s] = trans_probs[psi[t, s]] + np.log(self._emission_prob(X[t], s))
        
        # Backtrack
        path = np.zeros(n_obs, dtype=int)
        path[-1] = np.argmax(delta[-1])
        
        for t in range(n_obs - 2, -1, -1):
            path[t] = psi[t + 1, path[t + 1]]
        
        return path
