"""
Hyperparameter Optimizer

Advanced hyperparameter optimization using Optuna and other optimization techniques
with AI-powered search space recommendations.
"""

import optuna
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Callable, Tuple
from sklearn.model_selection import cross_val_score, StratifiedKFold, KFold
from sklearn.metrics import accuracy_score, mean_squared_error, r2_score
import xgboost as xgb
import lightgbm as lgb
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.svm import SVC, SVR
from sklearn.linear_model import LogisticRegression, LinearRegression
import warnings


class HyperparameterOptimizer:
    """
    Advanced hyperparameter optimizer using Optuna with intelligent
    search space recommendations based on model and data characteristics.
    """
    
    def __init__(self, n_trials: int = 100, timeout: Optional[int] = None):
        self.n_trials = n_trials
        self.timeout = timeout
        self.best_params = {}
        self.best_score = None
        self.optimization_history = []
        self.study = None
        
    def optimize_model(self, model_type: str, X: pd.DataFrame, y: pd.Series,
                      problem_type: str = 'classification',
                      cv_folds: int = 5, scoring: Optional[str] = None) -> Dict[str, Any]:
        """
        Optimize hyperparameters for a specific model type
        
        Args:
            model_type: Type of model to optimize
            X: Feature matrix
            y: Target variable
            problem_type: 'classification' or 'regression'
            cv_folds: Number of cross-validation folds
            scoring: Scoring metric (auto-selected if None)
            
        Returns:
            Dictionary containing optimization results
        """
        # Set default scoring based on problem type
        if scoring is None:
            scoring = 'accuracy' if problem_type == 'classification' else 'neg_mean_squared_error'
        
        # Create study
        direction = 'maximize' if scoring in ['accuracy', 'f1', 'r2'] else 'minimize'
        self.study = optuna.create_study(direction=direction)
        
        # Define objective function
        def objective(trial):
            # Get hyperparameters for this model type
            params = self._get_hyperparameters(trial, model_type, problem_type)
            
            # Create model with suggested parameters
            model = self._create_model(model_type, params, problem_type)
            
            # Cross-validation
            if problem_type == 'classification':
                cv = StratifiedKFold(n_splits=cv_folds, shuffle=True, random_state=42)
            else:
                cv = KFold(n_splits=cv_folds, shuffle=True, random_state=42)
            
            try:
                scores = cross_val_score(model, X, y, cv=cv, scoring=scoring, n_jobs=-1)
                return scores.mean()
            except Exception as e:
                # Return a bad score if model fails
                return -1.0 if direction == 'maximize' else float('inf')
        
        # Optimize
        self.study.optimize(objective, n_trials=self.n_trials, timeout=self.timeout)
        
        # Store results
        self.best_params = self.study.best_params
        self.best_score = self.study.best_value
        self.optimization_history = self.study.trials
        
        return {
            'best_params': self.best_params,
            'best_score': self.best_score,
            'n_trials': len(self.optimization_history),
            'optimization_history': self.optimization_history
        }
    
    def _get_hyperparameters(self, trial: optuna.Trial, model_type: str, 
                           problem_type: str) -> Dict[str, Any]:
        """Get hyperparameters for a specific model type"""
        
        if model_type.lower() in ['random_forest', 'rf']:
            return {
                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
                'max_depth': trial.suggest_int('max_depth', 3, 20),
                'min_samples_split': trial.suggest_int('min_samples_split', 2, 20),
                'min_samples_leaf': trial.suggest_int('min_samples_leaf', 1, 10),
                'max_features': trial.suggest_categorical('max_features', ['sqrt', 'log2', None])
            }
        
        elif model_type.lower() in ['xgboost', 'xgb']:
            return {
                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
                'max_depth': trial.suggest_int('max_depth', 3, 10),
                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),
                'subsample': trial.suggest_float('subsample', 0.6, 1.0),
                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0),
                'reg_alpha': trial.suggest_float('reg_alpha', 0, 1),
                'reg_lambda': trial.suggest_float('reg_lambda', 0, 1)
            }
        
        elif model_type.lower() in ['lightgbm', 'lgb']:
            return {
                'n_estimators': trial.suggest_int('n_estimators', 50, 500),
                'max_depth': trial.suggest_int('max_depth', 3, 10),
                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True),
                'num_leaves': trial.suggest_int('num_leaves', 20, 150),
                'subsample': trial.suggest_float('subsample', 0.6, 1.0),
                'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0),
                'reg_alpha': trial.suggest_float('reg_alpha', 0, 1),
                'reg_lambda': trial.suggest_float('reg_lambda', 0, 1)
            }
        
        elif model_type.lower() in ['svm']:
            params = {
                'C': trial.suggest_float('C', 0.1, 100, log=True)
            }
            if problem_type == 'classification':
                params['kernel'] = trial.suggest_categorical('kernel', ['rbf', 'poly', 'sigmoid'])
                if params['kernel'] in ['rbf', 'poly', 'sigmoid']:
                    params['gamma'] = trial.suggest_float('gamma', 1e-4, 1, log=True)
            else:
                params['kernel'] = trial.suggest_categorical('kernel', ['rbf', 'linear', 'poly'])
                if params['kernel'] in ['rbf', 'poly']:
                    params['gamma'] = trial.suggest_float('gamma', 1e-4, 1, log=True)
            return params
        
        elif model_type.lower() in ['logistic_regression']:
            return {
                'C': trial.suggest_float('C', 0.01, 100, log=True),
                'penalty': trial.suggest_categorical('penalty', ['l1', 'l2']),
                'solver': trial.suggest_categorical('solver', ['liblinear', 'saga'])
            }
        
        elif model_type.lower() in ['knn']:
            return {
                'n_neighbors': trial.suggest_int('n_neighbors', 1, 50),
                'weights': trial.suggest_categorical('weights', ['uniform', 'distance']),
                'p': trial.suggest_int('p', 1, 2)
            }
        
        else:
            return {}
    
    def _create_model(self, model_type: str, params: Dict[str, Any], 
                     problem_type: str):
        """Create a model instance with given parameters"""
        
        model_type = model_type.lower()
        
        if model_type in ['random_forest', 'rf']:
            if problem_type == 'classification':
                return RandomForestClassifier(**params, random_state=42, n_jobs=-1)
            else:
                return RandomForestRegressor(**params, random_state=42, n_jobs=-1)
        
        elif model_type in ['xgboost', 'xgb']:
            if problem_type == 'classification':
                return xgb.XGBClassifier(**params, random_state=42, eval_metric='logloss')
            else:
                return xgb.XGBRegressor(**params, random_state=42)
        
        elif model_type in ['lightgbm', 'lgb']:
            if problem_type == 'classification':
                return lgb.LGBMClassifier(**params, random_state=42, verbose=-1)
            else:
                return lgb.LGBMRegressor(**params, random_state=42, verbose=-1)
        
        elif model_type in ['svm']:
            if problem_type == 'classification':
                return SVC(**params, random_state=42)
            else:
                return SVR(**params)
        
        elif model_type in ['logistic_regression']:
            # Handle solver-penalty compatibility
            if params['penalty'] == 'l1' and params['solver'] not in ['liblinear', 'saga']:
                params['solver'] = 'liblinear'
            elif params['penalty'] == 'l2' and params['solver'] == 'liblinear' and 'dual' not in params:
                params['dual'] = False
            
            return LogisticRegression(**params, random_state=42, max_iter=1000)
        
        elif model_type in ['knn']:
            from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
            if problem_type == 'classification':
                return KNeighborsClassifier(**params)
            else:
                return KNeighborsRegressor(**params)
        
        else:
            raise ValueError(f"Unsupported model type: {model_type}")
    
    def get_optimization_report(self) -> Dict[str, Any]:
        """Get a comprehensive optimization report"""
        if not self.study:
            return {"error": "No optimization has been performed"}
        
        # Get best trial
        best_trial = self.study.best_trial
        
        # Get parameter importance
        try:
            importance = optuna.importance.get_param_importances(self.study)
        except:
            importance = {}
        
        # Get optimization history
        trials_df = self.study.trials_dataframe()
        
        return {
            'best_params': best_trial.params,
            'best_score': best_trial.value,
            'n_trials': len(self.study.trials),
            'parameter_importance': importance,
            'optimization_history': trials_df.to_dict('records') if not trials_df.empty else [],
            'study_summary': {
                'direction': self.study.direction.name,
                'best_trial_number': best_trial.number,
                'datetime_start': best_trial.datetime_start,
                'datetime_complete': best_trial.datetime_complete
            }
        }
    
    def plot_optimization_history(self):
        """Plot optimization history (requires matplotlib)"""
        try:
            import matplotlib.pyplot as plt
            
            if not self.study:
                print("No optimization study available")
                return
            
            # Plot optimization history
            fig, ax = plt.subplots(1, 2, figsize=(15, 5))
            
            # Trial values over time
            trials = [trial.value for trial in self.study.trials if trial.value is not None]
            ax[0].plot(trials)
            ax[0].set_title('Optimization History')
            ax[0].set_xlabel('Trial')
            ax[0].set_ylabel('Objective Value')
            
            # Best value over time
            best_values = []
            current_best = None
            for trial in self.study.trials:
                if trial.value is not None:
                    if current_best is None or (self.study.direction.name == 'MAXIMIZE' and trial.value > current_best) or \
                       (self.study.direction.name == 'MINIMIZE' and trial.value < current_best):
                        current_best = trial.value
                    best_values.append(current_best)
            
            ax[1].plot(best_values)
            ax[1].set_title('Best Value Over Time')
            ax[1].set_xlabel('Trial')
            ax[1].set_ylabel('Best Value')
            
            plt.tight_layout()
            plt.show()
            
        except ImportError:
            print("Matplotlib is required for plotting. Install with: pip install matplotlib")
    
    def suggest_next_parameters(self) -> Dict[str, Any]:
        """Suggest next set of parameters to try"""
        if not self.study:
            return {"error": "No optimization study available"}
        
        # Create a new trial to get suggestions
        trial = self.study.ask()
        return trial.params
    
    def get_hyperparameter_importance(self) -> Dict[str, float]:
        """Get hyperparameter importance scores"""
        if not self.study:
            return {}
        
        try:
            return optuna.importance.get_param_importances(self.study)
        except:
            return {}
    
    def save_study(self, filepath: str):
        """Save the optimization study to disk"""
        if not self.study:
            raise ValueError("No study to save")
        
        import joblib
        joblib.dump(self.study, filepath)
    
    def load_study(self, filepath: str):
        """Load an optimization study from disk"""
        import joblib
        self.study = joblib.load(filepath)
        self.best_params = self.study.best_params
        self.best_score = self.study.best_value
        self.optimization_history = self.study.trials


class EnsembleOptimizer:
    """
    Optimizer for ensemble methods including voting, stacking, and blending.
    """
    
    def __init__(self):
        self.ensemble_config = {}
        self.ensemble_performance = {}
    
    def optimize_voting_ensemble(self, models: List[Tuple[str, Any]], X: pd.DataFrame, 
                                y: pd.Series, problem_type: str = 'classification',
                                cv_folds: int = 5) -> Dict[str, Any]:
        """Optimize voting ensemble weights"""
        from sklearn.ensemble import VotingClassifier, VotingRegressor
        
        def objective(weights):
            # Normalize weights
            weights = np.array(weights)
            weights = weights / np.sum(weights)
            
            # Create voting ensemble
            if problem_type == 'classification':
                ensemble = VotingClassifier(
                    estimators=[(name, model) for name, model in models],
                    voting='soft',
                    weights=weights
                )
            else:
                ensemble = VotingRegressor(
                    estimators=[(name, model) for name, model in models],
                    weights=weights
                )
            
            # Cross-validation
            if problem_type == 'classification':
                cv = StratifiedKFold(n_splits=cv_folds, shuffle=True, random_state=42)
                scoring = 'accuracy'
            else:
                cv = KFold(n_splits=cv_folds, shuffle=True, random_state=42)
                scoring = 'neg_mean_squared_error'
            
            scores = cross_val_score(ensemble, X, y, cv=cv, scoring=scoring)
            return scores.mean()
        
        # Use simple grid search for weights
        best_score = -np.inf
        best_weights = None
        
        # Generate weight combinations
        n_models = len(models)
        for w1 in np.linspace(0.1, 1.0, 10):
            for w2 in np.linspace(0.1, 1.0, 10):
                if n_models == 2:
                    weights = [w1, w2]
                elif n_models == 3:
                    for w3 in np.linspace(0.1, 1.0, 10):
                        weights = [w1, w2, w3]
                        score = objective(weights)
                        if score > best_score:
                            best_score = score
                            best_weights = weights
                else:
                    weights = [w1, w2] + [1.0] * (n_models - 2)
                    score = objective(weights)
                    if score > best_score:
                        best_score = score
                        best_weights = weights
        
        return {
            'best_weights': best_weights,
            'best_score': best_score,
            'ensemble_type': 'voting'
        }
