"""
MRD HMM Regime Detection - Hidden Markov Model for Market Regimes
=================================================================

Implements a Hidden Markov Model for detecting market regimes
(Trend Up, Range, Trend Down) based on momentum features.

The HMM learns to identify distinct market states from
observable momentum indicators without explicit labels.

Key Features:
- 3-state Gaussian HMM for regime classification
- Probability-based regime assignment
- Regime duration tracking
- Regime transition analysis

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-25
Version: 1.0.0
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Union, Any
from dataclasses import dataclass, field
from pathlib import Path
import joblib
from loguru import logger
from enum import IntEnum

try:
    from hmmlearn.hmm import GaussianHMM
    HAS_HMMLEARN = True
except ImportError:
    HAS_HMMLEARN = False
    logger.warning("hmmlearn not available. Install with: pip install hmmlearn")


class MarketRegime(IntEnum):
    """Market regime states."""
    TREND_DOWN = 0
    RANGE = 1
    TREND_UP = 2


@dataclass
class RegimeHMMConfig:
    """Configuration for Regime HMM."""

    # Number of hidden states
    n_states: int = 3

    # HMM parameters
    n_iterations: int = 100
    convergence_threshold: float = 1e-4
    covariance_type: str = 'full'  # 'full', 'tied', 'diag', 'spherical'

    # Random state for reproducibility
    random_state: int = 42

    # Feature selection for HMM
    use_features: List[str] = field(default_factory=lambda: [
        'rsi_14_norm',
        'macd_histogram',
        'roc_10',
        'adx',
        'di_spread',
        'trend_distance',
        'trend_alignment',
        'bb_position_norm'
    ])

    # Regime classification thresholds
    regime_prob_threshold: float = 0.6


@dataclass
class RegimePrediction:
    """Prediction result from the HMM."""
    regime: int
    regime_name: str
    probabilities: Dict[str, float]
    confidence: float
    duration: int
    transition_prob_to_up: float
    transition_prob_to_down: float

    def to_dict(self) -> Dict:
        return {
            'regime': self.regime,
            'regime_name': self.regime_name,
            'probabilities': self.probabilities,
            'confidence': self.confidence,
            'duration': self.duration,
            'transition_prob_to_up': self.transition_prob_to_up,
            'transition_prob_to_down': self.transition_prob_to_down
        }


class RegimeHMM:
    """
    Hidden Markov Model for Market Regime Detection.

    Uses a Gaussian HMM with 3 hidden states to model market regimes:
    - State 0: Trend Down (bearish momentum)
    - State 1: Range (consolidation/low momentum)
    - State 2: Trend Up (bullish momentum)

    The model learns regime patterns from momentum features without
    requiring explicit regime labels.

    Usage:
        hmm = RegimeHMM()
        hmm.fit(features)

        # Predict current regime
        regime = hmm.predict_regime(features)

        # Get regime probabilities
        probs = hmm.get_regime_probabilities(features)
    """

    REGIME_NAMES = {
        0: 'trend_down',
        1: 'range',
        2: 'trend_up'
    }

    def __init__(self, config: Optional[RegimeHMMConfig] = None):
        """
        Initialize the Regime HMM.

        Args:
            config: HMM configuration
        """
        if not HAS_HMMLEARN:
            raise ImportError(
                "hmmlearn is required for RegimeHMM. "
                "Install with: pip install hmmlearn"
            )

        self.config = config or RegimeHMMConfig()
        self.model: Optional[GaussianHMM] = None
        self.feature_scaler_mean: Optional[np.ndarray] = None
        self.feature_scaler_std: Optional[np.ndarray] = None
        self._is_fitted = False
        self._regime_order: Optional[List[int]] = None

    def _init_model(self) -> GaussianHMM:
        """Initialize the HMM model."""
        return GaussianHMM(
            n_components=self.config.n_states,
            covariance_type=self.config.covariance_type,
            n_iter=self.config.n_iterations,
            tol=self.config.convergence_threshold,
            random_state=self.config.random_state,
            init_params='stmc',  # Initialize all parameters
            params='stmc'  # Train all parameters
        )

    def _preprocess_features(
        self,
        features: pd.DataFrame,
        fit_scaler: bool = False
    ) -> np.ndarray:
        """
        Preprocess features for HMM.

        Args:
            features: DataFrame with feature columns
            fit_scaler: Whether to fit the scaler (True for training)

        Returns:
            Scaled feature array
        """
        # Select relevant features
        use_cols = [c for c in self.config.use_features if c in features.columns]
        if not use_cols:
            use_cols = features.columns.tolist()
            logger.warning(f"No configured features found, using all: {len(use_cols)}")

        X = features[use_cols].values.astype(np.float64)

        # Handle NaN values
        X = np.nan_to_num(X, nan=0.0)

        # Scale features
        if fit_scaler:
            self.feature_scaler_mean = np.mean(X, axis=0)
            self.feature_scaler_std = np.std(X, axis=0) + 1e-8
            self._feature_columns = use_cols

        if self.feature_scaler_mean is not None:
            X = (X - self.feature_scaler_mean) / self.feature_scaler_std

        # Clip extreme values
        X = np.clip(X, -5, 5)

        return X

    def fit(
        self,
        features: pd.DataFrame,
        lengths: Optional[List[int]] = None
    ) -> 'RegimeHMM':
        """
        Fit the HMM to feature data.

        Args:
            features: DataFrame with momentum features
            lengths: Optional sequence lengths for multiple sequences

        Returns:
            Self (fitted model)
        """
        logger.info(f"Fitting RegimeHMM on {len(features)} samples")

        # Preprocess features
        X = self._preprocess_features(features, fit_scaler=True)

        # Initialize model
        self.model = self._init_model()

        # Fit the model
        if lengths is not None:
            self.model.fit(X, lengths)
        else:
            self.model.fit(X)

        self._is_fitted = True

        # Order regimes by mean ROC/momentum (so state 0 = bearish, state 2 = bullish)
        self._order_regimes(X)

        # Log model statistics
        self._log_model_stats()

        logger.info("RegimeHMM fitted successfully")

        return self

    def _order_regimes(self, X: np.ndarray):
        """
        Order regime states so that state 0=bearish, 1=range, 2=bullish.

        Uses the mean of momentum-related features in each state
        to determine the ordering.
        """
        # Predict states for all data
        states = self.model.predict(X)

        # Find momentum-related feature indices
        momentum_cols = ['roc_10', 'macd_histogram', 'di_spread', 'trend_distance']
        momentum_indices = [
            i for i, col in enumerate(self._feature_columns)
            if any(m in col for m in momentum_cols)
        ]

        if not momentum_indices:
            momentum_indices = list(range(min(3, X.shape[1])))

        # Calculate mean momentum for each state
        state_momentum = {}
        for state in range(self.config.n_states):
            mask = states == state
            if mask.sum() > 0:
                state_momentum[state] = np.mean(X[mask][:, momentum_indices])
            else:
                state_momentum[state] = 0

        # Order states by momentum (lowest = bearish, highest = bullish)
        ordered_states = sorted(state_momentum.keys(), key=lambda s: state_momentum[s])
        self._regime_order = ordered_states

        logger.info(f"Regime ordering by momentum: {self._regime_order}")
        logger.info(f"State momentums: {state_momentum}")

    def _map_state(self, state: int) -> int:
        """Map internal HMM state to ordered regime."""
        if self._regime_order is None:
            return state
        return self._regime_order.index(state)

    def _unmap_state(self, regime: int) -> int:
        """Map ordered regime back to internal HMM state."""
        if self._regime_order is None:
            return regime
        return self._regime_order[regime]

    def _log_model_stats(self):
        """Log model statistics after fitting."""
        if self.model is None:
            return

        logger.info("HMM Model Statistics:")
        logger.info(f"  States: {self.config.n_states}")
        logger.info(f"  Converged: {self.model.monitor_.converged}")
        logger.info(f"  Iterations: {len(self.model.monitor_.history)}")

        # Starting probabilities (mapped to ordered regimes)
        start_probs = {}
        for i, name in self.REGIME_NAMES.items():
            internal_state = self._unmap_state(i)
            start_probs[name] = self.model.startprob_[internal_state]
        logger.info(f"  Starting probabilities: {start_probs}")

        # Transition matrix (reordered)
        logger.info("  Transition matrix (rows=from, cols=to):")
        ordered_trans = self._get_ordered_transition_matrix()
        for i, row in enumerate(ordered_trans):
            logger.info(f"    {self.REGIME_NAMES[i]}: {row.round(3)}")

    def _get_ordered_transition_matrix(self) -> np.ndarray:
        """Get transition matrix with reordered states."""
        if self._regime_order is None:
            return self.model.transmat_

        n = self.config.n_states
        ordered_trans = np.zeros((n, n))

        for i in range(n):
            for j in range(n):
                from_internal = self._unmap_state(i)
                to_internal = self._unmap_state(j)
                ordered_trans[i, j] = self.model.transmat_[from_internal, to_internal]

        return ordered_trans

    def predict_regime(
        self,
        features: pd.DataFrame
    ) -> np.ndarray:
        """
        Predict the market regime for each sample.

        Args:
            features: DataFrame with momentum features

        Returns:
            Array of regime labels (0=trend_down, 1=range, 2=trend_up)
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")

        X = self._preprocess_features(features, fit_scaler=False)
        internal_states = self.model.predict(X)

        # Map to ordered regimes
        regimes = np.array([self._map_state(s) for s in internal_states])

        return regimes

    def get_regime_probabilities(
        self,
        features: pd.DataFrame
    ) -> np.ndarray:
        """
        Get probability distribution over regimes for each sample.

        Args:
            features: DataFrame with momentum features

        Returns:
            Array of shape (n_samples, n_states) with probabilities
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")

        X = self._preprocess_features(features, fit_scaler=False)

        # Get posterior probabilities
        _, posteriors = self.model.score_samples(X)

        # Reorder columns to match regime order
        ordered_posteriors = np.zeros_like(posteriors)
        for i in range(self.config.n_states):
            internal_state = self._unmap_state(i)
            ordered_posteriors[:, i] = posteriors[:, internal_state]

        return ordered_posteriors

    def get_regime_duration(
        self,
        states: np.ndarray
    ) -> np.ndarray:
        """
        Calculate the duration of the current regime at each point.

        Args:
            states: Array of regime states

        Returns:
            Array of durations (how many bars in current regime)
        """
        durations = np.ones_like(states)
        current_duration = 1

        for i in range(1, len(states)):
            if states[i] == states[i - 1]:
                current_duration += 1
            else:
                current_duration = 1
            durations[i] = current_duration

        return durations

    def predict_with_details(
        self,
        features: pd.DataFrame
    ) -> List[RegimePrediction]:
        """
        Get detailed regime predictions including probabilities and durations.

        Args:
            features: DataFrame with momentum features

        Returns:
            List of RegimePrediction objects
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")

        # Get regimes and probabilities
        regimes = self.predict_regime(features)
        probabilities = self.get_regime_probabilities(features)
        durations = self.get_regime_duration(regimes)

        # Get transition matrix for transition probabilities
        trans_matrix = self._get_ordered_transition_matrix()

        predictions = []
        for i in range(len(regimes)):
            regime = int(regimes[i])
            regime_name = self.REGIME_NAMES[regime]

            probs_dict = {
                self.REGIME_NAMES[j]: float(probabilities[i, j])
                for j in range(self.config.n_states)
            }

            confidence = float(probabilities[i, regime])

            pred = RegimePrediction(
                regime=regime,
                regime_name=regime_name,
                probabilities=probs_dict,
                confidence=confidence,
                duration=int(durations[i]),
                transition_prob_to_up=float(trans_matrix[regime, MarketRegime.TREND_UP]),
                transition_prob_to_down=float(trans_matrix[regime, MarketRegime.TREND_DOWN])
            )
            predictions.append(pred)

        return predictions

    def get_regime_statistics(
        self,
        features: pd.DataFrame
    ) -> Dict[str, Any]:
        """
        Get statistics about regimes in the data.

        Args:
            features: DataFrame with momentum features

        Returns:
            Dictionary with regime statistics
        """
        regimes = self.predict_regime(features)
        durations = self.get_regime_duration(regimes)

        stats = {
            'regime_counts': {},
            'regime_percentages': {},
            'avg_duration': {},
            'max_duration': {},
            'transition_matrix': self._get_ordered_transition_matrix().tolist()
        }

        for regime_id, regime_name in self.REGIME_NAMES.items():
            mask = regimes == regime_id
            count = mask.sum()
            stats['regime_counts'][regime_name] = int(count)
            stats['regime_percentages'][regime_name] = float(count / len(regimes) * 100)

            if count > 0:
                stats['avg_duration'][regime_name] = float(durations[mask].mean())
                stats['max_duration'][regime_name] = int(durations[mask].max())
            else:
                stats['avg_duration'][regime_name] = 0
                stats['max_duration'][regime_name] = 0

        return stats

    def save(self, path: str):
        """Save the model to disk."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save HMM model
        joblib.dump(self.model, path / 'hmm_model.joblib')

        # Save metadata
        metadata = {
            'config': self.config.__dict__,
            'feature_scaler_mean': self.feature_scaler_mean,
            'feature_scaler_std': self.feature_scaler_std,
            'feature_columns': self._feature_columns,
            'regime_order': self._regime_order,
            'version': '1.0.0'
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"Saved RegimeHMM to {path}")

    def load(self, path: str):
        """Load the model from disk."""
        path = Path(path)

        self.model = joblib.load(path / 'hmm_model.joblib')

        metadata = joblib.load(path / 'metadata.joblib')
        self.config = RegimeHMMConfig(**metadata['config'])
        self.feature_scaler_mean = metadata['feature_scaler_mean']
        self.feature_scaler_std = metadata['feature_scaler_std']
        self._feature_columns = metadata['feature_columns']
        self._regime_order = metadata['regime_order']

        self._is_fitted = True
        logger.info(f"Loaded RegimeHMM from {path}")


if __name__ == "__main__":
    # Test HMM regime detection
    from feature_engineering import MRDFeatureEngineer

    np.random.seed(42)
    n_samples = 2000

    # Create sample OHLCV data with trend patterns
    dates = pd.date_range('2024-01-01', periods=n_samples, freq='5min')

    # Create price with regime changes
    price = np.zeros(n_samples)
    price[0] = 2000

    regime_changes = [0, 500, 1000, 1500]
    regimes_true = ['up', 'down', 'range', 'up']

    for i in range(1, n_samples):
        # Determine current regime
        regime_idx = sum(1 for r in regime_changes if i >= r) - 1
        regime = regimes_true[regime_idx]

        if regime == 'up':
            drift = 0.0002
        elif regime == 'down':
            drift = -0.0002
        else:
            drift = 0

        price[i] = price[i-1] * (1 + drift + np.random.randn() * 0.0005)

    df = pd.DataFrame({
        'open': price * (1 + np.random.randn(n_samples) * 0.0002),
        'high': price * (1 + np.abs(np.random.randn(n_samples)) * 0.0005),
        'low': price * (1 - np.abs(np.random.randn(n_samples)) * 0.0005),
        'close': price,
        'volume': np.random.randint(100, 10000, n_samples)
    }, index=dates)

    # Ensure OHLC consistency
    df['high'] = df[['open', 'high', 'close']].max(axis=1)
    df['low'] = df[['open', 'low', 'close']].min(axis=1)

    # Compute features
    engineer = MRDFeatureEngineer()
    features = engineer.compute_all_features(df)

    # Train HMM
    print(f"\n{'='*60}")
    print("Testing RegimeHMM")
    print(f"{'='*60}")

    hmm = RegimeHMM()
    hmm.fit(features)

    # Get predictions
    regimes = hmm.predict_regime(features)
    probs = hmm.get_regime_probabilities(features)
    durations = hmm.get_regime_duration(regimes)

    # Statistics
    stats = hmm.get_regime_statistics(features)
    print(f"\nRegime Statistics:")
    for regime, count in stats['regime_counts'].items():
        pct = stats['regime_percentages'][regime]
        avg_dur = stats['avg_duration'][regime]
        print(f"  {regime}: {count} samples ({pct:.1f}%), avg duration: {avg_dur:.1f}")

    # Detailed predictions for last few samples
    print(f"\nLast 5 predictions:")
    predictions = hmm.predict_with_details(features.tail(5))
    for i, pred in enumerate(predictions):
        print(f"  {i}: {pred.regime_name} (conf={pred.confidence:.2f}, dur={pred.duration})")

    print("\nTest complete!")
