"""
MRD Feature Engineering - Momentum Regime Detection Features
=============================================================

Feature engineering module for momentum-based indicators used in
regime detection. Computes technical indicators that capture
momentum, trend strength, and market conditions.

Features computed:
- RSI (Relative Strength Index) at multiple periods
- MACD (Moving Average Convergence Divergence)
- ROC (Rate of Change) at multiple periods
- ADX (Average Directional Index) with +DI/-DI
- EMA crossover signals
- Trend strength relative to EMA 200

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
from dataclasses import dataclass, field
from loguru import logger


@dataclass
class MRDFeatureConfig:
    """Configuration for MRD feature engineering."""

    # RSI periods
    rsi_periods: List[int] = field(default_factory=lambda: [14, 28])

    # MACD parameters
    macd_fast: int = 12
    macd_slow: int = 26
    macd_signal: int = 9

    # ROC periods
    roc_periods: List[int] = field(default_factory=lambda: [5, 10, 20])

    # ADX period
    adx_period: int = 14

    # EMA crossover pairs (short, long)
    ema_pairs: List[Tuple[int, int]] = field(
        default_factory=lambda: [(9, 21), (21, 50)]
    )

    # Trend EMA
    trend_ema_period: int = 200

    # Smoothing for stability
    smooth_period: int = 3


class MRDFeatureEngineer:
    """
    Feature engineer for Momentum Regime Detection.

    Computes technical indicators that capture market momentum
    and regime characteristics for HMM-based regime detection.

    Features are designed to be stationary and suitable for
    Hidden Markov Model training.

    Usage:
        engineer = MRDFeatureEngineer()
        features = engineer.compute_all_features(df)

        # Or compute individual features:
        rsi = engineer.compute_rsi(df)
        macd, signal, hist = engineer.compute_macd(df)
    """

    def __init__(self, config: Optional[MRDFeatureConfig] = None):
        """
        Initialize the feature engineer.

        Args:
            config: Feature configuration
        """
        self.config = config or MRDFeatureConfig()
        self._feature_names: List[str] = []

    def compute_all_features(
        self,
        df: pd.DataFrame,
        include_raw_price: bool = False
    ) -> pd.DataFrame:
        """
        Compute all MRD features from OHLCV data.

        Args:
            df: DataFrame with OHLCV columns (open, high, low, close, volume)
            include_raw_price: Whether to include raw price features

        Returns:
            DataFrame with computed features
        """
        logger.info(f"Computing MRD features for {len(df)} samples")

        features = pd.DataFrame(index=df.index)

        # 1. RSI features
        rsi_features = self.compute_rsi(df)
        features = pd.concat([features, rsi_features], axis=1)

        # 2. MACD features
        macd, signal, histogram = self.compute_macd(df)
        features['macd'] = macd
        features['macd_signal'] = signal
        features['macd_histogram'] = histogram
        features['macd_histogram_change'] = histogram.diff()

        # 3. ROC features
        roc_features = self.compute_roc(df)
        features = pd.concat([features, roc_features], axis=1)

        # 4. ADX features
        adx, plus_di, minus_di = self.compute_adx(df)
        features['adx'] = adx
        features['plus_di'] = plus_di
        features['minus_di'] = minus_di
        features['di_spread'] = plus_di - minus_di
        features['di_ratio'] = plus_di / (minus_di + 1e-8)

        # 5. EMA crossover features
        crossover_features = self.compute_ema_crossovers(df)
        features = pd.concat([features, crossover_features], axis=1)

        # 6. Trend strength
        trend_features = self.compute_trend_strength(df)
        features = pd.concat([features, trend_features], axis=1)

        # 7. Additional momentum features
        momentum_features = self._compute_additional_momentum(df)
        features = pd.concat([features, momentum_features], axis=1)

        # 8. Volatility-adjusted features
        vol_features = self._compute_volatility_features(df)
        features = pd.concat([features, vol_features], axis=1)

        # Store feature names
        self._feature_names = features.columns.tolist()

        # Fill NaN values with forward fill then zero
        features = features.ffill().fillna(0)

        logger.info(f"Computed {len(self._feature_names)} features")

        return features

    def compute_rsi(
        self,
        df: pd.DataFrame,
        periods: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Compute RSI (Relative Strength Index) at multiple periods.

        RSI measures the magnitude of recent price changes to evaluate
        overbought or oversold conditions.

        Args:
            df: DataFrame with 'close' column
            periods: List of periods for RSI calculation

        Returns:
            DataFrame with RSI columns
        """
        periods = periods or self.config.rsi_periods
        close = df['close']

        rsi_features = pd.DataFrame(index=df.index)

        for period in periods:
            delta = close.diff()

            gain = delta.where(delta > 0, 0)
            loss = (-delta).where(delta < 0, 0)

            # Use EMA for more responsive RSI
            avg_gain = gain.ewm(span=period, adjust=False).mean()
            avg_loss = loss.ewm(span=period, adjust=False).mean()

            rs = avg_gain / (avg_loss + 1e-10)
            rsi = 100 - (100 / (1 + rs))

            # Normalize RSI to [-1, 1] for HMM
            rsi_normalized = (rsi - 50) / 50

            rsi_features[f'rsi_{period}'] = rsi
            rsi_features[f'rsi_{period}_norm'] = rsi_normalized

            # RSI rate of change
            rsi_features[f'rsi_{period}_roc'] = rsi.diff(3) / 100

        return rsi_features

    def compute_macd(
        self,
        df: pd.DataFrame,
        fast: Optional[int] = None,
        slow: Optional[int] = None,
        signal: Optional[int] = None
    ) -> Tuple[pd.Series, pd.Series, pd.Series]:
        """
        Compute MACD (Moving Average Convergence Divergence).

        MACD shows the relationship between two EMAs of a security's price.

        Args:
            df: DataFrame with 'close' column
            fast: Fast EMA period
            slow: Slow EMA period
            signal: Signal line period

        Returns:
            Tuple of (macd, signal_line, histogram)
        """
        fast = fast or self.config.macd_fast
        slow = slow or self.config.macd_slow
        signal_period = signal or self.config.macd_signal

        close = df['close']

        # Calculate EMAs
        ema_fast = close.ewm(span=fast, adjust=False).mean()
        ema_slow = close.ewm(span=slow, adjust=False).mean()

        # MACD line
        macd_line = ema_fast - ema_slow

        # Signal line
        signal_line = macd_line.ewm(span=signal_period, adjust=False).mean()

        # Histogram
        histogram = macd_line - signal_line

        # Normalize by price for comparability across assets
        price_scale = close.rolling(slow).mean()
        macd_normalized = macd_line / (price_scale + 1e-8) * 100
        signal_normalized = signal_line / (price_scale + 1e-8) * 100
        histogram_normalized = histogram / (price_scale + 1e-8) * 100

        return macd_normalized, signal_normalized, histogram_normalized

    def compute_roc(
        self,
        df: pd.DataFrame,
        periods: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Compute ROC (Rate of Change) at multiple periods.

        ROC measures the percentage change in price over a given period.

        Args:
            df: DataFrame with 'close' column
            periods: List of periods for ROC calculation

        Returns:
            DataFrame with ROC columns
        """
        periods = periods or self.config.roc_periods
        close = df['close']

        roc_features = pd.DataFrame(index=df.index)

        for period in periods:
            # Standard ROC
            roc = (close - close.shift(period)) / (close.shift(period) + 1e-10) * 100
            roc_features[f'roc_{period}'] = roc

            # Smoothed ROC for stability
            roc_smooth = roc.rolling(self.config.smooth_period).mean()
            roc_features[f'roc_{period}_smooth'] = roc_smooth

            # ROC acceleration (momentum of momentum)
            roc_features[f'roc_{period}_accel'] = roc.diff()

        return roc_features

    def compute_adx(
        self,
        df: pd.DataFrame,
        period: Optional[int] = None
    ) -> Tuple[pd.Series, pd.Series, pd.Series]:
        """
        Compute ADX (Average Directional Index) with +DI and -DI.

        ADX measures trend strength regardless of direction.
        +DI and -DI indicate the direction of the trend.

        Args:
            df: DataFrame with 'high', 'low', 'close' columns
            period: ADX period

        Returns:
            Tuple of (adx, plus_di, minus_di)
        """
        period = period or self.config.adx_period

        high = df['high']
        low = df['low']
        close = df['close']

        # True Range
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

        # Directional Movement
        plus_dm = high.diff()
        minus_dm = -low.diff()

        # Only count if positive and greater than the other
        plus_dm = plus_dm.where(
            (plus_dm > minus_dm) & (plus_dm > 0), 0
        )
        minus_dm = minus_dm.where(
            (minus_dm > plus_dm) & (minus_dm > 0), 0
        )

        # Smoothed averages using Wilder's smoothing
        atr = true_range.ewm(span=period, adjust=False).mean()
        plus_di_smooth = plus_dm.ewm(span=period, adjust=False).mean()
        minus_di_smooth = minus_dm.ewm(span=period, adjust=False).mean()

        # Directional Indicators
        plus_di = 100 * plus_di_smooth / (atr + 1e-10)
        minus_di = 100 * minus_di_smooth / (atr + 1e-10)

        # DX and ADX
        di_diff = abs(plus_di - minus_di)
        di_sum = plus_di + minus_di
        dx = 100 * di_diff / (di_sum + 1e-10)

        adx = dx.ewm(span=period, adjust=False).mean()

        return adx, plus_di, minus_di

    def compute_ema_crossovers(
        self,
        df: pd.DataFrame,
        ema_pairs: Optional[List[Tuple[int, int]]] = None
    ) -> pd.DataFrame:
        """
        Compute EMA crossover signals and distances.

        Generates signals for EMA crosses (9/21, 21/50) and measures
        the distance between EMAs as a trend indicator.

        Args:
            df: DataFrame with 'close' column
            ema_pairs: List of (short_period, long_period) tuples

        Returns:
            DataFrame with crossover features
        """
        ema_pairs = ema_pairs or self.config.ema_pairs
        close = df['close']

        crossover_features = pd.DataFrame(index=df.index)

        for short_period, long_period in ema_pairs:
            ema_short = close.ewm(span=short_period, adjust=False).mean()
            ema_long = close.ewm(span=long_period, adjust=False).mean()

            # Distance between EMAs (normalized by price)
            ema_distance = (ema_short - ema_long) / (close + 1e-10) * 100
            crossover_features[f'ema_{short_period}_{long_period}_dist'] = ema_distance

            # Crossover signal (1 = bullish cross, -1 = bearish cross)
            cross_up = (ema_short > ema_long) & (ema_short.shift(1) <= ema_long.shift(1))
            cross_down = (ema_short < ema_long) & (ema_short.shift(1) >= ema_long.shift(1))
            crossover_signal = cross_up.astype(int) - cross_down.astype(int)
            crossover_features[f'ema_{short_period}_{long_period}_cross'] = crossover_signal

            # Cumulative crossover state (1 = above, -1 = below)
            state = (ema_short > ema_long).astype(int) * 2 - 1
            crossover_features[f'ema_{short_period}_{long_period}_state'] = state

            # Time since last crossover (capped)
            cross_any = cross_up | cross_down
            crossover_features[f'ema_{short_period}_{long_period}_bars_since'] = (
                cross_any.groupby((cross_any != cross_any.shift()).cumsum()).cumcount()
            ).clip(upper=100)

            # Rate of change of distance
            crossover_features[f'ema_{short_period}_{long_period}_dist_roc'] = (
                ema_distance.diff(3)
            )

        return crossover_features

    def compute_trend_strength(
        self,
        df: pd.DataFrame,
        ema_period: Optional[int] = None
    ) -> pd.DataFrame:
        """
        Compute trend strength relative to EMA 200.

        Measures how far price is from long-term trend and
        the consistency of the trend direction.

        Args:
            df: DataFrame with 'close' column
            ema_period: EMA period for trend reference (default 200)

        Returns:
            DataFrame with trend strength features
        """
        ema_period = ema_period or self.config.trend_ema_period
        close = df['close']

        trend_features = pd.DataFrame(index=df.index)

        # EMA 200
        ema_200 = close.ewm(span=ema_period, adjust=False).mean()

        # Distance from EMA 200 (normalized)
        distance = (close - ema_200) / (ema_200 + 1e-10) * 100
        trend_features['trend_distance'] = distance

        # Trend state (above/below EMA 200)
        trend_features['trend_state'] = (close > ema_200).astype(int) * 2 - 1

        # EMA 200 slope (normalized)
        ema_slope = ema_200.diff(5) / (ema_200 + 1e-10) * 100
        trend_features['trend_slope'] = ema_slope

        # Price position in recent range relative to EMA
        rolling_high = df['high'].rolling(20).max()
        rolling_low = df['low'].rolling(20).min()
        ema_position = (ema_200 - rolling_low) / (rolling_high - rolling_low + 1e-10)
        trend_features['ema_position_in_range'] = ema_position

        # Trend consistency (how many bars in a row above/below EMA)
        above = close > ema_200
        trend_features['trend_consistency'] = (
            above.groupby((above != above.shift()).cumsum()).cumcount()
        ).clip(upper=50) * trend_features['trend_state']

        # Multi-period trend alignment
        ema_50 = close.ewm(span=50, adjust=False).mean()
        ema_100 = close.ewm(span=100, adjust=False).mean()

        # Trend alignment score (-3 to +3)
        alignment = (
            (close > ema_50).astype(int) +
            (close > ema_100).astype(int) +
            (close > ema_200).astype(int) +
            (ema_50 > ema_100).astype(int) +
            (ema_100 > ema_200).astype(int)
        )
        trend_features['trend_alignment'] = alignment - 2.5  # Center around 0

        return trend_features

    def _compute_additional_momentum(self, df: pd.DataFrame) -> pd.DataFrame:
        """Compute additional momentum indicators."""
        close = df['close']
        high = df['high']
        low = df['low']

        features = pd.DataFrame(index=df.index)

        # Williams %R
        highest_high = high.rolling(14).max()
        lowest_low = low.rolling(14).min()
        williams_r = (highest_high - close) / (highest_high - lowest_low + 1e-10) * -100
        features['williams_r'] = williams_r
        features['williams_r_norm'] = (williams_r + 50) / 50  # Normalize to [-1, 1]

        # Stochastic Oscillator
        stoch_k = (close - lowest_low) / (highest_high - lowest_low + 1e-10) * 100
        stoch_d = stoch_k.rolling(3).mean()
        features['stoch_k'] = stoch_k
        features['stoch_d'] = stoch_d
        features['stoch_diff'] = stoch_k - stoch_d

        # CCI (Commodity Channel Index)
        typical_price = (high + low + close) / 3
        sma_tp = typical_price.rolling(20).mean()
        mean_deviation = typical_price.rolling(20).apply(
            lambda x: np.abs(x - x.mean()).mean()
        )
        cci = (typical_price - sma_tp) / (0.015 * mean_deviation + 1e-10)
        features['cci'] = cci
        features['cci_norm'] = cci / 200  # Normalize roughly to [-1, 1]

        # Ultimate Oscillator
        bp = close - pd.concat([low, close.shift(1)], axis=1).min(axis=1)
        tr = pd.concat([
            high - low,
            abs(high - close.shift(1)),
            abs(low - close.shift(1))
        ], axis=1).max(axis=1)

        avg7 = bp.rolling(7).sum() / (tr.rolling(7).sum() + 1e-10)
        avg14 = bp.rolling(14).sum() / (tr.rolling(14).sum() + 1e-10)
        avg28 = bp.rolling(28).sum() / (tr.rolling(28).sum() + 1e-10)

        uo = 100 * (4 * avg7 + 2 * avg14 + avg28) / 7
        features['ultimate_osc'] = uo
        features['ultimate_osc_norm'] = (uo - 50) / 50

        return features

    def _compute_volatility_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Compute volatility-adjusted momentum features."""
        close = df['close']
        high = df['high']
        low = df['low']

        features = pd.DataFrame(index=df.index)

        # ATR for volatility normalization
        tr = pd.concat([
            high - low,
            abs(high - close.shift(1)),
            abs(low - close.shift(1))
        ], axis=1).max(axis=1)
        atr_14 = tr.rolling(14).mean()

        # Volatility ratio (current vs average)
        atr_50 = tr.rolling(50).mean()
        features['vol_ratio'] = atr_14 / (atr_50 + 1e-10)

        # Normalized range
        features['norm_range'] = (high - low) / (atr_14 + 1e-10)

        # Bollinger Band position
        sma_20 = close.rolling(20).mean()
        std_20 = close.rolling(20).std()
        bb_upper = sma_20 + 2 * std_20
        bb_lower = sma_20 - 2 * std_20
        bb_width = (bb_upper - bb_lower) / (sma_20 + 1e-10)
        bb_position = (close - bb_lower) / (bb_upper - bb_lower + 1e-10)

        features['bb_width'] = bb_width
        features['bb_position'] = bb_position
        features['bb_position_norm'] = bb_position * 2 - 1  # Normalize to [-1, 1]

        # Keltner Channel position
        ema_20 = close.ewm(span=20, adjust=False).mean()
        kc_upper = ema_20 + 2 * atr_14
        kc_lower = ema_20 - 2 * atr_14
        kc_position = (close - kc_lower) / (kc_upper - kc_lower + 1e-10)
        features['kc_position'] = kc_position
        features['kc_position_norm'] = kc_position * 2 - 1

        # Squeeze indicator (BB inside KC)
        squeeze = (bb_lower > kc_lower) & (bb_upper < kc_upper)
        features['squeeze'] = squeeze.astype(int)

        return features

    def get_feature_names(self) -> List[str]:
        """Get list of computed feature names."""
        return self._feature_names.copy()

    def get_feature_groups(self) -> Dict[str, List[str]]:
        """Get features organized by group."""
        groups = {
            'rsi': [n for n in self._feature_names if n.startswith('rsi')],
            'macd': [n for n in self._feature_names if n.startswith('macd')],
            'roc': [n for n in self._feature_names if n.startswith('roc')],
            'adx': [n for n in self._feature_names if any(x in n for x in ['adx', 'di_'])],
            'ema_crossover': [n for n in self._feature_names if n.startswith('ema_')],
            'trend': [n for n in self._feature_names if n.startswith('trend')],
            'momentum': [n for n in self._feature_names if any(
                x in n for x in ['williams', 'stoch', 'cci', 'ultimate']
            )],
            'volatility': [n for n in self._feature_names if any(
                x in n for x in ['vol_', 'bb_', 'kc_', 'squeeze', 'norm_range']
            )],
        }
        return groups


if __name__ == "__main__":
    # Test feature engineering
    import numpy as np

    np.random.seed(42)
    n_samples = 1000

    # Create sample OHLCV data
    dates = pd.date_range('2024-01-01', periods=n_samples, freq='5min')
    price = 2000 + np.cumsum(np.random.randn(n_samples) * 0.5)

    df = pd.DataFrame({
        'open': price + np.random.randn(n_samples) * 0.3,
        'high': price + np.abs(np.random.randn(n_samples)) * 1.0,
        'low': price - np.abs(np.random.randn(n_samples)) * 1.0,
        'close': price + np.random.randn(n_samples) * 0.3,
        '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)

    # Create feature engineer
    engineer = MRDFeatureEngineer()

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

    print(f"\n{'='*60}")
    print("MRD Feature Engineering Test")
    print(f"{'='*60}")
    print(f"Input samples: {len(df)}")
    print(f"Output features: {len(engineer.get_feature_names())}")
    print(f"\nFeature groups:")

    for group, names in engineer.get_feature_groups().items():
        print(f"  {group}: {len(names)} features")

    print(f"\nSample feature values (last row):")
    for col in features.columns[:10]:
        print(f"  {col}: {features[col].iloc[-1]:.4f}")

    print(f"\nFeature statistics:")
    print(features.describe().T[['mean', 'std', 'min', 'max']].head(15))
