"""
AutoML Model Selector

Intelligent model selection that automatically detects problem type,
suggests appropriate models, and provides AI-powered recommendations.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.svm import SVC, SVR
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.naive_bayes import GaussianNB
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.cluster import KMeans, DBSCAN
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.model_selection import cross_val_score, train_test_split
import xgboost as xgb
import lightgbm as lgb
import warnings
from abc import ABC, abstractmethod


class BaseModel(ABC):
    """Abstract base class for model wrappers"""
    
    def __init__(self, model, name: str):
        self.model = model
        self.name = name
        self.is_fitted = False
        self.performance_metrics = {}
        
    @abstractmethod
    def fit(self, X: pd.DataFrame, y: pd.Series):
        pass
    
    @abstractmethod
    def predict(self, X: pd.DataFrame):
        pass
    
    def get_performance_metrics(self) -> Dict[str, float]:
        return self.performance_metrics


class ClassificationModel(BaseModel):
    """Wrapper for classification models"""
    
    def fit(self, X: pd.DataFrame, y: pd.Series):
        self.model.fit(X, y)
        self.is_fitted = True
        self.classes_ = self.model.classes_
        
    def predict(self, X: pd.DataFrame):
        if not self.is_fitted:
            raise ValueError("Model must be fitted before prediction")
        return self.model.predict(X)
    
    def predict_proba(self, X: pd.DataFrame):
        if not self.is_fitted:
            raise ValueError("Model must be fitted before prediction")
        if hasattr(self.model, 'predict_proba'):
            return self.model.predict_proba(X)
        else:
            raise AttributeError("Model does not support probability predictions")


class RegressionModel(BaseModel):
    """Wrapper for regression models"""
    
    def fit(self, X: pd.DataFrame, y: pd.Series):
        self.model.fit(X, y)
        self.is_fitted = True
        
    def predict(self, X: pd.DataFrame):
        if not self.is_fitted:
            raise ValueError("Model must be fitted before prediction")
        return self.model.predict(X)


class ClusteringModel(BaseModel):
    """Wrapper for clustering models"""
    
    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
        self.model.fit(X)
        self.is_fitted = True
        if hasattr(self.model, 'labels_'):
            self.labels_ = self.model.labels_
        
    def predict(self, X: pd.DataFrame):
        if not self.is_fitted:
            raise ValueError("Model must be fitted before prediction")
        if hasattr(self.model, 'predict'):
            return self.model.predict(X)
        else:
            raise AttributeError("Model does not support prediction on new data")


class AutoMLSelector:
    """
    Intelligent AutoML model selector that automatically detects problem type,
    suggests appropriate models, and provides AI-powered recommendations.
    """
    
    def __init__(self):
        self.models = {}
        self.best_model = None
        self.problem_type = None
        self.model_rankings = []
        self.feature_importance = {}
        
    def detect_problem_type(self, y: pd.Series) -> str:
        """
        Automatically detect the type of ML problem
        
        Args:
            y: Target variable
            
        Returns:
            Problem type: 'classification', 'regression', or 'clustering'
        """
        if y is None:
            return 'clustering'
        
        if pd.api.types.is_numeric_dtype(y):
            n_unique = y.nunique()
            if n_unique <= 20:
                return 'classification'
            else:
                return 'regression'
        else:
            return 'classification'
    
    def get_candidate_models(self, problem_type: str, X: pd.DataFrame, 
                           y: Optional[pd.Series] = None) -> List[BaseModel]:
        """
        Get candidate models based on problem type and data characteristics
        
        Args:
            problem_type: Type of ML problem
            X: Feature matrix
            y: Target variable (optional for clustering)
            
        Returns:
            List of candidate models
        """
        models = []
        n_samples, n_features = X.shape
        
        if problem_type == 'classification':
            # Always include simple models
            models.append(ClassificationModel(LogisticRegression(max_iter=1000, random_state=42), 
                                            'Logistic Regression'))
            models.append(ClassificationModel(RandomForestClassifier(n_estimators=100, random_state=42), 
                                            'Random Forest'))
            models.append(ClassificationModel(DecisionTreeClassifier(random_state=42), 
                                            'Decision Tree'))
            
            # Add more complex models based on data size
            if n_samples > 1000:
                models.append(ClassificationModel(xgb.XGBClassifier(random_state=42, eval_metric='logloss'), 
                                                'XGBoost'))
                models.append(ClassificationModel(lgb.LGBMClassifier(random_state=42, verbose=-1), 
                                                'LightGBM'))
            
            if n_samples < 10000:  # SVM can be slow on large datasets
                models.append(ClassificationModel(SVC(random_state=42), 'SVM'))
            
            models.append(ClassificationModel(KNeighborsClassifier(), 'KNN'))
            models.append(ClassificationModel(GaussianNB(), 'Naive Bayes'))
            
        elif problem_type == 'regression':
            models.append(RegressionModel(LinearRegression(), 'Linear Regression'))
            models.append(RegressionModel(RandomForestRegressor(n_estimators=100, random_state=42), 
                                        'Random Forest'))
            models.append(RegressionModel(DecisionTreeRegressor(random_state=42), 
                                        'Decision Tree'))
            
            if n_samples > 1000:
                models.append(RegressionModel(xgb.XGBRegressor(random_state=42), 'XGBoost'))
                models.append(RegressionModel(lgb.LGBMRegressor(random_state=42, verbose=-1), 
                                            'LightGBM'))
            
            if n_samples < 10000:
                models.append(RegressionModel(SVR(), 'SVR'))
            
            models.append(RegressionModel(KNeighborsRegressor(), 'KNN'))
            
        elif problem_type == 'clustering':
            models.append(ClusteringModel(KMeans(n_clusters=5, random_state=42), 'K-Means'))
            models.append(ClusteringModel(DBSCAN(eps=0.5), 'DBSCAN'))
        
        return models
    
    def evaluate_models(self, X: pd.DataFrame, y: Optional[pd.Series] = None, 
                       cv_folds: int = 5, test_size: float = 0.2) -> Dict[str, Any]:
        """
        Evaluate multiple models and return performance rankings
        
        Args:
            X: Feature matrix
            y: Target variable
            cv_folds: Number of cross-validation folds
            test_size: Test set size for holdout evaluation
            
        Returns:
            Dictionary containing model evaluations and rankings
        """
        # Detect problem type
        self.problem_type = self.detect_problem_type(y)
        
        # Get candidate models
        candidate_models = self.get_candidate_models(self.problem_type, X, y)
        
        # Split data for evaluation
        if y is not None:
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=test_size, random_state=42, stratify=y if self.problem_type == 'classification' else None
            )
        else:
            X_train, X_test = train_test_split(X, test_size=test_size, random_state=42)
            y_train, y_test = None, None
        
        model_evaluations = []
        
        for model_wrapper in candidate_models:
            try:
                # Fit model
                model_wrapper.fit(X_train, y_train)
                
                # Evaluate based on problem type
                if self.problem_type in ['classification', 'regression']:
                    # Cross-validation scores
                    if self.problem_type == 'classification':
                        cv_scores = cross_val_score(model_wrapper.model, X_train, y_train, 
                                                 cv=cv_folds, scoring='accuracy')
                        test_predictions = model_wrapper.predict(X_test)
                        test_score = accuracy_score(y_test, test_predictions)
                        
                        metrics = {
                            'cv_mean': cv_scores.mean(),
                            'cv_std': cv_scores.std(),
                            'test_accuracy': test_score,
                            'classification_report': classification_report(y_test, test_predictions, output_dict=True)
                        }
                        
                    else:  # regression
                        cv_scores = cross_val_score(model_wrapper.model, X_train, y_train, 
                                                 cv=cv_folds, scoring='neg_mean_squared_error')
                        test_predictions = model_wrapper.predict(X_test)
                        test_mse = mean_squared_error(y_test, test_predictions)
                        test_r2 = r2_score(y_test, test_predictions)
                        
                        metrics = {
                            'cv_mean_mse': -cv_scores.mean(),
                            'cv_std_mse': cv_scores.std(),
                            'test_mse': test_mse,
                            'test_r2': test_r2,
                            'test_mae': mean_absolute_error(y_test, test_predictions)
                        }
                    
                    model_wrapper.performance_metrics = metrics
                    
                    # Feature importance
                    if hasattr(model_wrapper.model, 'feature_importances_'):
                        importance_dict = dict(zip(X.columns, model_wrapper.model.feature_importances_))
                        model_wrapper.feature_importance = importance_dict
                    
                else:  # clustering
                    # For clustering, use silhouette score
                    from sklearn.metrics import silhouette_score
                    try:
                        labels = model_wrapper.model.labels_
                        if len(np.unique(labels)) > 1:
                            silhouette_avg = silhouette_score(X_train, labels)
                            metrics = {'silhouette_score': silhouette_avg}
                        else:
                            metrics = {'silhouette_score': 0}
                    except:
                        metrics = {'silhouette_score': 0}
                    
                    model_wrapper.performance_metrics = metrics
                
                model_evaluations.append({
                    'model': model_wrapper,
                    'name': model_wrapper.name,
                    'metrics': metrics,
                    'status': 'success'
                })
                
            except Exception as e:
                model_evaluations.append({
                    'model': model_wrapper,
                    'name': model_wrapper.name,
                    'metrics': {},
                    'status': 'failed',
                    'error': str(e)
                })
        
        # Rank models
        self.model_rankings = self._rank_models(model_evaluations)
        self.models = {eval['name']: eval for eval in model_evaluations}
        
        # Select best model
        successful_models = [eval for eval in model_evaluations if eval['status'] == 'success']
        if successful_models:
            self.best_model = self.model_rankings[0]['model']
        
        return {
            'problem_type': self.problem_type,
            'model_evaluations': model_evaluations,
            'rankings': self.model_rankings,
            'best_model': self.best_model.name if self.best_model else None
        }
    
    def _rank_models(self, model_evaluations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Rank models based on performance metrics"""
        successful_models = [eval for eval in model_evaluations if eval['status'] == 'success']
        
        if self.problem_type == 'classification':
            # Sort by test accuracy, then by CV mean
            successful_models.sort(key=lambda x: (
                x['metrics'].get('test_accuracy', 0),
                x['metrics'].get('cv_mean', 0)
            ), reverse=True)
            
        elif self.problem_type == 'regression':
            # Sort by R² score (higher is better), then by MSE (lower is better)
            successful_models.sort(key=lambda x: (
                x['metrics'].get('test_r2', 0),
                -x['metrics'].get('test_mse', float('inf'))
            ), reverse=True)
            
        elif self.problem_type == 'clustering':
            # Sort by silhouette score
            successful_models.sort(key=lambda x: (
                x['metrics'].get('silhouette_score', 0)
            ), reverse=True)
        
        return successful_models
    
    def get_model_recommendations(self) -> List[str]:
        """Get AI-powered model recommendations"""
        recommendations = []
        
        if not self.model_rankings:
            return ["Please run evaluate_models() first"]
        
        best_model = self.model_rankings[0]
        
        # Best model recommendation
        recommendations.append(f"Best performing model: {best_model['name']}")
        
        # Problem-specific recommendations
        if self.problem_type == 'classification':
            if best_model['metrics'].get('test_accuracy', 0) < 0.8:
                recommendations.append("Consider feature engineering or ensemble methods to improve accuracy")
            
            # Check for overfitting
            cv_mean = best_model['metrics'].get('cv_mean', 0)
            test_acc = best_model['metrics'].get('test_accuracy', 0)
            if abs(cv_mean - test_acc) > 0.1:
                recommendations.append("Model may be overfitting. Consider regularization or more data")
                
        elif self.problem_type == 'regression':
            if best_model['metrics'].get('test_r2', 0) < 0.7:
                recommendations.append("Consider feature engineering or non-linear models to improve R²")
            
            # Check for overfitting
            cv_mse = best_model['metrics'].get('cv_mean_mse', 0)
            test_mse = best_model['metrics'].get('test_mse', 0)
            if abs(cv_mse - test_mse) / cv_mse > 0.2:
                recommendations.append("Model may be overfitting. Consider regularization")
        
        # Ensemble recommendations
        if len(self.model_rankings) >= 3:
            top_models = self.model_rankings[:3]
            if self.problem_type == 'classification':
                recommendations.append(f"Consider ensemble of: {', '.join([m['name'] for m in top_models])}")
            else:
                recommendations.append(f"Consider stacking or blending: {', '.join([m['name'] for m in top_models])}")
        
        # Feature importance recommendations
        if hasattr(best_model['model'], 'feature_importance'):
            importance = best_model['model'].feature_importance
            if importance:
                top_features = sorted(importance.items(), key=lambda x: x[1], reverse=True)[:5]
                recommendations.append(f"Top important features: {', '.join([f[0] for f in top_features])}")
        
        return recommendations
    
    def predict(self, X: pd.DataFrame) -> np.ndarray:
        """Make predictions using the best model"""
        if not self.best_model:
            raise ValueError("No model has been trained. Please run evaluate_models() first")
        
        return self.best_model.predict(X)
    
    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """Make probability predictions for classification"""
        if not self.best_model:
            raise ValueError("No model has been trained. Please run evaluate_models() first")
        
        if not isinstance(self.best_model, ClassificationModel):
            raise ValueError("Probability predictions are only available for classification models")
        
        return self.best_model.predict_proba(X)
    
    def get_feature_importance(self) -> Dict[str, float]:
        """Get feature importance from the best model"""
        if not self.best_model:
            raise ValueError("No model has been trained. Please run evaluate_models() first")
        
        if hasattr(self.best_model, 'feature_importance'):
            return self.best_model.feature_importance
        elif hasattr(self.best_model.model, 'feature_importances_'):
            # Try to get from underlying model if available
            if hasattr(self, 'X_columns'):
                return dict(zip(self.X_columns, self.best_model.model.feature_importances_))
        
        return {}
    
    def save_model(self, filepath: str):
        """Save the best model to disk"""
        if not self.best_model:
            raise ValueError("No model has been trained. Please run evaluate_models() first")
        
        import joblib
        joblib.dump(self.best_model.model, filepath)
    
    def load_model(self, filepath: str):
        """Load a model from disk"""
        import joblib
        loaded_model = joblib.load(filepath)
        
        # Determine model type based on problem type
        if self.problem_type == 'classification':
            self.best_model = ClassificationModel(loaded_model, 'Loaded Model')
        elif self.problem_type == 'regression':
            self.best_model = RegressionModel(loaded_model, 'Loaded Model')
        else:
            self.best_model = ClusteringModel(loaded_model, 'Loaded Model')
        
        self.best_model.is_fitted = True
