"""
Pipeline Creator

Automated scikit-learn pipeline creation with intelligent preprocessing
steps, model selection, and reproducible workflow management.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple, Union
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import (
    StandardScaler, MinMaxScaler, RobustScaler,
    OneHotEncoder, OrdinalEncoder, LabelEncoder,
    FunctionTransformer, PowerTransformer
)
from sklearn.impute import SimpleImputer, KNNImputer
from sklearn.feature_selection import (
    SelectKBest, SelectPercentile, RFE, RFECV,
    f_classif, f_regression, mutual_info_classif, mutual_info_regression
)
from sklearn.decomposition import PCA, TruncatedSVD
from sklearn.model_selection import cross_val_score
import joblib
import json
from datetime import datetime
import hashlib
import warnings
from dataclasses import dataclass, asdict


@dataclass
class PipelineConfig:
    """Configuration for pipeline creation"""
    name: str
    problem_type: str  # 'classification', 'regression', 'clustering'
    target_column: Optional[str] = None
    numeric_features: List[str] = None
    categorical_features: List[str] = None
    text_features: List[str] = None
    datetime_features: List[str] = None
    handle_missing: bool = True
    scaling_method: str = 'standard'  # 'standard', 'minmax', 'robust', 'none'
    encoding_method: str = 'auto'  # 'auto', 'onehot', 'ordinal', 'target'
    feature_selection: bool = False
    feature_selection_method: str = 'k_best'  # 'k_best', 'percentile', 'rfe'
    n_features_to_select: int = 10
    dimensionality_reduction: bool = False
    dim_reduction_method: str = 'pca'  # 'pca', 'svd'
    n_components: int = 0.95  # Can be int or float for variance ratio
    custom_transformers: List[str] = None
    
    def __post_init__(self):
        if self.numeric_features is None:
            self.numeric_features = []
        if self.categorical_features is None:
            self.categorical_features = []
        if self.text_features is None:
            self.text_features = []
        if self.datetime_features is None:
            self.datetime_features = []
        if self.custom_transformers is None:
            self.custom_transformers = []


@dataclass
class PipelineMetadata:
    """Metadata for pipeline versioning and tracking"""
    pipeline_id: str
    name: str
    version: str
    created_at: str
    config: Dict[str, Any]
    performance_metrics: Dict[str, float]
    feature_names: List[str]
    target_name: Optional[str]
    model_type: str
    data_hash: str
    pipeline_hash: str
    description: Optional[str] = None


class PipelineCreator:
    """
    Automated scikit-learn pipeline creator with intelligent preprocessing
    and comprehensive pipeline management.
    """
    
    def __init__(self):
        self.pipelines = {}
        self.metadata = {}
        self.current_pipeline = None
        self.current_config = None
        
    def create_pipeline(self, config: PipelineConfig, model: Any) -> Pipeline:
        """
        Create a complete scikit-learn pipeline
        
        Args:
            config: Pipeline configuration
            model: Scikit-learn compatible model
            
        Returns:
            Configured scikit-learn Pipeline
        """
        self.current_config = config
        
        # Create preprocessing steps
        preprocessing_steps = self._create_preprocessing_steps(config)
        
        # Create feature selection steps
        feature_selection_steps = self._create_feature_selection_steps(config)
        
        # Create dimensionality reduction steps
        dim_reduction_steps = self._create_dimensionality_reduction_steps(config)
        
        # Combine all steps
        all_steps = []
        if preprocessing_steps:
            all_steps.extend(preprocessing_steps)
        if feature_selection_steps:
            all_steps.extend(feature_selection_steps)
        if dim_reduction_steps:
            all_steps.extend(dim_reduction_steps)
        
        # Add model as final step
        all_steps.append(('model', model))
        
        # Create pipeline
        pipeline = Pipeline(all_steps)
        
        # Store pipeline and metadata
        pipeline_id = self._generate_pipeline_id(config, model)
        self.pipelines[pipeline_id] = pipeline
        self.current_pipeline = pipeline
        
        # Create metadata
        metadata = self._create_metadata(pipeline_id, config, model)
        self.metadata[pipeline_id] = metadata
        
        return pipeline
    
    def _create_preprocessing_steps(self, config: PipelineConfig) -> List[Tuple[str, Any]]:
        """Create preprocessing steps based on configuration"""
        steps = []
        
        if not config.numeric_features and not config.categorical_features:
            return steps
        
        # Create transformers for different feature types
        transformers = []
        
        # Numeric features preprocessing
        if config.numeric_features:
            numeric_transformer = self._create_numeric_transformer(config)
            transformers.append(('numeric', numeric_transformer, config.numeric_features))
        
        # Categorical features preprocessing
        if config.categorical_features:
            categorical_transformer = self._create_categorical_transformer(config)
            transformers.append(('categorical', categorical_transformer, config.categorical_features))
        
        # Text features preprocessing
        if config.text_features:
            text_transformer = self._create_text_transformer(config)
            transformers.append(('text', text_transformer, config.text_features))
        
        # Datetime features preprocessing
        if config.datetime_features:
            datetime_transformer = self._create_datetime_transformer(config)
            transformers.append(('datetime', datetime_transformer, config.datetime_features))
        
        if transformers:
            preprocessor = ColumnTransformer(transformers, remainder='drop')
            steps.append(('preprocessor', preprocessor))
        
        return steps
    
    def _create_numeric_transformer(self, config: PipelineConfig) -> Pipeline:
        """Create transformer for numeric features"""
        steps = []
        
        # Missing value handling
        if config.handle_missing:
            imputer = SimpleImputer(strategy='median')
            steps.append(('imputer', imputer))
        
        # Scaling
        if config.scaling_method == 'standard':
            scaler = StandardScaler()
            steps.append(('scaler', scaler))
        elif config.scaling_method == 'minmax':
            scaler = MinMaxScaler()
            steps.append(('scaler', scaler))
        elif config.scaling_method == 'robust':
            scaler = RobustScaler()
            steps.append(('scaler', scaler))
        
        return Pipeline(steps)
    
    def _create_categorical_transformer(self, config: PipelineConfig) -> Pipeline:
        """Create transformer for categorical features"""
        steps = []
        
        # Missing value handling
        if config.handle_missing:
            imputer = SimpleImputer(strategy='most_frequent')
            steps.append(('imputer', imputer))
        
        # Encoding
        if config.encoding_method == 'auto':
            # Auto-determine encoding method based on cardinality
            encoding_method = self._determine_encoding_method(config.categorical_features)
        else:
            encoding_method = config.encoding_method
        
        if encoding_method == 'onehot':
            encoder = OneHotEncoder(drop='first', sparse_output=False, handle_unknown='ignore')
            steps.append(('encoder', encoder))
        elif encoding_method == 'ordinal':
            encoder = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1)
            steps.append(('encoder', encoder))
        elif encoding_method == 'target':
            # Target encoding would require target variable
            # For now, fallback to ordinal
            encoder = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1)
            steps.append(('encoder', encoder))
        
        return Pipeline(steps)
    
    def _create_text_transformer(self, config: PipelineConfig) -> Pipeline:
        """Create transformer for text features"""
        steps = []
        
        # Basic text preprocessing
        def text_preprocessor(X):
            # Convert to string and basic cleaning
            if isinstance(X, pd.DataFrame):
                return X.applymap(lambda x: str(x).lower() if pd.notna(x) else '')
            else:
                return np.array([str(x).lower() if x is not None else '' for x in X])
        
        text_transformer = FunctionTransformer(text_preprocessor)
        steps.append(('cleaner', text_transformer))
        
        return Pipeline(steps)
    
    def _create_datetime_transformer(self, config: PipelineConfig) -> Pipeline:
        """Create transformer for datetime features"""
        def datetime_extractor(X):
            """Extract datetime features"""
            if isinstance(X, pd.DataFrame):
                result = pd.DataFrame()
                for col in X.columns:
                    if pd.api.types.is_datetime64_any_dtype(X[col]):
                        result[f'{col}_year'] = X[col].dt.year
                        result[f'{col}_month'] = X[col].dt.month
                        result[f'{col}_day'] = X[col].dt.day
                        result[f'{col}_hour'] = X[col].dt.hour
                        result[f'{col}_dayofweek'] = X[col].dt.dayofweek
                    else:
                        # Try to parse as datetime
                        try:
                            parsed = pd.to_datetime(X[col])
                            result[f'{col}_year'] = parsed.dt.year
                            result[f'{col}_month'] = parsed.dt.month
                            result[f'{col}_day'] = parsed.dt.day
                            result[f'{col}_hour'] = parsed.dt.hour
                            result[f'{col}_dayofweek'] = parsed.dt.dayofweek
                        except:
                            # If parsing fails, just use as is
                            result[col] = X[col]
                return result
            else:
                # Handle numpy array
                return X
        
        datetime_transformer = FunctionTransformer(datetime_extractor)
        return Pipeline([('extractor', datetime_transformer)])
    
    def _create_feature_selection_steps(self, config: PipelineConfig) -> List[Tuple[str, Any]]:
        """Create feature selection steps"""
        steps = []
        
        if not config.feature_selection:
            return steps
        
        if config.feature_selection_method == 'k_best':
            if config.problem_type in ['classification', 'clustering']:
                selector = SelectKBest(score_func=f_classif, k=config.n_features_to_select)
            else:
                selector = SelectKBest(score_func=f_regression, k=config.n_features_to_select)
            steps.append(('feature_selection', selector))
            
        elif config.feature_selection_method == 'percentile':
            if config.problem_type in ['classification', 'clustering']:
                selector = SelectPercentile(score_func=f_classif, percentile=config.n_features_to_select)
            else:
                selector = SelectPercentile(score_func=f_regression, percentile=config.n_features_to_select)
            steps.append(('feature_selection', selector))
            
        elif config.feature_selection_method == 'rfe':
            # RFE requires an estimator
            from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
            
            if config.problem_type in ['classification', 'clustering']:
                estimator = RandomForestClassifier(n_estimators=100, random_state=42)
            else:
                estimator = RandomForestRegressor(n_estimators=100, random_state=42)
            
            selector = RFE(estimator=estimator, n_features_to_select=config.n_features_to_select)
            steps.append(('feature_selection', selector))
        
        return steps
    
    def _create_dimensionality_reduction_steps(self, config: PipelineConfig) -> List[Tuple[str, Any]]:
        """Create dimensionality reduction steps"""
        steps = []
        
        if not config.dimensionality_reduction:
            return steps
        
        if config.dim_reduction_method == 'pca':
            if isinstance(config.n_components, float):
                # Variance ratio
                pca = PCA(n_components=config.n_components, random_state=42)
            else:
                # Fixed number of components
                pca = PCA(n_components=config.n_components, random_state=42)
            steps.append(('pca', pca))
            
        elif config.dim_reduction_method == 'svd':
            svd = TruncatedSVD(n_components=config.n_components, random_state=42)
            steps.append(('svd', svd))
        
        return steps
    
    def _determine_encoding_method(self, categorical_features: List[str]) -> str:
        """Determine best encoding method based on feature characteristics"""
        # This is a simplified version - in practice, you'd analyze the actual data
        if len(categorical_features) > 10:
            return 'ordinal'  # Many categorical features might be high cardinality
        else:
            return 'onehot'  # Default to one-hot for fewer features
    
    def _generate_pipeline_id(self, config: PipelineConfig, model: Any) -> str:
        """Generate unique pipeline ID"""
        config_str = json.dumps(asdict(config), sort_keys=True)
        model_str = str(type(model).__name__)
        combined = f"{config_str}_{model_str}_{datetime.now().isoformat()}"
        return hashlib.md5(combined.encode()).hexdigest()[:16]
    
    def _create_metadata(self, pipeline_id: str, config: PipelineConfig, model: Any) -> PipelineMetadata:
        """Create metadata for the pipeline"""
        return PipelineMetadata(
            pipeline_id=pipeline_id,
            name=config.name,
            version="1.0.0",
            created_at=datetime.now().isoformat(),
            config=asdict(config),
            performance_metrics={},
            feature_names=config.numeric_features + config.categorical_features,
            target_name=config.target_column,
            model_type=type(model).__name__,
            data_hash="",  # Would be computed from actual data
            pipeline_hash=pipeline_id
        )
    
    def fit_pipeline(self, pipeline_id: str, X: pd.DataFrame, y: pd.Series = None) -> Pipeline:
        """
        Fit a pipeline with data
        
        Args:
            pipeline_id: ID of the pipeline to fit
            X: Feature matrix
            y: Target variable
            
        Returns:
            Fitted pipeline
        """
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        pipeline = self.pipelines[pipeline_id]
        
        if y is not None:
            pipeline.fit(X, y)
        else:
            pipeline.fit(X)
        
        return pipeline
    
    def predict(self, pipeline_id: str, X: pd.DataFrame) -> np.ndarray:
        """Make predictions using a fitted pipeline"""
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        pipeline = self.pipelines[pipeline_id]
        return pipeline.predict(X)
    
    def predict_proba(self, pipeline_id: str, X: pd.DataFrame) -> np.ndarray:
        """Make probability predictions using a fitted pipeline"""
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        pipeline = self.pipelines[pipeline_id]
        if hasattr(pipeline, 'predict_proba'):
            return pipeline.predict_proba(X)
        else:
            raise ValueError("Pipeline does not support probability predictions")
    
    def evaluate_pipeline(self, pipeline_id: str, X: pd.DataFrame, y: pd.Series,
                         cv_folds: int = 5) -> Dict[str, float]:
        """Evaluate pipeline performance using cross-validation"""
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        pipeline = self.pipelines[pipeline_id]
        config = self.current_config
        
        # Determine scoring based on problem type
        if config.problem_type in ['classification', 'clustering']:
            scoring = 'accuracy'
        else:
            scoring = 'neg_mean_squared_error'
        
        # Perform cross-validation
        scores = cross_val_score(pipeline, X, y, cv=cv_folds, scoring=scoring)
        
        metrics = {
            f'cv_{scoring}_mean': scores.mean(),
            f'cv_{scoring}_std': scores.std(),
            f'cv_{scoring}_min': scores.min(),
            f'cv_{scoring}_max': scores.max()
        }
        
        # Update metadata
        if pipeline_id in self.metadata:
            self.metadata[pipeline_id].performance_metrics = metrics
        
        return metrics
    
    def save_pipeline(self, pipeline_id: str, filepath: str, save_metadata: bool = True):
        """Save pipeline and optionally metadata to disk"""
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        # Save pipeline
        joblib.dump(self.pipelines[pipeline_id], filepath)
        
        # Save metadata
        if save_metadata and pipeline_id in self.metadata:
            metadata_path = filepath.replace('.pkl', '_metadata.json')
            with open(metadata_path, 'w') as f:
                json.dump(asdict(self.metadata[pipeline_id]), f, indent=2)
    
    def load_pipeline(self, filepath: str, load_metadata: bool = True) -> str:
        """Load pipeline and optionally metadata from disk"""
        # Load pipeline
        pipeline = joblib.load(filepath)
        
        # Generate pipeline ID
        pipeline_id = hashlib.md5(filepath.encode()).hexdigest()[:16]
        self.pipelines[pipeline_id] = pipeline
        
        # Load metadata
        if load_metadata:
            metadata_path = filepath.replace('.pkl', '_metadata.json')
            try:
                with open(metadata_path, 'r') as f:
                    metadata_dict = json.load(f)
                self.metadata[pipeline_id] = PipelineMetadata(**metadata_dict)
            except FileNotFoundError:
                # Create basic metadata if file not found
                self.metadata[pipeline_id] = PipelineMetadata(
                    pipeline_id=pipeline_id,
                    name="Loaded Pipeline",
                    version="1.0.0",
                    created_at=datetime.now().isoformat(),
                    config={},
                    performance_metrics={},
                    feature_names=[],
                    target_name=None,
                    model_type="Unknown",
                    data_hash="",
                    pipeline_hash=pipeline_id
                )
        
        return pipeline_id
    
    def get_pipeline_summary(self, pipeline_id: str) -> Dict[str, Any]:
        """Get comprehensive summary of a pipeline"""
        if pipeline_id not in self.pipelines:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        
        pipeline = self.pipelines[pipeline_id]
        metadata = self.metadata.get(pipeline_id)
        
        summary = {
            'pipeline_id': pipeline_id,
            'steps': [],
            'n_steps': len(pipeline.steps),
            'parameters': pipeline.get_params()
        }
        
        # Add step details
        for name, transformer in pipeline.steps:
            step_info = {
                'name': name,
                'type': type(transformer).__name__,
                'parameters': transformer.get_params() if hasattr(transformer, 'get_params') else {}
            }
            summary['steps'].append(step_info)
        
        # Add metadata if available
        if metadata:
            summary['metadata'] = asdict(metadata)
        
        return summary
    
    def compare_pipelines(self, pipeline_ids: List[str], X: pd.DataFrame, 
                         y: pd.Series) -> Dict[str, Any]:
        """Compare multiple pipelines on the same dataset"""
        comparison = {
            'pipelines': {},
            'best_pipeline': None,
            'ranking': []
        }
        
        results = []
        
        for pipeline_id in pipeline_ids:
            if pipeline_id not in self.pipelines:
                continue
            
            # Evaluate pipeline
            metrics = self.evaluate_pipeline(pipeline_id, X, y)
            
            pipeline_info = {
                'pipeline_id': pipeline_id,
                'metrics': metrics,
                'name': self.metadata[pipeline_id].name if pipeline_id in self.metadata else pipeline_id
            }
            
            comparison['pipelines'][pipeline_id] = pipeline_info
            results.append(pipeline_info)
        
        # Rank pipelines
        if results:
            # Use primary metric for ranking
            primary_metric = 'cv_accuracy_mean' if self.current_config.problem_type in ['classification', 'clustering'] else 'cv_neg_mean_squared_error_mean'
            
            if primary_metric.startswith('cv_neg_'):
                # For negative metrics, higher is better (less negative)
                results.sort(key=lambda x: x['metrics'].get(primary_metric, 0), reverse=True)
            else:
                # For positive metrics, higher is better
                results.sort(key=lambda x: x['metrics'].get(primary_metric, 0), reverse=True)
            
            comparison['ranking'] = results
            comparison['best_pipeline'] = results[0]['pipeline_id']
        
        return comparison
    
    def auto_create_pipeline(self, X: pd.DataFrame, y: pd.Series = None,
                           model: Any = None, name: str = "auto_pipeline") -> str:
        """
        Automatically create a pipeline based on data characteristics
        
        Args:
            X: Feature matrix
            y: Target variable
            model: Model to use (auto-selected if None)
            name: Pipeline name
            
        Returns:
            Pipeline ID
        """
        # Analyze data characteristics
        numeric_features = X.select_dtypes(include=[np.number]).columns.tolist()
        categorical_features = X.select_dtypes(include=['object', 'category']).columns.tolist()
        datetime_features = X.select_dtypes(include=['datetime64']).columns.tolist()
        
        # Determine problem type
        if y is not None:
            if pd.api.types.is_numeric_dtype(y):
                if y.nunique() <= 20:
                    problem_type = 'classification'
                else:
                    problem_type = 'regression'
            else:
                problem_type = 'classification'
        else:
            problem_type = 'clustering'
        
        # Auto-select model if not provided
        if model is None:
            model = self._auto_select_model(problem_type, len(X), len(numeric_features))
        
        # Create configuration
        config = PipelineConfig(
            name=name,
            problem_type=problem_type,
            target_column=y.name if y is not None else None,
            numeric_features=numeric_features,
            categorical_features=categorical_features,
            datetime_features=datetime_features,
            handle_missing=True,
            scaling_method='standard' if numeric_features else 'none',
            encoding_method='auto',
            feature_selection=len(numeric_features) > 20,
            feature_selection_method='k_best',
            n_features_to_select=min(50, len(numeric_features)),
            dimensionality_reduction=len(numeric_features) > 100,
            dim_reduction_method='pca',
            n_components=0.95
        )
        
        # Create pipeline
        pipeline = self.create_pipeline(config, model)
        pipeline_id = list(self.pipelines.keys())[-1]  # Get the last created pipeline ID
        
        return pipeline_id
    
    def _auto_select_model(self, problem_type: str, n_samples: int, n_features: int) -> Any:
        """Auto-select a model based on problem characteristics"""
        if problem_type in ['classification', 'clustering']:
            from sklearn.ensemble import RandomForestClassifier
            return RandomForestClassifier(n_estimators=100, random_state=42)
        else:
            from sklearn.ensemble import RandomForestRegressor
            return RandomForestRegressor(n_estimators=100, random_state=42)
