fx-quant-research/tradingbot/models/model_manager.py

"""Model training, feature selection, evaluation and prediction manager."""
import time
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.feature_selection import RFE, SelectFromModel
import lightgbm as lgb
from boruta import BorutaPy
from joblib import Parallel, delayed, dump, load
import optuna
import shap

warnings.filterwarnings("ignore")


class GoldModelManager:
    """
    Manages model training, evaluation, and prediction for 5-minute gold trading.
    Implements various feature selection techniques, training strategies,
    and signal generation approaches.
    """
    def __init__(
        self,
        n_splits=5,
        test_size=0.2,
        feature_selection_method='boruta',
        max_features=25,
        ensemble_models=3,
        volatility_based_models=True,
        n_jobs=-1,
        model_path='models',
        random_state=42
    ):
        """
        Initialize the model manager

        Parameters:
        -----------
        n_splits : int
            Number of splits for time series cross-validation
        test_size : float
            Proportion of data to use for final testing
        feature_selection_method : str
            Method for feature selection ('boruta', 'rfe', 'importance', 'pca')
        max_features : int
            Maximum number of features to select
        ensemble_models : int
            Number of models in the ensemble
        volatility_based_models : bool
            Whether to train separate models for different volatility regimes
        n_jobs : int
            Number of parallel jobs for training
        model_path : str
            Directory to save trained models
        random_state : int
            Random seed for reproducibility
        """
        self.n_splits = n_splits
        self.test_size = test_size
        self.feature_selection_method = feature_selection_method
        self.max_features = max_features
        self.ensemble_models = ensemble_models
        self.volatility_based_models = volatility_based_models
        self.n_jobs = n_jobs
        self.random_state = random_state
        self.model_path = Path(model_path)
        self.model_path.mkdir(exist_ok=True)

        # Initialize components
        self.scaler = StandardScaler()
        self.models = {}
        self.meta_model = None
        self.feature_selector = None
        self.selected_features = None
        self.feature_importance = None
        self.pca = None
    def validate_data(self, X, y):
        """
        Validate data before modeling to ensure all classes are represented
        """
        print("Validating data...")

        # Check for missing values
        missing_count = X.isna().sum().sum()
        if missing_count > 0:
            print(f"Warning: {missing_count} missing values detected in features")

        # Check for infinities - only for numeric columns
        numeric_cols = X.select_dtypes(include=['number']).columns
        if len(numeric_cols) > 0:
            inf_count = np.isinf(X[numeric_cols].values).sum()
            if inf_count > 0:
                print(f"Warning: {inf_count} infinity values detected in features")

        # Check class distribution
        unique_classes = np.unique(y)
        class_counts = {cls: np.sum(y == cls) for cls in unique_classes}

        print("Class distribution:")
        for cls, count in sorted(class_counts.items()):
            print(f"  Class {cls}: {count} samples ({count/len(y):.2%})")

        # Verify all expected classes are present
        expected_classes = set(range(-2, 3))  # -2, -1, 0, 1, 2
        missing_classes = expected_classes - set(unique_classes)

        if missing_classes:
            print(f"Warning: Missing classes in data: {missing_classes}")
            print("Model will be trained with available classes only.")

        # Check for severe imbalance
        min_class_count = min(class_counts.values())
        if min_class_count < 10:
            print(f"Warning: Some classes have very few samples (min: {min_class_count})")

        # Check feature values for numeric columns only
        for col in numeric_cols:
            try:
                col_min = X[col].min()
                col_max = X[col].max()
                col_mean = X[col].mean()
                col_std = X[col].std()

                # Check for suspiciously high values
                if col_max > 1e6 or col_min < -1e6:
                    print(f"Warning: Feature '{col}' has extreme values: min={col_min}, max={col_max}")

                # Check for very low variance
                if col_std < 1e-6:
                    print(f"Warning: Feature '{col}' has very low variance: std={col_std}")
            except Exception as e:
                print(f"Warning: Could not analyze feature '{col}': {str(e)}")

        # For categorical columns, check cardinality
        cat_cols = X.select_dtypes(include=['category', 'object']).columns
        for col in cat_cols:
            try:
                n_unique = X[col].nunique()
                print(f"Categorical feature '{col}' has {n_unique} unique values")
            except Exception as e:
                print(f"Warning: Could not analyze categorical feature '{col}': {str(e)}")

        return True

    def prepare_data(self, data, remove_cols=None):
        """
        Prepare data for modeling by separating features from target
        and removing unwanted columns

        Parameters:
        -----------
        data : pd.DataFrame
            DataFrame with features and target
        remove_cols : list
            List of columns to remove (not used as features)

        Returns:
        --------
        X : pd.DataFrame
            Feature matrix
        y : pd.Series
            Target variable
        """
        if remove_cols is None:
            # Default columns to remove (original data and non-feature columns)
            remove_cols = [
                'open', 'high', 'low', 'close', 'volume',
                'future_return', 'strong_threshold', 'weak_threshold'
            ]

        # Create copy to avoid modifying original
        df = data.copy()

        # Ensure target column exists
        if 'target' not in df.columns:
            raise ValueError("Target column 'target' not found in data")

        # Extract target
        y = df['target']

        # Remove target and other non-feature columns
        X = df.drop(['target'] + remove_cols, axis=1, errors='ignore')

        # Handle remaining non-numeric columns
        for col in X.select_dtypes(include=['object', 'category']).columns:
            X[col] = X[col].astype('category')

        return X, y

    def _create_time_series_splits(self, X, y, train_ratio=0.8, validation_ratio=0.1):
        """
        Create time-series aware data splits for proper backtesting

        Parameters:
        -----------
        X : pd.DataFrame
            Feature matrix
        y : pd.Series
            Target variable
        train_ratio : float
            Ratio of data to use for training
        validation_ratio : float
            Ratio of data to use for validation (from end of train)

        Returns:
        --------
        X_train, X_val, X_test, y_train, y_val, y_test : DataFrames/Series
            Split data components
        """
        # Ensure indices match between X and y
        assert X.index.equals(y.index), "X and y must have the same index"

        # Calculate split points
        n = len(X)
        train_end = int(n * train_ratio)
        val_end = train_end + int(n * validation_ratio)

        # Split data
        X_train = X.iloc[:train_end]
        X_val = X.iloc[train_end:val_end]
        X_test = X.iloc[val_end:]

        y_train = y.iloc[:train_end]
        y_val = y.iloc[train_end:val_end]
        y_test = y.iloc[val_end:]

        return X_train, X_val, X_test, y_train, y_val, y_test

    def _get_optimal_lightgbm_params(self, X_train, y_train, X_val, y_val):
        """
        Optimize LightGBM hyperparameters using Optuna with proper class weight handling
        """
        def objective(trial):
            param = {
                'boosting_type': 'gbdt',
                'objective': 'multiclass',
                'num_class': 5,
                'metric': 'multi_logloss',
                'num_leaves': trial.suggest_int('num_leaves', 10, 100),
                'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.1, log=True),
                'feature_fraction': trial.suggest_float('feature_fraction', 0.5, 1.0),
                'bagging_fraction': trial.suggest_float('bagging_fraction', 0.5, 1.0),
                'bagging_freq': trial.suggest_int('bagging_freq', 1, 10),
                'min_child_samples': trial.suggest_int('min_child_samples', 5, 100),
                'reg_alpha': trial.suggest_float('reg_alpha', 1e-8, 10.0, log=True),
                'reg_lambda': trial.suggest_float('reg_lambda', 1e-8, 10.0, log=True),
                'verbose': -1,
                'random_state': self.random_state
            }

            # Get unique classes in the training data
            unique_classes = np.unique(y_train)

            # Create weights for each class that exists in the data
            class_weights = {}
            for cls in range(-2, 3):  # -2, -1, 0, 1, 2
                if cls in unique_classes:
                    # Calculate actual weight based on class frequency
                    count = np.sum(y_train == cls)
                    class_weights[cls] = 1.0 / max(count, 1)
                else:
                    # For missing classes, assign a moderate weight
                    # (average of existing weights would be another option)
                    class_weights[cls] = 1.0

            # Normalize weights to sum to number of classes
            total_weight = sum(class_weights.values())
            class_weights = {k: (v / total_weight) * 5 for k, v in class_weights.items()}

            # Train model
            model = lgb.LGBMClassifier(**param, class_weight=class_weights)
            model.fit(
                X_train, y_train,
                eval_set=[(X_val, y_val)],
                eval_metric='multi_logloss',
                callbacks=[lgb.early_stopping(50, verbose=False)],
            )

            # Evaluate on validation set
            preds = model.predict(X_val)
            return accuracy_score(y_val, preds)

        # Create and run the study
        study = optuna.create_study(direction='maximize')
        study.optimize(objective, n_trials=50)

        return study.best_params

    def select_features(self, X, y, method=None):
        """
        Select the most important features using the specified method

        Parameters:
        -----------
        X : pd.DataFrame
            Feature matrix
        y : pd.Series
            Target variable
        method : str, optional
            Feature selection method (if None, use the instance's method)

        Returns:
        --------
        selected_features : list
            List of selected feature names
        """
        if method is None:
            method = self.feature_selection_method

        print(f"Selecting features using {method} method...")
        if method == 'boruta':
            # Boruta feature selection
            try:
                # Ensure positive value for n_estimators
                base_model = lgb.LGBMClassifier(
                    n_jobs=self.n_jobs,
                    random_state=self.random_state,
                    n_estimators=100,  # Explicit positive value
                    verbose=-1
                )

                # Properly handle class weights for all possible classes
                unique_classes = np.unique(y)
                class_weights = {}

                for cls in range(-2, 3):  # -2, -1, 0, 1, 2
                    if cls in unique_classes:
                        count = np.sum(y == cls)
                        class_weights[cls] = 1.0 / max(count, 1)
                    else:
                        class_weights[cls] = 1.0  # Default weight for missing classes

                # Normalize weights
                total_weight = sum(class_weights.values())
                class_weights = {k: (v / total_weight) * 5 for k, v in class_weights.items()}

                base_model.set_params(class_weight=class_weights)

                self.feature_selector = BorutaPy(
                    estimator=base_model,
                    n_estimators=100,  # Explicit positive value instead of 'auto'
                    max_iter=100,
                    verbose=0,
                    random_state=self.random_state
                )

                # Convert to numpy arrays
                X_values = X.values.astype(np.float64)  # Ensure float type
                y_values = y.values

                # Check for NaNs and infinities
                if np.isnan(X_values).any() or np.isinf(X_values).any():
                    print("Warning: NaN or Inf values detected in features, cleaning...")
                    X_values = np.nan_to_num(X_values, nan=0, posinf=0, neginf=0)

                # Fit Boruta
                self.feature_selector.fit(X_values, y_values)
                selected_mask = self.feature_selector.support_
                self.selected_features = X.columns[selected_mask].tolist()

            except Exception as e:
                print(f"Boruta feature selection failed: {str(e)}")
                print("Falling back to importance-based feature selection...")
                method = 'importance'  # Fall back to importance-based selection

        if method == 'rfe':
            # Recursive Feature Elimination
            base_model = lgb.LGBMClassifier(
                n_jobs=self.n_jobs,
                random_state=self.random_state
            )

            self.feature_selector = RFE(
                estimator=base_model,
                n_features_to_select=min(self.max_features, X.shape[1]),
                step=1,
                verbose=0
            )

            # Fit RFE
            self.feature_selector.fit(X, y)
            selected_mask = self.feature_selector.support_
            self.selected_features = X.columns[selected_mask].tolist()

        if method == 'importance':
            # Feature importance-based selection
            base_model = lgb.LGBMClassifier(
                n_estimators=100,
                n_jobs=self.n_jobs,
                random_state=self.random_state,
                verbose=-1
            )

            # Properly handle class weights
            unique_classes = np.unique(y)
            class_weights = {}

            for cls in range(-2, 3):  # -2, -1, 0, 1, 2
                if cls in unique_classes:
                    count = np.sum(y == cls)
                    class_weights[cls] = 1.0 / max(count, 1)
                else:
                    class_weights[cls] = 1.0

            # Normalize weights
            total_weight = sum(class_weights.values())
            class_weights = {k: (v / total_weight) * 5 for k, v in class_weights.items()}

            base_model.set_params(class_weight=class_weights)

            # Train model
            base_model.fit(X, y)

            # Get feature importance
            importance = base_model.feature_importances_
            self.feature_importance = pd.DataFrame({
                'feature': X.columns,
                'importance': importance
            }).sort_values('importance', ascending=False)

            # Select top features
            self.selected_features = self.feature_importance['feature'].iloc[:min(self.max_features, X.shape[1])].tolist()

            # Create a SelectFromModel as feature_selector for consistency
            self.feature_selector = SelectFromModel(
                base_model,
                max_features=min(self.max_features, X.shape[1])
            )
            self.feature_selector.fit(X, y)

        if method == 'pca':
            # PCA-based dimensionality reduction
            # Scale the data first
            X_scaled = self.scaler.fit_transform(X)

            # Create and fit PCA
            self.pca = PCA(n_components=min(self.max_features, X.shape[1]))
            self.pca.fit(X_scaled)

            # Get explained variance
            explained_variance = self.pca.explained_variance_ratio_
            cumulative_variance = np.cumsum(explained_variance)

            # Determine number of components for 95% variance
            n_components = np.argmax(cumulative_variance >= 0.95) + 1
            n_components = min(n_components, self.max_features)

            # Create new PCA with optimal components
            self.pca = PCA(n_components=n_components)
            self.pca.fit(X_scaled)

            # PCA doesn't select features but transforms them
            # For API consistency, we'll return feature names
            self.selected_features = X.columns.tolist()

        else:
            # No feature selection, use all features
            self.selected_features = X.columns.tolist()

        print(f"Selected {len(self.selected_features)} features")
        return self.selected_features

    def train_base_models(self, X_train, y_train, X_val, y_val):
        """
        Train the base models for the ensemble with proper class weight handling
        """
        print("Training base models...")
        models = {}

        # Optimize hyperparameters once for shared configuration
        best_params = self._get_optimal_lightgbm_params(X_train, y_train, X_val, y_val)
        print(f"Optimized hyperparameters: {best_params}")

        # Get unique classes in the training data
        unique_classes = np.unique(y_train)
        print(f"Classes present in training data: {unique_classes}")

        # Create weights for each class that exists in the data
        class_weights = {}
        for cls in range(-2, 3):  # -2, -1, 0, 1, 2
            if cls in unique_classes:
                # Calculate actual weight based on class frequency
                count = np.sum(y_train == cls)
                class_weights[cls] = 1.0 / max(count, 1)
                print(f"Class {cls}: {count} samples, weight: {class_weights[cls]:.4f}")
            else:
                # For missing classes, assign a moderate weight
                class_weights[cls] = 1.0
                print(f"Class {cls}: 0 samples (missing), weight: 1.0000")

        # Normalize weights to sum to number of classes
        total_weight = sum(class_weights.values())
        class_weights = {k: (v / total_weight) * 5 for k, v in class_weights.items()}

        print("Normalized class weights:")
        for k, v in class_weights.items():
            print(f"Class {k}: {v:.4f}")

        # Train ensemble models with different seeds and subset of features
        for i in range(self.ensemble_models):
            # Create a unique subset of features for diversity
            if len(self.selected_features) > 10:
                # Select 80-90% of features randomly for each model
                sample_size = int(len(self.selected_features) * (0.8 + 0.1 * np.random.random()))
                features = np.random.choice(self.selected_features, size=sample_size, replace=False)
            else:
                features = self.selected_features

            # Adjust random seed for diversity
            model_seed = self.random_state + i

            # Create model with optimized parameters
            model = lgb.LGBMClassifier(
                **best_params,
                random_state=model_seed,
                n_jobs=self.n_jobs,
                class_weight=class_weights
            )

            # Fit model
            model.fit(
                X_train[features], y_train,
                eval_set=[(X_val[features], y_val)],
                eval_metric='multi_logloss',
                callbacks=[lgb.early_stopping(50, verbose=False)],
            )

            # Store model and its features
            models[f'base_model_{i}'] = {
                'model': model,
                'features': features.tolist() if isinstance(features, np.ndarray) else features
            }

            print(f"  Trained base model {i+1}/{self.ensemble_models}")

        # If using volatility-based models, train separate models for each regime
        if self.volatility_based_models and 'volatility_regime' in X_train.columns:
            volatility_regimes = X_train['volatility_regime'].unique()

            for regime in volatility_regimes:
                # Get data for this regime
                regime_mask = X_train['volatility_regime'] == regime
                if regime_mask.sum() < 1000:  # Skip if too few samples
                    continue

                X_regime = X_train[regime_mask]
                y_regime = y_train[regime_mask]

                # Validation data for this regime
                val_regime_mask = X_val['volatility_regime'] == regime
                X_val_regime = X_val[val_regime_mask]
                y_val_regime = y_val[val_regime_mask]

                if len(X_val_regime) < 100:  # Skip if too few validation samples
                    continue

                # Get unique classes for this regime
                regime_unique_classes = np.unique(y_regime)

                # Create weights for each class that exists in this regime
                regime_class_weights = {}
                for cls in range(-2, 3):  # -2, -1, 0, 1, 2
                    if cls in regime_unique_classes:
                        # Calculate actual weight based on class frequency
                        count = np.sum(y_regime == cls)
                        regime_class_weights[cls] = 1.0 / max(count, 1)
                    else:
                        # For missing classes, assign a default weight
                        regime_class_weights[cls] = 1.0

                # Normalize weights
                regime_total_weight = sum(regime_class_weights.values())
                regime_class_weights = {k: (v / regime_total_weight) * 5 for k, v in regime_class_weights.items()}

                # Create and train model
                model = lgb.LGBMClassifier(
                    **best_params,
                    random_state=self.random_state,
                    n_jobs=self.n_jobs,
                    class_weight=regime_class_weights
                )

                # Fit model
                model.fit(
                    X_regime[self.selected_features], y_regime,
                    eval_set=[(X_val_regime[self.selected_features], y_val_regime)],
                    eval_metric='multi_logloss',
                    callbacks=[lgb.early_stopping(50, verbose=False)],
                )

                # Store model
                models[f'regime_model_{int(regime)}'] = {
                    'model': model,
                    'features': self.selected_features,
                    'regime': regime
                }

                print(f"  Trained model for volatility regime {int(regime)}")

        self.models = models
        return models

    def train_meta_model(self, X_val, y_val, X_test=None, y_test=None):
        """
        Train a meta-model on the predictions of base models

        Parameters:
        -----------
        X_val, y_val : Validation data for base model predictions
        X_test, y_test : Optional test data

        Returns:
        --------
        meta_model : trained meta-model
        """
        print("Training meta-model...")

        try:
            # Generate predictions from all base models
            base_predictions = self._get_ensemble_predictions(X_val)

            # Create meta-features for training
            meta_features = pd.DataFrame(base_predictions)

            # Add volatility regime if available
            if 'volatility_regime' in X_val.columns:
                meta_features['volatility_regime'] = X_val['volatility_regime'].values

            # Check for NaN or inf values
            nan_count = np.isnan(meta_features.values).sum()
            inf_count = np.isinf(meta_features.values).sum()

            if nan_count > 0 or inf_count > 0:
                print(f"Warning: Meta-features contain {nan_count} NaN and {inf_count} inf values")
                print("Replacing with zeros...")
                meta_features = meta_features.fillna(0)
                meta_features = meta_features.replace([np.inf, -np.inf], 0)

            # Get unique classes in validation data
            unique_classes = np.unique(y_val)
            print(f"Classes present in validation data: {unique_classes}")

            # Create and train meta-model
            meta_model = lgb.LGBMClassifier(
                n_estimators=100,
                learning_rate=0.05,
                num_leaves=31,
                random_state=self.random_state,
                n_jobs=self.n_jobs,
                objective='multiclass',
                num_class=5  # Explicitly set to 5 for all classes (-2 to 2)
            )

            # Create weights for each class
            class_weights = {}
            for cls in range(-2, 3):  # -2, -1, 0, 1, 2
                if cls in unique_classes:
                    # Calculate actual weight based on class frequency
                    count = np.sum(y_val == cls)
                    class_weights[cls] = 1.0 / max(count, 1)
                    print(f"Class {cls}: {count} samples, weight: {class_weights[cls]:.4f}")
                else:
                    # For missing classes, assign a default weight
                    class_weights[cls] = 1.0
                    print(f"Class {cls}: 0 samples (missing), weight: 1.0000")

            # Convert class weights to the format expected by LightGBM (0-4 index)
            indexed_weights = {}
            for cls in range(-2, 3):
                idx = cls + 2  # Convert class value to index: -2->0, -1->1, 0->2, 1->3, 2->4
                indexed_weights[idx] = class_weights[cls]

            # Normalize weights
            total_weight = sum(indexed_weights.values())
            indexed_weights = {k: (v / total_weight) * 5 for k, v in indexed_weights.items()}

            print("Adjusted class weights for meta-model:")
            for k, v in indexed_weights.items():
                print(f"Index {k} (Class {k-2}): {v:.4f}")

            # Set class weights
            meta_model.set_params(class_weight=indexed_weights)

            # Fit meta-model
            meta_model.fit(meta_features, y_val)

            # Evaluate meta-model if test data is provided
            if X_test is not None and y_test is not None:
                test_predictions = self._get_ensemble_predictions(X_test)
                test_meta_features = pd.DataFrame(test_predictions)

                if 'volatility_regime' in X_test.columns:
                    test_meta_features['volatility_regime'] = X_test['volatility_regime'].values

                # Replace NaN or inf values
                test_meta_features = test_meta_features.fillna(0)
                test_meta_features = test_meta_features.replace([np.inf, -np.inf], 0)

                y_pred = meta_model.predict(test_meta_features)
                accuracy = accuracy_score(y_test, y_pred)
                report = classification_report(y_test, y_pred)

                print(f"Meta-model test accuracy: {accuracy:.4f}")
                print("Classification report:")
                print(report)

            self.meta_model = meta_model
            return meta_model

        except Exception as e:
            import traceback
            print(f"Error training meta-model: {str(e)}")
            print(f"Detailed traceback: {traceback.format_exc()}")
            return None

    def _get_ensemble_predictions(self, X):
        """
        Get predictions from all base models with improved error handling

        Parameters:
        -----------
        X : pd.DataFrame
            Feature matrix

        Returns:
        --------
        predictions : dict
            Dictionary with predictions from each model
        """
        predictions = {}

        # Get predictions from each base model
        for name, model_info in self.models.items():
            try:
                model = model_info['model']
                features = model_info['features']

                # Make sure we're only using features that exist in X
                available_features = [f for f in features if f in X.columns]
                if len(available_features) != len(features):
                    missing_features = set(features) - set(available_features)
                    print(f"Warning: {len(missing_features)} features missing for model {name}")
                    if len(missing_features) <= 5:
                        print(f"Missing features: {missing_features}")

                # For regime models, only predict on matching regime data
                if 'regime' in model_info and 'volatility_regime' in X.columns:
                    regime = model_info['regime']
                    regime_mask = X['volatility_regime'] == regime

                    # Skip if no data for this regime
                    if not regime_mask.any():
                        print(f"Skipping regime model {name} - no matching data")
                        # Initialize with zeros for this model for all possible classes
                        for i in range(5):  # 5 classes (-2 to 2)
                            predictions[f"{name}_class_{i-2}"] = np.zeros(len(X))
                        continue

                    # Initialize predictions with zeros
                    proba = np.zeros((len(X), 5))  # Always use 5 classes (-2 to 2)

                    # Get predictions for matching regime
                    try:
                        regime_proba = model.predict_proba(X.loc[regime_mask, available_features])

                        # Handle case where model has fewer than 5 classes
                        if regime_proba.shape[1] < 5:
                            temp_proba = np.zeros((regime_proba.shape[0], 5))
                            # Map classes to correct positions in the 5-class array
                            classes = model.classes_ + 2  # Adjust to 0-4 indices
                            for i, cls_idx in enumerate(classes):
                                temp_proba[:, cls_idx] = regime_proba[:, i]
                            regime_proba = temp_proba

                        # Place predictions in the correct indices
                        regime_indices = np.where(regime_mask)[0]
                        proba[regime_indices] = regime_proba

                    except Exception as e:
                        print(f"Error predicting for regime {regime}: {str(e)}")
                        # Keep zeros for this regime
                else:
                    # Regular model prediction on all data
                    try:
                        raw_proba = model.predict_proba(X[available_features])

                        # Initialize with zeros for all 5 classes
                        proba = np.zeros((len(X), 5))

                        # Handle case where model has fewer than 5 classes
                        if raw_proba.shape[1] < 5:
                            # Map classes to correct positions in the 5-class array
                            classes = model.classes_ + 2  # Adjust to 0-4 indices
                            for i, cls_idx in enumerate(classes):
                                proba[:, cls_idx] = raw_proba[:, i]
                        else:
                            proba = raw_proba

                    except Exception as e:
                        print(f"Error predicting with model {name}: {str(e)}")
                        # Keep zeros for this model

                # Store class probabilities
                for i in range(5):  # 5 classes (-2 to 2)
                    predictions[f"{name}_class_{i-2}"] = proba[:, i]

            except Exception as e:
                import traceback
                print(f"Error processing model {name}: {str(e)}")
                print(f"Traceback: {traceback.format_exc()}")

                # Initialize with zeros for all classes
                for i in range(5):  # 5 classes (-2 to 2)
                    predictions[f"{name}_class_{i-2}"] = np.zeros(len(X))

        return predictions

    def fit(self, data, remove_cols=None):
        """
        Complete model training pipeline with improved error handling

        Parameters:
        -----------
        data : pd.DataFrame
            DataFrame with features and target
        remove_cols : list, optional
            List of columns to remove (not used as features)

        Returns:
        --------
        self : for method chaining
        """
        print("Starting model training pipeline...")

        try:
            # Prepare data
            X, y = self.prepare_data(data, remove_cols)

            # Basic data validation
            self.validate_data(X, y)

            # Train/validation/test split
            X_train, X_val, X_test, y_train, y_val, y_test = self._create_time_series_splits(X, y)
            print(f"Data split: train={len(X_train)}, validation={len(X_val)}, test={len(X_test)}")

            # Feature selection
            try:
                selected_features = self.select_features(X_train, y_train)
            except Exception as e:
                print(f"Feature selection error: {str(e)}")
                print("Using all features as fallback")
                self.selected_features = X_train.columns.tolist()

            # Train base models
            try:
                self.train_base_models(X_train, y_train, X_val, y_val)
            except Exception as e:
                import traceback
                print(f"Base model training error: {str(e)}")
                print(f"Traceback: {traceback.format_exc()}")
                raise

            # Train meta-model
            try:
                self.train_meta_model(X_val, y_val, X_test, y_test)
            except Exception as e:
                import traceback
                print(f"Meta-model training error: {str(e)}")
                print(f"Traceback: {traceback.format_exc()}")
                raise

            # Save models
            try:
                self.save_models()
            except Exception as e:
                print(f"Model saving error: {str(e)}")

            # Final evaluation
            try:
                self.evaluate(X_test, y_test)
            except Exception as e:
                print(f"Evaluation error: {str(e)}")

            return self

        except Exception as e:
            import traceback
            print(f"Error in model training pipeline: {str(e)}")
            print(f"Traceback: {traceback.format_exc()}")
            return self

    def predict(self, X):
        """
        Generate predictions using the trained ensemble with improved error handling

        Parameters:
        -----------
        X : pd.DataFrame
            Feature matrix

        Returns:
        --------
        predictions : pd.Series
            Class predictions
        probabilities : pd.DataFrame
            Class probabilities
        """
        try:
            if self.meta_model is None:
                raise ValueError("No trained meta-model available for prediction")

            # Get base model predictions
            base_predictions = self._get_ensemble_predictions(X)
            meta_features = pd.DataFrame(base_predictions)

            # Add volatility regime if available
            if 'volatility_regime' in X.columns:
                meta_features['volatility_regime'] = X['volatility_regime'].values

            # Handle missing values
            meta_features = meta_features.fillna(0)
            meta_features = meta_features.replace([np.inf, -np.inf], 0)

            # Generate predictions
            class_proba = self.meta_model.predict_proba(meta_features)

            # Convert to DataFrame with appropriate class labels
            class_names = [f'class_{i-2}' for i in range(5)]  # class_-2 to class_2
            probabilities = pd.DataFrame(
                class_proba,
                columns=class_names,
                index=X.index
            )

            # Get class predictions
            raw_predictions = self.meta_model.predict(meta_features)
            predictions = pd.Series(
                raw_predictions,
                index=X.index,
                name='prediction'
            )

            return predictions, probabilities

        except Exception as e:
            import traceback
            print(f"Error in prediction: {str(e)}")
            print(f"Traceback: {traceback.format_exc()}")

            # Return empty predictions as fallback
            empty_predictions = pd.Series(index=X.index, name='prediction')
            empty_probabilities = pd.DataFrame(index=X.index)

            return empty_predictions, empty_probabilities

    def evaluate(self, X_test, y_test):
        """
        Evaluate the model on test data

        Parameters:
        -----------
        X_test : pd.DataFrame
            Test feature matrix
        y_test : pd.Series
            Test target

        Returns:
        --------
        metrics : dict
            Dictionary of evaluation metrics
        """
        print("\nEvaluating model performance...")

        try:
            # Generate predictions
            y_pred, y_proba = self.predict(X_test)

            # Check for unique classes in predictions and actual
            print(f"Unique classes in test data: {np.unique(y_test)}")
            print(f"Unique classes in predictions: {np.unique(y_pred)}")

            # Calculate metrics
            accuracy = accuracy_score(y_test, y_pred)

            # Handle case where some classes might be missing
            # Get all possible classes from both actual and predicted
            all_classes = sorted(set(np.unique(y_test)) | set(np.unique(y_pred)))

            # Generate report with all possible classes
            report = classification_report(y_test, y_pred, labels=all_classes, output_dict=True)

            # For confusion matrix, we need to use the same classes
            conf_matrix = confusion_matrix(y_test, y_pred, labels=all_classes)

            # Print results
            print(f"Test accuracy: {accuracy:.4f}")
            print("Classification report:")
            print(classification_report(y_test, y_pred, labels=all_classes))

            # Plot confusion matrix
            plt.figure(figsize=(10, 8))
            class_labels = [f"Class {cls}" for cls in all_classes]
            sns.heatmap(
                conf_matrix,
                annot=True,
                fmt='d',
                cmap='Blues',
                xticklabels=class_labels,
                yticklabels=class_labels
            )
            plt.title('Confusion Matrix')
            plt.xlabel('Predicted')
            plt.ylabel('Actual')
            plt.tight_layout()
            plt.savefig(self.model_path / 'confusion_matrix.png')

            # Create metrics dictionary
            metrics = {
                'accuracy': accuracy,
                'report': report,
                'confusion_matrix': conf_matrix
            }

            return metrics

        except Exception as e:
            import traceback
            print(f"Error in evaluation: {str(e)}")
            print(f"Traceback: {traceback.format_exc()}")

            # Return basic metrics
            return {
                'accuracy': 0.0,
                'report': {},
                'confusion_matrix': np.array([[0]])
            }

    def save_models(self):
        """Save trained models and metadata to disk"""
        # Create model directory if it doesn't exist
        self.model_path.mkdir(exist_ok=True)

        # Save base models
        for name, model_info in self.models.items():
            model = model_info['model']
            features = model_info['features']

            # Save model
            dump(model, self.model_path / f'{name}.joblib')

            # Save feature list
            with open(self.model_path / f'{name}_features.txt', 'w') as f:
                f.write('\n'.join(features))

        # Save meta-model
        if self.meta_model is not None:
            dump(self.meta_model, self.model_path / 'meta_model.joblib')

        # Save feature selector or PCA
        if self.feature_selector is not None:
            dump(self.feature_selector, self.model_path / 'feature_selector.joblib')

        if self.pca is not None:
            dump(self.pca, self.model_path / 'pca.joblib')

        # Save selected features
        if self.selected_features is not None:
            with open(self.model_path / 'selected_features.txt', 'w') as f:
                f.write('\n'.join(self.selected_features))

        # Save feature importance if available
        if self.feature_importance is not None:
            self.feature_importance.to_csv(self.model_path / 'feature_importance.csv', index=False)

        print(f"Models and metadata saved to {self.model_path}")

    def load_models(self):
        """Load trained models and metadata from disk"""
        # Check if model directory exists
        if not self.model_path.exists():
            raise FileNotFoundError(f"Model directory {self.model_path} not found")

        # Load base models
        self.models = {}
        for model_file in self.model_path.glob('base_model_*.joblib'):
            name = model_file.stem
            features_file = self.model_path / f'{name}_features.txt'

            if features_file.exists():
                with open(features_file, 'r') as f:
                    features = f.read().splitlines()

                model = load(model_file)
                self.models[name] = {
                    'model': model,
                    'features': features
                }

        # Load regime models
        for model_file in self.model_path.glob('regime_model_*.joblib'):
            name = model_file.stem
            features_file = self.model_path / f'{name}_features.txt'

            if features_file.exists():
                with open(features_file, 'r') as f:
                    features = f.read().splitlines()

                model = load(model_file)
                regime = int(name.split('_')[-1])
                self.models[name] = {
                    'model': model,
                    'features': features,
                    'regime': regime
                }

        # Load meta-model
        meta_model_file = self.model_path / 'meta_model.joblib'
        if meta_model_file.exists():
            self.meta_model = load(meta_model_file)

        # Load feature selector or PCA
        feature_selector_file = self.model_path / 'feature_selector.joblib'
        if feature_selector_file.exists():
            self.feature_selector = load(feature_selector_file)

        pca_file = self.model_path / 'pca.joblib'
        if pca_file.exists():
            self.pca = load(pca_file)

        # Load selected features
        selected_features_file = self.model_path / 'selected_features.txt'
        if selected_features_file.exists():
            with open(selected_features_file, 'r') as f:
                self.selected_features = f.read().splitlines()

        # Load feature importance
        feature_importance_file = self.model_path / 'feature_importance.csv'
        if feature_importance_file.exists():
            self.feature_importance = pd.read_csv(feature_importance_file)

        print(f"Loaded {len(self.models)} models and metadata from {self.model_path}")