"""Feature engineering and the end-to-end feature pipeline."""
import numpy as np
import pandas as pd
import talib
from sklearn.base import BaseEstimator, TransformerMixin
from hmmlearn import hmm
from scipy import stats
import warnings

from tradingbot.data.processor import GoldDataProcessor

warnings.filterwarnings("ignore")


class GoldFeatureEngineer(BaseEstimator, TransformerMixin):
    """
    Feature engineering for 5-minute gold trading data.
    Implements market microstructure, volatility, indicator adjustments,
    and time-based features.
    """
    def __init__(self, forecast_horizon=6, volatility_regime_states=3, look_back_window=120):
        """
        Initialize the feature engineer

        Parameters:
        -----------
        forecast_horizon : int
            Number of bars to look ahead for target creation (default: 6 bars = 30 minutes)
        volatility_regime_states : int
            Number of states for HMM volatility regime detection
        look_back_window : int
            Window size for rolling computations (default: 120 bars = 10 hours)
        """
        self.forecast_horizon = forecast_horizon
        self.volatility_regime_states = volatility_regime_states
        self.look_back_window = look_back_window
        self.regime_model = hmm.GaussianHMM(
            n_components=volatility_regime_states,
            random_state=42,
            covariance_type="diag"
        )

    def fit(self, X, y=None):
        """
        Fit method to comply with sklearn transformer interface
        Learns volatility regimes from the data
        """
        # Ensure X is a DataFrame
        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)

        # Detect volatility regimes for later transformation
        self._fit_volatility_regimes(X)

        return self

    def _clean_features(self, data):
        """
        Clean feature data by handling infinities, NaNs, and extreme values
        """
        # Replace infinities with NaN
        data = data.replace([np.inf, -np.inf], np.nan)

        # For each feature column, handle NaNs and clip extreme values
        for col in data.columns:
            if col not in ['open', 'high', 'low', 'close', 'volume', 'target', 'is_weekend']:
                # Fill NaNs with median (or forward fill if median is NaN)
                median_val = data[col].median()
                if pd.isna(median_val):
                    data[col] = data[col].fillna(method='ffill')
                    if data[col].isna().any():
                        data[col] = data[col].fillna(method='bfill')
                        if data[col].isna().any():
                            data[col] = data[col].fillna(0)
                else:
                    data[col] = data[col].fillna(median_val)

                # Clip extreme values to 5 standard deviations from mean
                if not data[col].isna().all():
                    mean_val = data[col].mean()
                    std_val = data[col].std()
                    if std_val > 0:  # Only clip if standard deviation is positive
                        data[col] = data[col].clip(
                            lower=mean_val - 5*std_val,
                            upper=mean_val + 5*std_val
                        )

        return data

    def transform(self, X):
        """
        Transform the data by adding engineered features
        """
        # Ensure X is a DataFrame and create a copy to avoid modifying original
        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)

        data = X.copy()

        # Add all feature groups
        data = self._add_market_microstructure_features(data)
        data = self._add_volatility_features(data)
        data = self._add_indicator_features(data)
        data = self._add_time_based_features(data)

        # Create target variable if close column exists
        if 'close' in data.columns:
            data = self._create_target_variable(data)

        data = self._clean_features(data)

        # Drop rows with NaN values resulting from indicators that need lookback
        data = data.dropna()

        return data

    def _fit_volatility_regimes(self, data):
        """
        Fit HMM model to detect volatility regimes
        """
        # Calculate log returns
        if 'close' in data.columns:
            returns = np.log(data['close'] / data['close'].shift(1)).dropna()

            # Calculate rolling volatility features for HMM
            rolling_vol = returns.rolling(window=20).std().dropna()
            rolling_range = ((data['high'] / data['low'] - 1)
                             .rolling(window=20).mean().dropna())

            # Combine features for regime detection
            features = pd.DataFrame({
                'returns_volatility': rolling_vol,
                'price_range': rolling_range
            }).dropna()

            if len(features) > self.volatility_regime_states * 5:  # Ensure enough data to fit
                # Normalize features for HMM
                features = (features - features.mean()) / features.std()

                # Fit the HMM model
                self.regime_model.fit(features.values)

    def _add_market_microstructure_features(self, data):
        """
        Add market microstructure features:
        - Order flow imbalance
        - Price velocity
        - Relative volume analysis
        """
        # Order Flow Imbalance (approximated from OHLC)
        data['bar_sentiment'] = np.where(
            data['close'] > data['open'],
            (data['close'] - data['open']) / (data['high'] - data['low'] + 1e-8),
            -1 * (data['open'] - data['close']) / (data['high'] - data['low'] + 1e-8)
        )

        # Smooth bar sentiment
        data['smooth_sentiment'] = data['bar_sentiment'].rolling(5).mean()

        # Price Velocity - Rate of change over multiple timeframes
        for window in [1, 3, 5, 15]:
            data[f'price_velocity_{window}'] = data['close'].pct_change(window) / (window * 5)  # Normalize by minutes

        # Acceleration (change in velocity)
        data['price_acceleration'] = data['price_velocity_3'].diff()

        # Relative Volume Analysis
        # Calculate average volume by time of day first
        data['hour'] = data.index.hour
        data['minute'] = data.index.minute
        data['time_of_day'] = data['hour'] * 60 + data['minute']

        # Group by time of day and calculate average volume
        avg_volume_by_time = data.groupby('time_of_day')['volume'].transform('mean')
        data['relative_volume'] = data['volume'] / (avg_volume_by_time + 1e-8)  # Avoid division by zero

        # Volume momentum
        data['volume_momentum'] = data['volume'].pct_change(5)

        # Drop temporary columns
        data = data.drop(['hour', 'minute', 'time_of_day'], axis=1)

        return data

    def _add_volatility_features(self, data):
        """
        Add volatility-related features:
        - Micro-volatility clusters
        - Bollinger Band Width
        - Volatility regime detection
        """
        # ATR with shorter periods for micro-volatility
        for period in [5, 10]:
            data[f'atr_{period}'] = talib.ATR(
                data['high'].values,
                data['low'].values,
                data['close'].values,
                timeperiod=period
            )
            # Normalize ATR by price level
            data[f'atr_{period}_pct'] = data[f'atr_{period}'] / data['close']

        # Bollinger Band Width (20 periods, 2 std)
        upper, middle, lower = talib.BBANDS(
            data['close'].values,
            timeperiod=20,
            nbdevup=2,
            nbdevdn=2
        )
        data['bb_width'] = (upper - lower) / middle

        # Rate of change of BB width
        data['bb_width_change'] = data['bb_width'].pct_change(3)

        # Volatility Regime Detection
        if hasattr(self, 'regime_model') and hasattr(self.regime_model, 'transmat_'):
            # Calculate the same features used during fitting
            returns = np.log(data['close'] / data['close'].shift(1))
            rolling_vol = returns.rolling(window=20).std()
            rolling_range = (data['high'] / data['low'] - 1).rolling(window=20).mean()

            # Combine and normalize features
            features = pd.DataFrame({
                'returns_volatility': rolling_vol,
                'price_range': rolling_range
            })

            # Handle NaN values
            features = features.fillna(method='bfill')

            # Normalize using the same approach as in fitting
            features = (features - features.mean()) / features.std()

            # Predict regimes where we have data
            valid_idx = ~features.isnull().any(axis=1)
            if valid_idx.sum() > 0:
                regimes = self.regime_model.predict(features[valid_idx].values)

                # Create a series with index aligned to original data
                regime_series = pd.Series(index=data.index, dtype='float64')
                regime_series.loc[features[valid_idx].index] = regimes

                # Forward fill regime values
                data['volatility_regime'] = regime_series.fillna(method='ffill')
            else:
                # If we don't have enough data, use a default regime
                data['volatility_regime'] = 1
        else:
            # If regime model isn't fitted, use a simpler approach
            returns = np.log(data['close'] / data['close'].shift(1))
            vol = returns.rolling(window=20).std() * np.sqrt(252 * 288)  # Annualized (288 5-min bars/day)
            data['volatility_regime'] = pd.qcut(
                vol,
                q=self.volatility_regime_states,
                labels=False,
                duplicates='drop'
            ).fillna(self.volatility_regime_states // 2)

        return data

    def _add_indicator_features(self, data):
        """
        Add adjusted technical indicators specific for 5-minute data:
        - Fast & Slow EMAs
        - RSI
        - Stochastic Oscillator
        - MACD
        """
        # EMAs with shorter periods
        for period in [10, 20, 50]:
            data[f'ema_{period}'] = talib.EMA(data['close'].values, timeperiod=period)

            # Add relative position to EMA
            data[f'close_to_ema_{period}'] = (data['close'] / data[f'ema_{period}'] - 1) * 100

        # EMA Cross Features
        data['ema_10_20_cross'] = data['ema_10'] - data['ema_20']
        data['ema_10_50_cross'] = data['ema_10'] - data['ema_50']

        # Shorter-period RSI
        data['rsi_7'] = talib.RSI(data['close'].values, timeperiod=7)
        data['rsi_14'] = talib.RSI(data['close'].values, timeperiod=14)

        # RSI momentum and mean-reversion features
        data['rsi_7_change'] = data['rsi_7'].diff(3)
        data['rsi_divergence'] = data['rsi_14'] - data['rsi_7']

        # Stochastic Oscillator - Fast and Slow
        slowk, slowd = talib.STOCH(
            data['high'].values,
            data['low'].values,
            data['close'].values,
            fastk_period=5,
            slowk_period=3,
            slowk_matype=0,
            slowd_period=3,
            slowd_matype=0
        )
        data['stoch_k_fast'] = slowk
        data['stoch_d_fast'] = slowd

        slowk, slowd = talib.STOCH(
            data['high'].values,
            data['low'].values,
            data['close'].values,
            fastk_period=14,
            slowk_period=3,
            slowk_matype=0,
            slowd_period=3,
            slowd_matype=0
        )
        data['stoch_k_slow'] = slowk
        data['stoch_d_slow'] = slowd

        # Stochastic crossover signal
        data['stoch_crossover'] = data['stoch_k_fast'] - data['stoch_d_fast']

        # MACD with faster settings
        macd, macdsignal, macdhist = talib.MACD(
            data['close'].values,
            fastperiod=8,
            slowperiod=17,
            signalperiod=9
        )
        data['macd'] = macd
        data['macd_signal'] = macdsignal
        data['macd_hist'] = macdhist

        # Add momentum oscillator
        data['mom_10'] = talib.MOM(data['close'].values, timeperiod=10)

        # Add Average Directional Index for trend strength
        data['adx_14'] = talib.ADX(
            data['high'].values,
            data['low'].values,
            data['close'].values,
            timeperiod=14
        )

        return data

    def _add_time_based_features(self, data):
        """
        Add time-based features:
        - Time of day (hour, minute)
        - Day of week
        - Session (Tokyo, London, New York)
        """
        # Extract basic time components
        data['hour'] = data.index.hour
        data['minute'] = data.index.minute
        data['day_of_week'] = data.index.dayofweek

        # Create cyclical time features (circular encoding to handle continuity)
        # For hour of day (24 hours)
        data['hour_sin'] = np.sin(2 * np.pi * data['hour'] / 24)
        data['hour_cos'] = np.cos(2 * np.pi * data['hour'] / 24)

        # For minute within hour
        data['minute_sin'] = np.sin(2 * np.pi * data['minute'] / 60)
        data['minute_cos'] = np.cos(2 * np.pi * data['minute'] / 60)

        # For day of week (5-day trading week)
        data['day_sin'] = np.sin(2 * np.pi * data['day_of_week'] / 7)
        data['day_cos'] = np.cos(2 * np.pi * data['day_of_week'] / 7)

        # Create Trading Session markers
        # Define trading sessions (UTC times)
        # Tokyo: 00:00-09:00
        # London: 08:00-17:00
        # New York: 13:00-22:00

        # Convert hour to UTC assuming index is in UTC
        # If index is not in UTC, this would need to be adjusted
        utc_hour = data['hour']

        # Trading session flags
        data['tokyo_session'] = ((utc_hour >= 0) & (utc_hour < 9)).astype(int)
        data['london_session'] = ((utc_hour >= 8) & (utc_hour < 17)).astype(int)
        data['ny_session'] = ((utc_hour >= 13) & (utc_hour < 22)).astype(int)

        # Session overlap flags
        data['tokyo_london_overlap'] = ((utc_hour >= 8) & (utc_hour < 9)).astype(int)
        data['london_ny_overlap'] = ((utc_hour >= 13) & (utc_hour < 17)).astype(int)

        # Drop original time columns
        data = data.drop(['hour', 'minute', 'day_of_week'], axis=1)

        return data

    def _create_target_variable(self, data):
        """
        Create target variable for prediction:
        Multi-class classification based on future price movement
        """
        # Calculate future return over forecast_horizon
        future_close = data['close'].shift(-self.forecast_horizon)
        future_return = (future_close / data['close'] - 1) * 100  # Percentage return

        # Create multi-class target
        # Calculate dynamic thresholds based on recent volatility
        volatility = data['close'].pct_change().rolling(window=20).std() * 100  # Convert to percentage

        # Define thresholds as a function of volatility
        # More volatile periods have wider thresholds
        strong_threshold = volatility * 0.75  # 75% of recent volatility
        weak_threshold = volatility * 0.25    # 25% of recent volatility

        # Ensure minimum thresholds
        strong_threshold = np.maximum(strong_threshold, 0.15)  # At least 0.15%
        weak_threshold = np.maximum(weak_threshold, 0.05)     # At least 0.05%

        # Create targets
        conditions = [
            future_return > strong_threshold,
            (future_return > weak_threshold) & (future_return <= strong_threshold),
            (future_return >= -weak_threshold) & (future_return <= weak_threshold),
            (future_return < -weak_threshold) & (future_return >= -strong_threshold),
            future_return < -strong_threshold
        ]

        choices = [2, 1, 0, -1, -2]  # Strong Up, Weak Up, Sideways, Weak Down, Strong Down

        data['target'] = np.select(conditions, choices, default=np.nan)

        # Store the thresholds used for reference
        data['strong_threshold'] = strong_threshold
        data['weak_threshold'] = weak_threshold

        # Also store raw future return for potential regression tasks
        data['future_return'] = future_return

        return data

class GoldFeaturePipeline:
    """
    Complete pipeline for gold 5-minute data preprocessing and feature engineering
    """
    def __init__(self, forecast_horizon=6, volatility_regime_states=3, remove_outliers=True):
        """
        Initialize the complete pipeline

        Parameters:
        -----------
        forecast_horizon : int
            Number of bars to look ahead for target creation
        volatility_regime_states : int
            Number of states for HMM volatility regime detection
        remove_outliers : bool
            Whether to remove price outliers
        """
        self.data_processor = GoldDataProcessor(
            fill_gaps=True,
            remove_outliers=remove_outliers,
            handle_after_hours=True
        )

        self.feature_engineer = GoldFeatureEngineer(
            forecast_horizon=forecast_horizon,
            volatility_regime_states=volatility_regime_states
        )

    def process(self, file_path):
        """
        Process data from file through the complete pipeline

        Parameters:
        -----------
        file_path : str
            Path to the CSV file with 5-minute OHLCV data

        Returns:
        --------
        processed_data : pandas.DataFrame
            Processed data with all features and target
        """
        # Load and preprocess data
        data = self.data_processor.load_data(file_path)

        # Apply feature engineering
        self.feature_engineer.fit(data)
        processed_data = self.feature_engineer.transform(data)

        # Return the processed data
        return processed_data