"""
Feature Engineering Component.

Handles adding base features (cyclical, imbalance, TA) and selecting features
using Logistic Regression (L1) and Variance Inflation Factor (VIF).
"""

import pandas as pd
import numpy as np
import logging
import json

from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import SelectFromModel
from statsmodels.stats.outliers_influence import variance_inflation_factor
import statsmodels.api as sm

# Import TA library functions directly
from ta.volatility import AverageTrueRange
from ta.momentum import RSIIndicator
from ta.trend import EMAIndicator, MACD

logger = logging.getLogger(__name__)

_EPS = 1e-6

class FeatureEngineer:
    """Encapsulates feature creation and selection logic."""

    def __init__(self, minimal_whitelist: list):
        """
        Initialize the FeatureEngineer.

        Args:
            minimal_whitelist (list): A base list of features considered essential.
        """
        self.minimal_whitelist = minimal_whitelist
        logger.info(f"FeatureEngineer initialized with minimal whitelist: {self.minimal_whitelist}")

    def _add_cyclical_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Adds sine and cosine transformations of the hour."""
        if isinstance(df.index, pd.DatetimeIndex):
            timestamp_source = df.index
            logger.info("Adding cyclical hour features (sin/cos)...")
            df['hour_sin'] = np.sin(2 * np.pi * timestamp_source.hour / 24)
            df['hour_cos'] = np.cos(2 * np.pi * timestamp_source.hour / 24)
        else:
            logger.warning("Index is not DatetimeIndex. Skipping cyclical hour features.")
            # Add placeholders if needed by downstream code, though it's better to ensure datetime index upstream
            df['hour_sin'] = 0.0
            df['hour_cos'] = 1.0 
        return df

    def _add_imbalance_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Add Chaikin AD line, signed volume imbalance, gap imbalance."""
        logger.info("Adding imbalance features...")
        if not {"open", "high", "low", "close", "volume"}.issubset(df.columns):
            logger.warning("Missing required columns for imbalance features. Skipping.")
            return df

        try:
            clv = ((df["close"] - df["low"]) - (df["high"] - df["close"])) / (
                df["high"] - df["low"] + _EPS
            )
            df["chaikin_AD_10"] = (clv * df["volume"]).rolling(10).sum()

            signed_vol = np.where(df["close"] >= df["open"], df["volume"], -df["volume"])
            df["svi_10"] = pd.Series(signed_vol, index=df.index).rolling(10).sum()

            med_vol = df["volume"].rolling(50).median()
            gap_up = (df["low"] > df["high"].shift(1)) & (df["volume"] > 2 * med_vol)
            gap_dn = (df["high"] < df["low"].shift(1)) & (df["volume"] > 2 * med_vol)
            df["gap_imbalance"] = gap_up.astype(int) - gap_dn.astype(int)

            # Fill NaNs introduced - use fillna(0) for simplicity here
            # More robust filling (bfill/ffill) might be needed depending on feature use
            df.fillna({"chaikin_AD_10": 0, "svi_10": 0, "gap_imbalance": 0}, inplace=True)
            logger.info("Successfully added imbalance features.")
        except Exception as e:
            logger.error(f"Error calculating imbalance features: {e}", exc_info=True)

        return df

    def _add_ta_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Adds TA features using the 'ta' library."""
        logger.info("Adding TA features...")
        required_cols = {'open', 'high', 'low', 'close', 'volume'}
        if not required_cols.issubset(df.columns):
             logger.warning(f"Missing required columns for TA features ({required_cols - set(df.columns)}). Skipping TA.")
             return df

        # Apply shift(1) to prevent lookahead bias in TA features based on close
        # Features will be calculated based on data up to t-1
        df_shifted = df.shift(1)
        df_ta = pd.DataFrame(index=df.index) # Create empty DF to store results aligned with original index

        try:
            # Calculate returns first (use shifted close)
            # Fill NaNs robustly before pct_change on the *shifted* data
            close_filled = df_shifted["close"].bfill().ffill()
            df_ta["return_1m"] = close_filled.pct_change()
            df_ta["return_15m"] = close_filled.pct_change(15)
            df_ta["return_60m"] = close_filled.pct_change(60)

            # Calculate TA features using ta library on *shifted* data
            df_ta["ATR_14"] = AverageTrueRange(df_shifted['high'], df_shifted['low'], df_shifted['close'], window=14).average_true_range()

            # Daily volatility (use calculated 1m return)
            df_ta["volatility_14d"] = (
                df_ta["return_1m"].rolling(60 * 24 * 14, min_periods=30).std() # rough 14d for 1‑min bars
            )

            # EMA 10 / 50 + MACD using ta library (on shifted close)
            df_ta["EMA_10"] = EMAIndicator(df_shifted["close"], 10).ema_indicator()
            df_ta["EMA_50"] = EMAIndicator(df_shifted["close"], 50).ema_indicator()
            macd = MACD(df_shifted["close"], window_slow=26, window_fast=12, window_sign=9)
            df_ta["MACD"] = macd.macd()
            df_ta["MACD_signal"] = macd.macd_signal()

            # RSI 14 using ta library (on shifted close)
            df_ta["RSI_14"] = RSIIndicator(df_shifted["close"], window=14).rsi()

            # Handle potential NaNs introduced by TA calculations
            df_ta.bfill(inplace=True)
            df_ta.ffill(inplace=True)

            # Add the calculated TA features back to the original df
            ta_cols_to_add = [col for col in df_ta.columns if col not in df.columns]
            for col in ta_cols_to_add:
                df[col] = df_ta[col]
            logger.info("Successfully added TA features.")

        except Exception as e:
            logger.error(f"Error calculating TA features: {e}", exc_info=True)

        return df

    def add_base_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Adds a standard set of base features: cyclical, imbalance, and TA.

        Args:
            df (pd.DataFrame): Input DataFrame with OHLCV data and DatetimeIndex.

        Returns:
            pd.DataFrame: DataFrame with added features.
        """
        logger.info("--- Adding Base Features --- ")
        df_out = df.copy()
        df_out = self._add_cyclical_features(df_out)
        df_out = self._add_imbalance_features(df_out)
        df_out = self._add_ta_features(df_out)

        # Ensure minimal whitelist columns exist, fill with 0 if missing after calculation errors
        for col in self.minimal_whitelist:
            if col not in df_out.columns:
                logger.warning(f"Minimal whitelist feature '{col}' not found after calculations. Adding column filled with 0.")
                df_out[col] = 0.0

        logger.info(f"Base feature engineering complete. DataFrame shape: {df_out.shape}")
        return df_out

    def select_features(self, X_train_raw: pd.DataFrame, y_dir_train: pd.Series, vif_threshold: float = 10.0, logreg_c: float = 0.1) -> list:
        """
        Performs feature selection using Logistic Regression (L1) followed by VIF filtering.

        Args:
            X_train_raw (pd.DataFrame): Raw training features (should include minimal whitelist + others).
            y_dir_train (pd.Series): Binary direction target for training.
            vif_threshold (float): VIF threshold for multicollinearity removal.
            logreg_c (float): Inverse of regularization strength for Logistic Regression.

        Returns:
            list: The final list of selected feature names.
        """
        logger.info("--- Selecting Features (LogReg L1 + VIF) ---")

        if X_train_raw is None or X_train_raw.empty or y_dir_train is None or y_dir_train.empty:
            logger.error("Input data for feature selection is missing or empty. Returning minimal whitelist.")
            return self.minimal_whitelist
        
        initial_features = X_train_raw.columns.tolist()
        logger.info(f"Starting selection from {len(initial_features)} raw features.")

        final_whitelist = self.minimal_whitelist # Default to minimal if errors occur

        try:
            # --- LogReg L1 Selection --- 
            logger.info(f"Performing Logistic Regression (L1, C={logreg_c}) selection...")
            
            # Handle potential NaN/Inf values robustly before fitting
            X_train_processed = X_train_raw.copy()
            X_train_processed.replace([np.inf, -np.inf], np.nan, inplace=True)
            # Check if imputation is needed
            if X_train_processed.isnull().any().any():
                logger.warning("NaNs detected in training features for LogReg. Imputing with column median.")
                # Impute with median, could use other strategies (mean, IterativeImputer)
                X_train_processed = X_train_processed.fillna(X_train_processed.median())
                # Check if any NaNs remain (e.g., all values were NaN)
                if X_train_processed.isnull().any().any():
                     logger.warning("NaNs still present after median imputation. Filling remaining NaNs with 0.")
                     X_train_processed.fillna(0, inplace=True)

            logreg_selector = LogisticRegression(
                penalty='l1',
                solver='liblinear',
                C=logreg_c,
                random_state=42,
                max_iter=1000,
                class_weight='balanced' # Added for potentially imbalanced targets
            )

            feature_selector = SelectFromModel(estimator=logreg_selector, prefit=False, threshold=-np.inf) # Select all features with non-zero coefficients
            feature_selector.fit(X_train_processed, y_dir_train)

            selected_features_mask = feature_selector.get_support()
            selected_features_names = X_train_raw.columns[selected_features_mask].tolist()
            logger.info(f"Features selected by LogReg L1: {selected_features_names}")

            # Combine with the minimal whitelist
            candidate_whitelist_set = set(self.minimal_whitelist).union(set(selected_features_names))
            candidate_whitelist = sorted(list(candidate_whitelist_set))
            logger.info(f"Candidate whitelist after LogReg ({len(candidate_whitelist)} features): {candidate_whitelist}")

            # --- VIF Filtering --- 
            logger.info(f"Performing VIF filtering (threshold={vif_threshold}) on candidate features...")
            
            # Prepare data subset for VIF calculation
            X_vif = X_train_raw[candidate_whitelist].copy()

            # Handle NaN/Inf robustly before VIF
            X_vif.replace([np.inf, -np.inf], np.nan, inplace=True)
            if X_vif.isnull().any().any():
                logger.warning("NaNs detected in features for VIF. Imputing with median.")
                X_vif = X_vif.fillna(X_vif.median())
                if X_vif.isnull().any().any():
                     logger.warning("NaNs still present after median imputation for VIF. Filling remaining NaNs with 0.")
                     X_vif.fillna(0, inplace=True)

            features_after_vif = list(candidate_whitelist) # Start with the candidate list
            
            # Iteratively remove features with high VIF
            while len(features_after_vif) > 1:
                # Add constant for VIF calculation
                X_vif_subset = X_vif[features_after_vif]
                X_vif_const = sm.add_constant(X_vif_subset, prepend=False, has_constant='raise')
                
                try:
                    vif_data = pd.DataFrame()
                    vif_data["feature"] = X_vif_const.columns[:-1] # Exclude the added constant
                    vif_values = [variance_inflation_factor(X_vif_const.values, i)
                                  for i in range(X_vif_const.shape[1] - 1)]
                    vif_data["VIF"] = vif_values
                except Exception as vif_e:
                    logger.error(f"Error calculating VIF: {vif_e}. Stopping VIF filtering.")
                    # Return the features identified so far before the error
                    final_whitelist = features_after_vif 
                    logger.warning(f"Returning {len(final_whitelist)} features identified before VIF error: {final_whitelist}")
                    return final_whitelist
                
                max_vif = vif_data["VIF"].max()
                # Check for infinite VIF separately, often indicates perfect multicollinearity
                if np.isinf(max_vif):
                    feature_to_remove = vif_data.loc[vif_data["VIF"].idxmax(), "feature"]
                    logger.warning(f"Removing feature '{feature_to_remove}' due to infinite VIF.")
                elif max_vif > vif_threshold:
                    feature_to_remove = vif_data.loc[vif_data["VIF"].idxmax(), "feature"]
                    logger.info(f"Removing feature '{feature_to_remove}' due to high VIF ({max_vif:.2f} > {vif_threshold})...")
                else:
                    logger.info(f"VIF filtering complete. Max VIF = {max_vif:.2f} <= {vif_threshold}.")
                    break # All VIFs are below threshold
                
                features_after_vif.remove(feature_to_remove)
            
            final_whitelist = features_after_vif
            if len(final_whitelist) == 0:
                logger.warning("VIF filtering removed all features. Falling back to minimal whitelist.")
                final_whitelist = self.minimal_whitelist
            else:
                logger.info(f"Final whitelist after VIF filtering ({len(final_whitelist)} features): {final_whitelist}")

        except Exception as e:
            logger.error(f"Feature selection or VIF filtering failed: {e}", exc_info=True)
            logger.warning("Falling back to the minimal_whitelist.")
            final_whitelist = self.minimal_whitelist
        
        # Ensure final whitelist is a list of strings
        if not isinstance(final_whitelist, list):
            final_whitelist = list(final_whitelist)
            
        return final_whitelist

    def prune_features(self, df: pd.DataFrame, whitelist: list) -> pd.DataFrame:
        """
        Selects columns from a DataFrame based on a provided whitelist.

        Args:
            df (pd.DataFrame): The input DataFrame.
            whitelist (list): List of column names to keep.

        Returns:
            pd.DataFrame: DataFrame containing only the whitelisted columns.
        """
        if df is None or df.empty:
             logger.warning("Input DataFrame for pruning is None or empty. Returning empty DataFrame.")
             return pd.DataFrame()
             
        logger.debug(f"Pruning DataFrame (shape {df.shape}) to whitelist: {whitelist}")
        
        # Find columns present in both DataFrame and whitelist
        cols_to_keep = [c for c in whitelist if c in df.columns]
        missing_cols = [c for c in whitelist if c not in df.columns]
        
        if missing_cols:
            logger.warning(f"The following whitelisted columns were not found in the DataFrame: {missing_cols}")
        
        if not cols_to_keep:
            logger.error("No columns from the whitelist found in the DataFrame. Returning empty DataFrame.")
            return pd.DataFrame()
            
        # Select the columns
        df_pruned = df[cols_to_keep].copy()
        logger.debug(f"Pruned DataFrame shape: {df_pruned.shape}")
        
        # Verification (optional but good practice)
        if set(df_pruned.columns) != set(cols_to_keep):
             logger.error(f"Pruning verification failed: Output columns {set(df_pruned.columns)} != Expected {set(cols_to_keep)}")
             # Potentially return original df or raise error depending on desired strictness
             # For now, return the potentially incorrect df_pruned

        return df_pruned 