"""
Auto Preprocessor Module

Intelligent automated preprocessing with AI-powered recommendations
for handling missing values, scaling, encoding, and feature engineering.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from sklearn.preprocessing import (
    StandardScaler, MinMaxScaler, RobustScaler, 
    LabelEncoder, OneHotEncoder, OrdinalEncoder
)
from sklearn.impute import SimpleImputer, KNNImputer
from sklearn.feature_selection import (
    SelectKBest, f_classif, f_regression, 
    mutual_info_classif, mutual_info_regression
)
from sklearn.decomposition import PCA
import warnings


class AutoPreprocessor:
    """
    Intelligent automated preprocessor that applies appropriate
    preprocessing steps based on data analysis.
    """
    
    def __init__(self):
        self.preprocessing_steps = []
        self.fitted_transformers = {}
        self.feature_types = {}
        self.preprocessing_log = []
        
    def fit(self, df: pd.DataFrame, feature_types: Dict[str, str], 
            target_column: Optional[str] = None) -> 'AutoPreprocessor':
        """
        Fit the preprocessor based on data analysis
        
        Args:
            df: Input DataFrame
            feature_types: Dictionary mapping column names to feature types
            target_column: Name of target column
            
        Returns:
            Self for method chaining
        """
        self.feature_types = feature_types
        self.preprocessing_steps = []
        self.fitted_transformers = {}
        
        # Separate features and target
        if target_column and target_column in df.columns:
            X = df.drop(columns=[target_column])
            y = df[target_column]
        else:
            X = df.copy()
            y = None
        
        # Process each feature type
        processed_X = X.copy()
        
        # Handle missing values
        processed_X = self._handle_missing_values(processed_X)
        
        # Encode categorical variables
        processed_X = self._encode_categorical(processed_X)
        
        # Scale numerical variables
        processed_X = self._scale_numerical(processed_X, y)
        
        # Feature selection/engineering
        processed_X = self._feature_selection(processed_X, y)
        
        return self
    
    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Transform new data using fitted preprocessing steps
        
        Args:
            df: Input DataFrame to transform
            
        Returns:
            Transformed DataFrame
        """
        if not self.fitted_transformers:
            raise ValueError("Preprocessor must be fitted before transforming data")
        
        result = df.copy()
        
        # Apply transformations in order
        for step_name, transformer in self.fitted_transformers.items():
            if step_name.startswith('imputer_'):
                cols = [col for col in transformer.feature_names_in_ if col in result.columns]
                if cols:
                    result[cols] = transformer.transform(result[cols])
            elif step_name.startswith('encoder_'):
                if hasattr(transformer, 'categories_'):
                    cols = transformer.feature_names_in_
                    if all(col in result.columns for col in cols):
                        encoded = transformer.transform(result[cols])
                        if hasattr(transformer, 'get_feature_names_out'):
                            feature_names = transformer.get_feature_names_out()
                        else:
                            feature_names = [f"{col}_encoded" for col in cols]
                        
                        # Remove original columns and add encoded ones
                        result = result.drop(columns=cols)
                        for i, name in enumerate(feature_names):
                            result[name] = encoded[:, i] if encoded.ndim > 1 else encoded
            elif step_name.startswith('scaler_'):
                cols = [col for col in transformer.feature_names_in_ if col in result.columns]
                if cols:
                    result[cols] = transformer.transform(result[cols])
            elif step_name.startswith('selector_'):
                if hasattr(transformer, 'get_support'):
                    selected_features = result.columns[transformer.get_support()]
                    result = result[selected_features]
        
        return result
    
    def fit_transform(self, df: pd.DataFrame, feature_types: Dict[str, str],
                     target_column: Optional[str] = None) -> pd.DataFrame:
        """Fit and transform in one step"""
        return self.fit(df, feature_types, target_column).transform(df)
    
    def _handle_missing_values(self, df: pd.DataFrame) -> pd.DataFrame:
        """Handle missing values based on feature type and missing percentage"""
        result = df.copy()
        
        for col in df.columns:
            missing_percentage = df[col].isnull().sum() / len(df)
            
            if missing_percentage > 0:
                feature_type = self.feature_types.get(col, 'unknown')
                
                if missing_percentage < 0.05:  # < 5% missing
                    if feature_type == 'numerical':
                        imputer = SimpleImputer(strategy='median')
                    else:
                        imputer = SimpleImputer(strategy='most_frequent')
                    
                    result[col] = imputer.fit_transform(df[[col]]).ravel()
                    self.fitted_transformers[f'imputer_{col}'] = imputer
                    self.preprocessing_log.append(f"Applied {imputer.strategy} imputation to {col}")
                    
                elif missing_percentage < 0.3:  # 5-30% missing
                    if feature_type == 'numerical':
                        imputer = KNNImputer(n_neighbors=5)
                        result[col] = imputer.fit_transform(df[[col]]).ravel()
                        self.fitted_transformers[f'imputer_{col}'] = imputer
                        self.preprocessing_log.append(f"Applied KNN imputation to {col}")
                    else:
                        imputer = SimpleImputer(strategy='most_frequent')
                        result[col] = imputer.fit_transform(df[[col]]).ravel()
                        self.fitted_transformers[f'imputer_{col}'] = imputer
                        self.preprocessing_log.append(f"Applied mode imputation to {col}")
                else:  # > 30% missing
                    # Create missing indicator
                    result[f'{col}_missing'] = df[col].isnull().astype(int)
                    self.preprocessing_log.append(f"Created missing indicator for {col}")
        
        return result
    
    def _encode_categorical(self, df: pd.DataFrame) -> pd.DataFrame:
        """Encode categorical variables"""
        result = df.copy()
        
        categorical_cols = [col for col in df.columns 
                          if self.feature_types.get(col) == 'categorical']
        
        for col in categorical_cols:
            n_unique = df[col].nunique()
            
            if n_unique == 2:  # Binary
                encoder = LabelEncoder()
                result[col] = encoder.fit_transform(df[col].astype(str))
                self.fitted_transformers[f'encoder_{col}'] = encoder
                self.preprocessing_log.append(f"Applied label encoding to {col}")
                
            elif n_unique <= 10:  # Low cardinality
                encoder = OneHotEncoder(drop='first', sparse_output=False)
                encoded = encoder.fit_transform(df[[col]].astype(str))
                feature_names = [f"{col}_{cat}" for cat in encoder.categories_[0][1:]]
                
                # Add encoded columns
                for i, name in enumerate(feature_names):
                    result[name] = encoded[:, i]
                
                result = result.drop(columns=[col])
                self.fitted_transformers[f'encoder_{col}'] = encoder
                self.preprocessing_log.append(f"Applied one-hot encoding to {col}")
                
            else:  # High cardinality
                encoder = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1)
                result[col] = encoder.fit_transform(df[[col]].astype(str)).ravel()
                self.fitted_transformers[f'encoder_{col}'] = encoder
                self.preprocessing_log.append(f"Applied ordinal encoding to {col}")
        
        return result
    
    def _scale_numerical(self, df: pd.DataFrame, y: Optional[pd.Series] = None) -> pd.DataFrame:
        """Scale numerical variables"""
        result = df.copy()
        
        numerical_cols = [col for col in df.columns 
                         if self.feature_types.get(col) == 'numerical']
        
        if not numerical_cols:
            return result
        
        # Choose scaling method based on data distribution
        if y is not None and pd.api.types.is_numeric_dtype(y):
            # For regression tasks, use robust scaling if outliers are present
            outlier_count = 0
            for col in numerical_cols:
                Q1 = df[col].quantile(0.25)
                Q3 = df[col].quantile(0.75)
                IQR = Q3 - Q1
                outliers = ((df[col] < (Q1 - 1.5 * IQR)) | (df[col] > (Q3 + 1.5 * IQR))).sum()
                outlier_count += outliers
            
            if outlier_count > len(df) * 0.05:  # More than 5% outliers
                scaler = RobustScaler()
                scaling_method = "robust"
            else:
                scaler = StandardScaler()
                scaling_method = "standard"
        else:
            # For classification or no target, use standard scaling
            scaler = StandardScaler()
            scaling_method = "standard"
        
        result[numerical_cols] = scaler.fit_transform(df[numerical_cols])
        self.fitted_transformers[f'scaler_numerical'] = scaler
        self.preprocessing_log.append(f"Applied {scaling_method} scaling to numerical features")
        
        return result
    
    def _feature_selection(self, df: pd.DataFrame, y: Optional[pd.Series] = None) -> pd.DataFrame:
        """Apply feature selection if beneficial"""
        if y is None or len(df.columns) < 10:
            return df
        
        result = df.copy()
        
        # Determine if it's classification or regression
        if pd.api.types.is_numeric_dtype(y) and y.nunique() > 20:
            task_type = 'regression'
            score_func = f_regression
        else:
            task_type = 'classification'
            score_func = f_classif
        
        # Apply feature selection if we have many features
        if len(df.columns) > 20:
            # Select top k features
            k = min(50, len(df.columns) // 2)
            selector = SelectKBest(score_func=score_func, k=k)
            
            # Ensure y is properly encoded for classification
            if task_type == 'classification':
                from sklearn.preprocessing import LabelEncoder
                le = LabelEncoder()
                y_encoded = le.fit_transform(y.astype(str))
            else:
                y_encoded = y
            
            try:
                selector.fit(df, y_encoded)
                selected_features = df.columns[selector.get_support()]
                result = result[selected_features]
                self.fitted_transformers['selector_features'] = selector
                self.preprocessing_log.append(f"Selected {len(selected_features)} features using SelectKBest")
            except Exception as e:
                self.preprocessing_log.append(f"Feature selection failed: {str(e)}")
        
        return result
    
    def get_preprocessing_pipeline(self) -> Dict[str, Any]:
        """Get the complete preprocessing pipeline for reproducibility"""
        return {
            'steps': self.preprocessing_steps,
            'transformers': {k: str(v) for k, v in self.fitted_transformers.items()},
            'feature_types': self.feature_types,
            'log': self.preprocessing_log
        }
    
    def get_feature_importance_after_preprocessing(self, df: pd.DataFrame, 
                                                  y: pd.Series) -> Dict[str, float]:
        """Get feature importance after preprocessing"""
        if not self.fitted_transformers:
            raise ValueError("Preprocessor must be fitted first")
        
        # Transform the data
        X_transformed = self.transform(df)
        
        # Use mutual information for feature importance
        if pd.api.types.is_numeric_dtype(y) and y.nunique() > 20:
            importance = mutual_info_regression(X_transformed, y)
        else:
            importance = mutual_info_classif(X_transformed, y)
        
        feature_importance = dict(zip(X_transformed.columns, importance))
        
        # Sort by importance
        sorted_importance = dict(sorted(feature_importance.items(), 
                                      key=lambda x: x[1], reverse=True))
        
        return sorted_importance
