"""Regime Detection with Hidden Markov Model + Strategy Switching."""
import numpy as np
import pandas as pd
from sklearn.mixture import GaussianMixture
from typing import Dict, List, Optional
import warnings
warnings.filterwarnings('ignore')

try:
    from hmmlearn.hmm import GaussianHMM
    HMM_AVAILABLE = True
except ImportError:
    HMM_AVAILABLE = False


class RegimeDetectorHMM:
    """Market regime detection with strategy parameters per regime."""
    
    def __init__(self, n_regimes: int = 3):
        self.n_regimes = n_regimes
        self.model = None
        self.regime_names = {}
        self.regime_history = []
        
    def fit(self, returns: pd.Series, volatility: Optional[pd.Series] = None):
        features = pd.DataFrame({'returns': returns.fillna(0)})
        features['volatility'] = volatility.fillna(0) if volatility is not None else returns.fillna(0).rolling(21).std() * np.sqrt(252)
        features = features.dropna()
        
        if HMM_AVAILABLE and len(features) > 63:
            self.model = GaussianHMM(n_components=self.n_regimes, covariance_type='full', n_iter=100, random_state=42)
            self.model.fit(features.values)
            means = self.model.means_[:, 0]
            order = np.argsort(means)[::-1]
            self.regime_names = {order[0]: 'bull', order[1]: 'neutral', order[2]: 'bear'}
        else:
            self.model = None
            
    def predict(self, returns: pd.Series, volatility: Optional[pd.Series] = None) -> pd.Series:
        features = pd.DataFrame({'returns': returns.fillna(0)})
        features['volatility'] = volatility.fillna(0) if volatility is not None else returns.fillna(0).rolling(21).std() * np.sqrt(252)
        features = features.dropna()
        
        if self.model is not None and len(features) > 10:
            states = self.model.predict(features.values)
            regime_series = pd.Series([self.regime_names.get(s, 'neutral') for s in states], index=features.index)
        else:
            regime_series = pd.Series('neutral', index=features.index)
            for idx in features.index:
                r = features.loc[idx, 'returns']
                v = features.loc[idx, 'volatility']
                v_med = features['volatility'].median()
                if v > v_med * 1.5: regime_series.loc[idx] = 'high_vol'
                elif r > 0.001: regime_series.loc[idx] = 'bull'
                elif r < -0.001: regime_series.loc[idx] = 'bear'
                else: regime_series.loc[idx] = 'neutral'
        
        self.regime_history = regime_series
        return regime_series
    
    def get_regime_strategy(self, regime: str) -> Dict:
        strategies = {
            'bull': {'risk_aversion': 1.0, 'momentum_weight': 0.7, 'mean_reversion_weight': 0.1, 'max_leverage': 1.5, 'hedge_ratio': 0.0},
            'bear': {'risk_aversion': 3.0, 'momentum_weight': 0.2, 'mean_reversion_weight': 0.5, 'max_leverage': 0.5, 'hedge_ratio': 0.5},
            'high_vol': {'risk_aversion': 4.0, 'momentum_weight': 0.3, 'mean_reversion_weight': 0.3, 'max_leverage': 0.3, 'hedge_ratio': 0.7},
            'neutral': {'risk_aversion': 2.0, 'momentum_weight': 0.4, 'mean_reversion_weight': 0.4, 'max_leverage': 1.0, 'hedge_ratio': 0.2}
        }
        return strategies.get(regime, strategies['neutral'])
    
    def get_regime_stats(self, returns: pd.Series) -> pd.DataFrame:
        if len(self.regime_history) == 0:
            return pd.DataFrame()
        stats = []
        for regime in self.regime_history.unique():
            mask = self.regime_history == regime
            r = returns.reindex(self.regime_history.index)[mask].dropna()
            if len(r) > 0:
                stats.append({'regime': regime, 'n_days': len(r), 'mean_return': r.mean() * 252,
                              'volatility': r.std() * np.sqrt(252), 'max_drawdown': (r.cumsum() - r.cumsum().cummax()).min()})
        return pd.DataFrame(stats)