@浙大疏锦行

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import get_scorer
import joblib
import json
import warnings
warnings.filterwarnings('ignore')


class DataLoader(ABC):
    """抽象数据加载器"""
    @abstractmethod
    def load(self, data_path: str, **kwargs) -> Any:
        pass
    
    @abstractmethod
    def validate(self, data: Any) -> bool:
        pass


class FeatureEngineer(BaseEstimator, TransformerMixin):
    """特征工程处理器"""
    def __init__(self, config: Dict = None):
        self.config = config or {}
        self.feature_names = []
        
    def fit(self, X, y=None):
        # 计算需要学习的特征（如标准化参数、编码映射等）
        if self.config.get('normalize'):
            self.mean_ = X.mean(axis=0) #计算均值
            self.std_ = X.std(axis=0) #计算标准差
        return self
    
    def transform(self, X):
        X_transformed = X.copy()
        
        # 数值特征处理 标准化
        if self.config.get('normalize'):
            X_transformed = (X_transformed - self.mean_) / (self.std_ + 1e-8)
            
        # 特征选择（示例）
        if self.config.get('feature_selection'):
            # 这里可以添加特征选择逻辑
            pass
            
        return X_transformed


class ModelFactory:
    """模型工厂，创建和管理模型"""
    
    _model_registry = {}
    
    @classmethod
    def register_model(cls, name: str, model_class):
        cls._model_registry[name] = model_class
    
    @classmethod
    def create_model(cls, 
                     model_type: str, 
                     task_type: str,
                     **params) -> Any:
        """根据任务类型创建模型"""
        
        model_configs = {
            'classification': {
                'xgboost': 'xgboost.XGBClassifier',
                'lightgbm': 'lightgbm.LGBMClassifier',
                'random_forest': 'sklearn.ensemble.RandomForestClassifier',
                'logistic': 'sklearn.linear_model.LogisticRegression',
                'svm': 'sklearn.svm.SVC'
            },
            'regression': {
                'xgboost': 'xgboost.XGBRegressor',
                'lightgbm': 'lightgbm.LGBMRegressor',
                'random_forest': 'sklearn.ensemble.RandomForestRegressor',
                'linear': 'sklearn.linear_model.LinearRegression',
                'svm': 'sklearn.svm.SVR'
            },
            'clustering': {
                'kmeans': 'sklearn.cluster.KMeans',
                'dbscan': 'sklearn.cluster.DBSCAN',
                'gmm': 'sklearn.mixture.GaussianMixture'
            }
        }
        
        if task_type not in model_configs:
            raise ValueError(f"Unsupported task type: {task_type}")
        
        if model_type not in model_configs[task_type]:
            raise ValueError(f"Unsupported model type: {model_type}")
        
        # 动态导入模型类
        module_path, class_name = model_configs[task_type][model_type].rsplit('.', 1)
        module = __import__(module_path, fromlist=[class_name])
        model_class = getattr(module, class_name)
        
        return model_class(**params)


class ModelSelector:
    """模型选择器，自动选择最佳模型"""
    
    def __init__(self, 
                 task_type: str,
                 metric: str = None,
                 cv: int = 5):
        self.task_type = task_type
        self.metric = metric or self._default_metric()
        self.cv = cv
        self.best_model = None
        self.best_score = -np.inf
        
    def _default_metric(self) -> str:
        """根据任务类型返回默认评估指标"""
        metrics = {
            'classification': 'accuracy',
            'binary_classification': 'roc_auc',
            'regression': 'neg_mean_squared_error',
            'clustering': 'silhouette'
        }
        return metrics.get(self.task_type, 'accuracy')
    
    def select(self, X, y, candidates: List[Dict]) -> Any:
        """从候选模型中选取最佳模型"""
        
        for candidate in candidates:
            model_type = candidate['type']
            params = candidate.get('params', {})
            
            try:
                model = ModelFactory.create_model(
                    model_type, 
                    self.task_type, 
                    **params
                )
                
                # 交叉验证评估
                if self.task_type in ['classification', 'regression']:
                    scores = cross_val_score(
                        model, X, y, 
                        cv=self.cv, 
                        scoring=self.metric
                    )
                    score = scores.mean()
                    
                    if score > self.best_score:
                        self.best_score = score
                        self.best_model = model
                        
            except Exception as e:
                print(f"Model {model_type} failed: {e}")
                continue
        
        if self.best_model is None:
            raise ValueError("No suitable model found")
            
        return self.best_model


class MLPipeline:
    """通用机器学习Pipeline"""
    
    def __init__(self, 
                 config: Dict = None,
                 task_type: str = 'auto'):
        """
        初始化Pipeline
        
        Args:
            config: 配置字典
            task_type: 任务类型 ('auto', 'classification', 'regression', 'clustering')
        """
        self.config = config or {}
        self.task_type = self._detect_task_type(task_type)
        self.data_loader = None
        self.feature_engineer = None
        self.model = None
        self.metrics = {}
        self.history = []
        
    def _detect_task_type(self, task_type: str) -> str:
        """自动检测或验证任务类型"""
        if task_type != 'auto':
            return task_type
            
        # 这里可以添加自动检测逻辑
        return 'classification'  # 默认
    
    def load_data(self, 
                  data_path: str, 
                  loader: Optional[DataLoader] = None,
                  **kwargs) -> Tuple:
        """加载数据"""
        
        if loader is None:
            # 默认使用CSV加载器
            data = pd.read_csv(data_path, **kwargs)
        else:
            data = loader.load(data_path, **kwargs)
        
        # 分离特征和标签
        if 'target' in kwargs:
            target_col = kwargs['target']
            X = data.drop(columns=[target_col])
            y = data[target_col]
        else:
            X = data
            y = None
            
        self._validate_data(X, y)
        
        # 存储原始数据
        self.X_raw = X.copy()
        self.y_raw = y.copy() if y is not None else None
        
        return X, y
    
    def _validate_data(self, X, y):
        """数据验证"""
        if X is None or len(X) == 0:
            raise ValueError("Empty feature data")
        
        # 检查缺失值
        if isinstance(X, pd.DataFrame):
            missing = X.isnull().sum().sum()
            if missing > 0:
                print(f"Warning: Found {missing} missing values")
        
        # 检查数据类型
        if not isinstance(X, (pd.DataFrame, np.ndarray)):
            raise TypeError("X must be DataFrame or ndarray")
    
    def preprocess(self, 
                   X: Any,
                   y: Optional[Any] = None,
                   feature_config: Dict = None) -> Any:
        """数据预处理和特征工程"""
        
        feature_config = feature_config or self.config.get('features', {})
        self.feature_engineer = FeatureEngineer(feature_config)
        
        # 拟合特征处理器（如果需要y）
        if hasattr(self.feature_engineer, 'fit'):
            self.feature_engineer.fit(X, y)
        
        # 转换特征
        X_processed = self.feature_engineer.transform(X)
        
        # 特征名称（如果适用）
        if hasattr(X_processed, 'columns'):
            self.feature_names = X_processed.columns.tolist()
        elif hasattr(self.feature_engineer, 'feature_names'):
            self.feature_names = self.feature_engineer.feature_names
        
        return X_processed
    
    def split_data(self, 
                   X: Any, 
                   y: Any,
                   test_size: float = 0.2,
                   random_state: int = 42,
                   stratify: bool = True) -> Tuple:
        """划分训练集和测试集"""
        
        stratify_y = y if stratify and self.task_type == 'classification' else None
        
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, 
            test_size=test_size,
            random_state=random_state,
            stratify=stratify_y
        )
        
        return X_train, X_test, y_train, y_test
    
    def select_model(self,
                    X: Any,
                    y: Any,
                    model_candidates: List[Dict] = None) -> Any:
        """自动选择模型"""
        
        if model_candidates is None:
            # 默认候选模型
            if self.task_type == 'classification':
                model_candidates = [
                    {'type': 'random_forest', 'params': {'n_estimators': 100}},
                    {'type': 'xgboost', 'params': {'n_estimators': 100}},
                    {'type': 'logistic', 'params': {'max_iter': 1000}}
                ]
            elif self.task_type == 'regression':
                model_candidates = [
                    {'type': 'random_forest', 'params': {'n_estimators': 100}},
                    {'type': 'xgboost', 'params': {'n_estimators': 100}},
                    {'type': 'linear', 'params': {}}
                ]
            else:
                model_candidates = [
                    {'type': 'kmeans', 'params': {'n_clusters': 3}}
                ]
        
        selector = ModelSelector(
            task_type=self.task_type,
            metric=self.config.get('metric'),
            cv=self.config.get('cv_folds', 5)
        )
        
        self.model = selector.select(X, y, model_candidates)
        print(f"Selected model: {type(self.model).__name__}, CV score: {selector.best_score:.4f}")
        
        return self.model
    
    def train(self, 
              X_train: Any,
              y_train: Any,
              model: Optional[Any] = None,
              **kwargs):
        """训练模型"""
        
        if model is not None:
            self.model = model
            
        if self.model is None:
            raise ValueError("No model specified. Use select_model() or provide a model.")
        
        # 训练模型
        self.model.fit(X_train, y_train, **kwargs)
        
        # 记录训练历史
        self.history.append({
            'step': 'train',
            'model': type(self.model).__name__,
            'params': self.model.get_params() if hasattr(self.model, 'get_params') else {}
        })
        
        return self.model
    
    def evaluate(self, 
                 X_test: Any,
                 y_test: Any,
                 metrics: List[str] = None) -> Dict:
        """评估模型"""
        
        if self.model is None:
            raise ValueError("Model not trained yet")
        
        if metrics is None:
            metrics = self._get_default_metrics()
        
        y_pred = self.predict(X_test)
        
        # 计算指标
        self.metrics = {}
        for metric_name in metrics:
            try:
                if metric_name in ['accuracy', 'precision', 'recall', 'f1', 'roc_auc']:
                    from sklearn.metrics import get_scorer
                    scorer = get_scorer(metric_name)
                    score = scorer(self.model, X_test, y_test)
                elif metric_name == 'mse':
                    from sklearn.metrics import mean_squared_error
                    score = mean_squared_error(y_test, y_pred)
                elif metric_name == 'mae':
                    from sklearn.metrics import mean_absolute_error
                    score = mean_absolute_error(y_test, y_pred)
                elif metric_name == 'silhouette':
                    from sklearn.metrics import silhouette_score
                    score = silhouette_score(X_test, y_pred)
                else:
                    continue
                    
                self.metrics[metric_name] = score
                
            except Exception as e:
                print(f"Failed to compute {metric_name}: {e}")
        
        # 打印评估结果
        print("\n" + "="*50)
        print("Model Evaluation Results:")
        print("="*50)
        for metric, value in self.metrics.items():
            print(f"{metric}: {value:.4f}")
        print("="*50)
        
        return self.metrics
    
    def _get_default_metrics(self) -> List[str]:
        """获取默认评估指标"""
        
        if self.task_type in ['classification', 'binary_classification']:
            return ['accuracy', 'precision', 'recall', 'f1', 'roc_auc']
        elif self.task_type == 'regression':
            return ['mse', 'mae', 'r2']
        elif self.task_type == 'clustering':
            return ['silhouette']
        else:
            return ['accuracy']
    
    def predict(self, X: Any) -> Any:
        """预测"""
        
        if self.model is None:
            raise ValueError("Model not trained yet")
        
        # 预处理（如果已定义特征处理器）
        if self.feature_engineer is not None:
            X = self.feature_engineer.transform(X)
        
        return self.model.predict(X)
    
    def predict_proba(self, X: Any) -> Any:
        """预测概率（分类任务）"""
        
        if not hasattr(self.model, 'predict_proba'):
            raise AttributeError("Model does not support probability predictions")
        
        if self.feature_engineer is not None:
            X = self.feature_engineer.transform(X)
        
        return self.model.predict_proba(X)
    
    def save(self, 
             path: str,
             save_data: bool = False,
             save_features: bool = True):
        """保存pipeline"""
        
        save_dict = {
            'model': self.model,
            'feature_engineer': self.feature_engineer,
            'config': self.config,
            'task_type': self.task_type,
            'metrics': self.metrics,
            'feature_names': self.feature_names
        }
        
        if save_data:
            save_dict.update({
                'X_raw': self.X_raw,
                'y_raw': self.y_raw
            })
        
        # 保存模型和pipeline
        joblib.dump(save_dict, f"{path}_pipeline.pkl")
        
        # 保存配置为JSON
        with open(f"{path}_config.json", 'w') as f:
            json.dump({
                'config': self.config,
                'task_type': self.task_type,
                'metrics': self.metrics,
                'feature_names': self.feature_names
            }, f, indent=2)
        
        print(f"Pipeline saved to {path}_pipeline.pkl")
    
    @classmethod
    def load(cls, path: str) -> 'MLPipeline':
        """加载pipeline"""
        
        data = joblib.load(f"{path}_pipeline.pkl")
        
        pipeline = cls(config=data['config'], task_type=data['task_type'])
        pipeline.model = data['model']
        pipeline.feature_engineer = data['feature_engineer']
        pipeline.metrics = data.get('metrics', {})
        pipeline.feature_names = data.get('feature_names', [])
        
        if 'X_raw' in data:
            pipeline.X_raw = data['X_raw']
            pipeline.y_raw = data['y_raw']
        
        return pipeline
    
    def get_feature_importance(self) -> Optional[Dict]:
        """获取特征重要性"""
        
        if hasattr(self.model, 'feature_importances_'):
            importances = self.model.feature_importances_
            return dict(zip(self.feature_names, importances))
        elif hasattr(self.model, 'coef_'):
            coef = self.model.coef_
            return dict(zip(self.feature_names, coef.flatten()))
        
        return None





# 注册自定义组件
class CSVLoader(DataLoader):
    """CSV数据加载器"""
    def load(self, data_path: str, **kwargs):
        return pd.read_csv(data_path, **kwargs)
    
    def validate(self, data: Any) -> bool:
        return isinstance(data, pd.DataFrame) and not data.empty

# ==================== 使用示例 ====================

def example_classification():
    """分类任务示例"""
    from sklearn.datasets import load_iris
    
    # 加载数据
    data = load_iris()
    X, y = data.data, data.target
    
    # 创建pipeline
    config = {
        'features': {
            'normalize': True
        },
        'metric': 'accuracy',
        'cv_folds': 5
    }
    
    pipeline = MLPipeline(config=config, task_type='classification')
    
    # 预处理
    X_processed = pipeline.preprocess(X, y)
    
    # 划分数据
    X_train, X_test, y_train, y_test = pipeline.split_data(X_processed, y)
    
    # 选择模型
    pipeline.select_model(X_train, y_train)
    
    # 训练
    pipeline.train(X_train, y_train)
    
    # 评估
    metrics = pipeline.evaluate(X_test, y_test)
    
    # 预测
    predictions = pipeline.predict(X_test[:5])
    print(f"\nSample predictions: {predictions}")
    
    # 保存
    pipeline.save('iris_model')
    
    return pipeline


def example_regression():
    """回归任务示例"""
    from sklearn.datasets import load_diabetes
    
    data = load_diabetes()
    X, y = data.data, data.target
    
    pipeline = MLPipeline(task_type='regression')
    
    # 完整流程
    X_processed = pipeline.preprocess(X, y)
    X_train, X_test, y_train, y_test = pipeline.split_data(X_processed, y)
    
    # 自定义候选模型
    candidates = [
        {'type': 'xgboost', 'params': {'n_estimators': 200, 'learning_rate': 0.1}},
        {'type': 'random_forest', 'params': {'n_estimators': 200}},
        {'type': 'linear', 'params': {}}
    ]
    
    pipeline.select_model(X_train, y_train, candidates)
    pipeline.train(X_train, y_train)
    pipeline.evaluate(X_test, y_test)
    
    return pipeline