src/ml_pipeline/optimization/hyperparameter_tuner.py

import traceback
import numpy as np
import pandas as pd
import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime

from src.ml_pipeline.core.model_factory import ModelFactory
from ..config.ml_config import MLConfig
from ..core.base_model import BaseMLModel
from ..core.resource_monitor import ResourceMonitor

from sklearn.model_selection import (
    GridSearchCV, RandomizedSearchCV,
    TimeSeriesSplit, KFold, StratifiedKFold,
    GroupKFold, LeaveOneOut, ShuffleSplit,
    cross_val_score
)

# Optuna for Bayesian optimisation
try:
    import optuna
    from optuna.samplers import TPESampler
    from optuna.pruners import MedianPruner
    OPTUNA_AVAILABLE = True
except ImportError:
    OPTUNA_AVAILABLE = False

logger = logging.getLogger(__name__)


class HyperparameterOptimizer:
    """Hyperparameter optimisation using various search methods."""

    def __init__(self, config: MLConfig):
        self.config = config
        self.optimization_results = {}
        self.best_params_history = {}
        self.studies = {}
        self.cache_dir = f"{config.cache_dir}/optimization"
        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
        self.resource_monitor = ResourceMonitor()

    def create_cv_strategy(self, n_splits: int = None):
        """Create cross-validation strategy based on configuration."""
        if n_splits is None:
            n_splits = min(5, self.config.cv_folds)

        if self.config.cv_method == 'timeseries':
            return TimeSeriesSplit(n_splits=n_splits)
        elif self.config.cv_method == 'stratified':
            if self.config.stratify:
                y_temp = None
                if hasattr(self, 'y_train'):
                    y_temp = self.y_train

                if y_temp is not None and len(y_temp) > 0:
                    y_bins = pd.qcut(y_temp, q=min(5, len(y_temp) // 10), labels=False, duplicates='drop')
                    return StratifiedKFold(n_splits=n_splits, shuffle=True,
                                          random_state=self.config.random_state)

        return KFold(n_splits=n_splits, shuffle=True,
                    random_state=self.config.random_state)

    def validate_parameters(self, model_name: str, params: Dict):
        """Validate and correct model parameters."""
        valid_params = {}

        try:
            if model_name == 'random_forest':
                if params.get('bootstrap') == False and 'max_samples' in params:
                    if params['max_samples'] is not None:
                        logger.warning("max_samples ignored when bootstrap=False")
                        params['max_samples'] = None

            elif model_name == 'hist_gradient_boosting':
                if 'max_bins' in params:
                    max_bins = params['max_bins']
                    if max_bins is not None and (max_bins < 2 or max_bins > 255):
                        params['max_bins'] = min(max(2, max_bins), 255)
                        logger.warning(f"max_bins adjusted to {params['max_bins']}")

            valid_params = params

        except Exception as e:
            logger.debug(f"Parameter validation for {model_name}: {e}")

        return valid_params

    def grid_search(self, model, X_train, y_train, param_grid: Dict, cv=None):
        """Perform grid search hyperparameter optimisation."""
        logger.info(f"Starting Grid Search for {model.model_name}...")

        self.resource_monitor.start()

        if cv is None:
            cv = self.create_cv_strategy()

        start_time = datetime.now()

        try:
            grid_search = GridSearchCV(
                estimator=model.model,
                param_grid=param_grid,
                cv=cv,
                scoring='neg_mean_squared_error',
                n_jobs=min(2, self.config.n_jobs),
                verbose=1 if self.config.verbose else 0,
                refit=True,
                error_score='raise',
                return_train_score=True
            )

            self.resource_monitor.snapshot('grid_search_start')
            grid_search.fit(X_train, y_train)
            self.resource_monitor.snapshot('grid_search_complete')

            results = {
                'best_params': grid_search.best_params_,
                'best_score': grid_search.best_score_,
                'best_estimator': grid_search.best_estimator_,
                'cv_results': grid_search.cv_results_,
                'scorer': grid_search.scorer_,
                'n_splits': grid_search.n_splits_,
                'refit_time': grid_search.refit_time_,
                'rank': grid_search.cv_results_['rank_test_score'].tolist(),
                'mean_test_score': grid_search.cv_results_['mean_test_score'].tolist(),
                'std_test_score': grid_search.cv_results_['std_test_score'].tolist(),
                'method': 'grid',
                'resource_usage': self.resource_monitor.get_summary()
            }

            model.model = grid_search.best_estimator_
            model.best_params = grid_search.best_params_

            execution_time = (datetime.now() - start_time).total_seconds()

            logger.info(f"✓ Grid Search completed in {execution_time:.2f} sec")
            logger.info(f"  Best parameters: {grid_search.best_params_}")
            logger.info(f"  Best score: {-grid_search.best_score_:.4f}")

            self.optimization_results[model.model_name] = results
            self.best_params_history[model.model_name] = grid_search.best_params_

            self._save_optimization_results(model.model_name, results)

            return results

        except Exception as e:
            logger.error(f"✗ Grid Search failed: {e}")
            traceback.print_exc()
            return {}

    def random_search(self, model, X_train, y_train, param_distributions: Dict,
                     n_iter: int = 20, cv=None):
        """Perform random search hyperparameter optimisation."""
        logger.info(f"Starting Random Search for {model.model_name}...")

        self.resource_monitor.start()

        if cv is None:
            cv = self.create_cv_strategy()

        start_time = datetime.now()

        try:
            random_search = RandomizedSearchCV(
                estimator=model.model,
                param_distributions=param_distributions,
                n_iter=n_iter,
                cv=cv,
                scoring='neg_mean_squared_error',
                n_jobs=min(2, self.config.n_jobs),
                verbose=1 if self.config.verbose else 0,
                random_state=self.config.random_state,
                refit=True,
                error_score='raise',
                return_train_score=True
            )

            self.resource_monitor.snapshot('random_search_start')
            random_search.fit(X_train, y_train)
            self.resource_monitor.snapshot('random_search_complete')

            results = {
                'best_params': random_search.best_params_,
                'best_score': random_search.best_score_,
                'best_estimator': random_search.best_estimator_,
                'cv_results': random_search.cv_results_,
                'scorer': random_search.scorer_,
                'n_splits': random_search.n_splits_,
                'refit_time': random_search.refit_time_,
                'rank': random_search.cv_results_['rank_test_score'].tolist(),
                'mean_test_score': random_search.cv_results_['mean_test_score'].tolist(),
                'std_test_score': random_search.cv_results_['std_test_score'].tolist(),
                'method': 'random',
                'resource_usage': self.resource_monitor.get_summary()
            }

            model.model = random_search.best_estimator_
            model.best_params = random_search.best_params_

            execution_time = (datetime.now() - start_time).total_seconds()

            logger.info(f"✓ Random Search completed in {execution_time:.2f} sec")
            logger.info(f"  Best parameters: {random_search.best_params_}")
            logger.info(f"  Best score: {-random_search.best_score_:.4f}")

            self.optimization_results[model.model_name] = results
            self.best_params_history[model.model_name] = random_search.best_params_

            self._save_optimization_results(model.model_name, results)

            return results

        except Exception as e:
            logger.error(f"✗ Random Search failed: {e}")
            traceback.print_exc()
            return {}

    def bayesian_optimization(self, model, X_train, y_train, param_space: Dict = None,
                             n_trials: int = None, cv=None):
        """Perform Bayesian optimisation using Optuna."""
        logger.info(f"Starting Bayesian Optimisation for {model.model_name}...")

        if not OPTUNA_AVAILABLE:
            logger.error("Optuna not installed. Install with: pip install optuna")
            return {}

        self.resource_monitor.start()

        if n_trials is None:
            n_trials = self.config.n_trials

        if cv is None:
            cv = self.create_cv_strategy()

        if param_space is None:
            param_space = self.config.param_search_space.get(model.model_name, {})

        start_time = datetime.now()

        try:
            def objective(trial):
                params = {}

                for param_name, param_config in param_space.items():
                    if isinstance(param_config, list):
                        params[param_name] = trial.suggest_categorical(param_name, param_config)
                    elif isinstance(param_config, tuple) and len(param_config) == 2:
                        if all(isinstance(x, int) for x in param_config):
                            params[param_name] = trial.suggest_int(
                                param_name, param_config[0], param_config[1]
                            )
                        else:
                            params[param_name] = trial.suggest_float(
                                param_name, param_config[0], param_config[1], log=True
                            )
                    elif isinstance(param_config, dict):
                        if param_config.get('type') == 'categorical':
                            params[param_name] = trial.suggest_categorical(
                                param_name, param_config['values']
                            )
                        elif param_config.get('type') == 'int':
                            params[param_name] = trial.suggest_int(
                                param_name,
                                param_config['low'],
                                param_config['high'],
                                log=param_config.get('log', False),
                                step=param_config.get('step', 1)
                            )
                        elif param_config.get('type') == 'float':
                            params[param_name] = trial.suggest_float(
                                param_name,
                                param_config['low'],
                                param_config['high'],
                                log=param_config.get('log', False)
                            )

                params = self.validate_parameters(model.model_name, params)

                trial_model = ModelFactory.create_model(model.model_name, self.config)
                trial_model.create_model(**params)

                try:
                    cv_scores = cross_val_score(
                        trial_model.model, X_train, y_train,
                        cv=cv,
                        scoring='neg_mean_squared_error',
                        n_jobs=min(2, self.config.n_jobs),
                        error_score='raise'
                    )

                    return cv_scores.mean()
                except Exception as e:
                    raise optuna.TrialPruned(f"Validation error: {str(e)}")

            study_name = f"{model.model_name}_bayesian"
            storage_path = f"sqlite:///{self.cache_dir}/{study_name}.db"

            sampler = TPESampler(seed=self.config.random_state)
            pruner = MedianPruner(n_startup_trials=10, n_warmup_steps=5)

            study = optuna.create_study(
                study_name=study_name,
                direction='maximize',
                sampler=sampler,
                pruner=pruner,
                storage=storage_path if self.config.use_cache else None,
                load_if_exists=True
            )

            self.studies[model.model_name] = study

            self.resource_monitor.snapshot('bayesian_start')
            study.optimize(
                objective,
                n_trials=n_trials,
                timeout=self.config.timeout_per_trial,
                show_progress_bar=True,
                gc_after_trial=True
            )
            self.resource_monitor.snapshot('bayesian_complete')

            best_params = study.best_params
            best_value = study.best_value

            best_params = self.validate_parameters(model.model_name, best_params)

            model.create_model(**best_params)
            model.fit(X_train, y_train)
            model.best_params = best_params

            execution_time = (datetime.now() - start_time).total_seconds()

            logger.info(f"✓ Bayesian Optimisation completed in {execution_time:.2f} sec")
            logger.info(f"  Best parameters: {best_params}")
            logger.info(f"  Best score: {-best_value:.4f}")

            results = {
                'best_params': best_params,
                'best_score': best_value,
                'study': study,
                'trials_dataframe': study.trials_dataframe(),
                'best_trial': study.best_trial,
                'n_trials': len(study.trials),
                'method': 'bayesian',
                'resource_usage': self.resource_monitor.get_summary()
            }

            self.optimization_results[model.model_name] = results
            self.best_params_history[model.model_name] = best_params

            if self.config.save_plots:
                self._plot_optuna_results(model.model_name, study)

            self._save_optimization_results(model.model_name, results)

            return results

        except Exception as e:
            logger.error(f"✗ Bayesian Optimisation failed: {e}")
            traceback.print_exc()
            return {}

    def _plot_optuna_results(self, model_name: str, study):
        """Generate and save visualisations of Optuna optimisation results."""
        if study is None:
            return

        try:
            from optuna.visualization import plot_optimization_history, plot_param_importances

            plot_dir = f"{self.config.ml_results_dir}/plots/hyperparameters"
            Path(plot_dir).mkdir(parents=True, exist_ok=True)

            fig1 = plot_optimization_history(study)
            fig1.update_layout(title=f"Optimisation History: {model_name}")
            fig1.write_html(f"{plot_dir}/{model_name}_optuna_history.html")

            try:
                import kaleido
                fig1.write_image(f"{plot_dir}/{model_name}_optuna_history.png",
                               width=1200, height=800)
            except:
                logger.debug("kaleido not installed, skipping PNG export")

            try:
                fig2 = plot_param_importances(study)
                fig2.update_layout(title=f"Parameter Importance: {model_name}")
                fig2.write_html(f"{plot_dir}/{model_name}_optuna_importance.html")

                try:
                    fig2.write_image(f"{plot_dir}/{model_name}_optuna_importance.png",
                                   width=1200, height=800)
                except:
                    pass
            except:
                pass

            logger.info(f"✓ Optuna visualisations saved for {model_name}")

        except Exception as e:
            logger.warning(f"Failed to save Optuna visualisations: {e}")

    def _save_optimization_results(self, model_name: str, results: Dict):
        """Cache optimisation results to disk."""
        if not self.config.use_cache:
            return

        try:
            results_to_save = results.copy()
            for key in ['study', 'best_estimator', 'scorer']:
                if key in results_to_save:
                    del results_to_save[key]

            cache_path = f"{self.cache_dir}/{model_name}_optimization.json"
            with open(cache_path, 'w', encoding='utf-8') as f:
                json.dump(results_to_save, f, indent=4, ensure_ascii=False)

        except Exception as e:
            logger.debug(f"Failed to save optimisation cache: {e}")

    def optimize_model(self, model, X_train, y_train, method: str = None,
                      param_space: Dict = None):
        """Perform hyperparameter optimisation using specified method."""
        if not self.config.enable_hyperparameter_tuning:
            return {}

        if method is None:
            method = self.config.tuning_method

        if param_space is None:
            param_space = self.config.param_search_space.get(model.model_name, {})

        if not param_space:
            logger.warning(f"No parameter space defined for {model.model_name}")
            return {}

        if method == 'grid':
            return self.grid_search(model, X_train, y_train, param_space)
        elif method == 'random':
            return self.random_search(model, X_train, y_train, param_space)
        elif method == 'bayesian':
            return self.bayesian_optimization(model, X_train, y_train, param_space)
        else:
            logger.warning(f"Unknown optimisation method: {method}")
            return {}

    def get_optimization_summary(self) -> pd.DataFrame:
        """Generate summary DataFrame of all optimisation results."""
        summary_data = []

        for model_name, results in self.optimization_results.items():
            if results:
                summary_data.append({
                    'model': model_name,
                    'method': results.get('method', 'unknown'),
                    'best_score': -results.get('best_score', 0),
                    'n_trials': results.get('n_trials', 0),
                    'training_time': results.get('resource_usage', {}).get('duration_sec', 0),
                    'best_params': json.dumps(results.get('best_params', {}))
                })

        return pd.DataFrame(summary_data)