import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_score, recall_score
from sklearn.preprocessing import StandardScaler
import xgboost as xgb
import warnings
warnings.filterwarnings('ignore')

class ModelTrainer:
    def __init__(self):
        self.scaler = StandardScaler()
        
    def train_model(self, features, labels, model_type='random_forest'):
        """Train machine learning model for trading signals"""
        try:
            if len(features) == 0 or len(labels) == 0:
                raise ValueError("No data available for training")
            
            # Handle infinite and NaN values
            features = features.replace([np.inf, -np.inf], np.nan)
            features = features.fillna(features.mean())
            
            # Split data
            X_train, X_test, y_train, y_test = train_test_split(
                features, labels, test_size=0.2, random_state=42, stratify=labels
            )
            
            # Scale features
            X_train_scaled = self.scaler.fit_transform(X_train)
            X_test_scaled = self.scaler.transform(X_test)
            
            # Train model
            if model_type == 'random_forest':
                model = self.train_random_forest(X_train_scaled, y_train)
            elif model_type == 'xgboost':
                model = self.train_xgboost(X_train_scaled, y_train)
            else:
                raise ValueError(f"Unsupported model type: {model_type}")
            
            # Evaluate model
            metrics = self.evaluate_model(model, X_test_scaled, y_test)
            
            # Feature importance
            feature_importance = self.get_feature_importance(model, features.columns)
            
            # Create final model pipeline that includes scaling
            final_model = ModelPipeline(self.scaler, model)
            
            return final_model, features.columns.tolist(), metrics
            
        except Exception as e:
            print(f"Error training model: {str(e)}")
            raise e
    
    def train_random_forest(self, X_train, y_train):
        """Train Random Forest model with hyperparameter tuning"""
        try:
            # Define parameter grid for tuning
            param_grid = {
                'n_estimators': [100, 200],
                'max_depth': [10, 15, None],
                'min_samples_split': [2, 5],
                'min_samples_leaf': [1, 2]
            }
            
            # Create base model
            rf = RandomForestClassifier(random_state=42, class_weight='balanced')
            
            # Grid search with cross-validation
            grid_search = GridSearchCV(
                rf, param_grid, cv=5, scoring='accuracy', n_jobs=-1
            )
            
            grid_search.fit(X_train, y_train)
            
            return grid_search.best_estimator_
            
        except Exception as e:
            print(f"Error training Random Forest: {str(e)}")
            # Fallback to basic model
            return RandomForestClassifier(
                n_estimators=100, random_state=42, class_weight='balanced'
            ).fit(X_train, y_train)
    
    def train_xgboost(self, X_train, y_train):
        """Train XGBoost model"""
        try:
            # Convert labels to ensure compatibility
            label_mapping = {-1: 0, 0: 1, 1: 2}
            y_train_mapped = np.array([label_mapping[label] for label in y_train])
            
            model = xgb.XGBClassifier(
                n_estimators=100,
                max_depth=6,
                learning_rate=0.1,
                random_state=42,
                eval_metric='mlogloss'
            )
            
            model.fit(X_train, y_train_mapped)
            
            # Wrap model to handle label mapping
            return XGBoostWrapper(model, label_mapping)
            
        except Exception as e:
            print(f"Error training XGBoost: {str(e)}")
            raise e
    
    def evaluate_model(self, model, X_test, y_test):
        """Evaluate model performance"""
        try:
            y_pred = model.predict(X_test)
            
            # Calculate metrics
            accuracy = accuracy_score(y_test, y_pred)
            precision = precision_score(y_test, y_pred, average='weighted', zero_division=0)
            recall = recall_score(y_test, y_pred, average='weighted', zero_division=0)
            
            metrics = {
                'accuracy': accuracy,
                'precision': precision,
                'recall': recall,
                'classification_report': classification_report(y_test, y_pred, zero_division=0),
                'confusion_matrix': confusion_matrix(y_test, y_pred).tolist()
            }
            
            return metrics
            
        except Exception as e:
            print(f"Error evaluating model: {str(e)}")
            return {'accuracy': 0.0, 'precision': 0.0, 'recall': 0.0}
    
    def get_feature_importance(self, model, feature_names):
        """Get feature importance from trained model"""
        try:
            if hasattr(model, 'feature_importances_'):
                importance = model.feature_importances_
            elif hasattr(model, 'model') and hasattr(model.model, 'feature_importances_'):
                importance = model.model.feature_importances_
            else:
                return {}
            
            feature_importance = dict(zip(feature_names, importance))
            
            # Sort by importance
            sorted_importance = dict(sorted(feature_importance.items(), 
                                          key=lambda x: x[1], reverse=True))
            
            return sorted_importance
            
        except Exception as e:
            print(f"Error getting feature importance: {str(e)}")
            return {}

class ModelPipeline:
    """Pipeline that combines scaling and model prediction"""
    def __init__(self, scaler, model):
        self.scaler = scaler
        self.model = model
    
    def predict(self, X):
        """Make predictions with scaling"""
        X_scaled = self.scaler.transform(X)
        return self.model.predict(X_scaled)
    
    def predict_proba(self, X):
        """Make probability predictions with scaling"""
        X_scaled = self.scaler.transform(X)
        return self.model.predict_proba(X_scaled)

class XGBoostWrapper:
    """Wrapper for XGBoost to handle label mapping"""
    def __init__(self, model, label_mapping):
        self.model = model
        self.label_mapping = label_mapping
        self.reverse_mapping = {v: k for k, v in label_mapping.items()}
    
    def predict(self, X):
        """Predict with label remapping"""
        pred_mapped = self.model.predict(X)
        return np.array([self.reverse_mapping[pred] for pred in pred_mapped])
    
    def predict_proba(self, X):
        """Predict probabilities"""
        return self.model.predict_proba(X)