"""
Customer Churn Dataset - Model Training Script
==============================================
Provided by: RSK World
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Contact Page: https://rskworld.in/contact.php
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, AdaBoostClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV, RandomizedSearchCV, StratifiedKFold
from sklearn.metrics import (accuracy_score, precision_score, recall_score, f1_score, 
                            confusion_matrix, classification_report, roc_auc_score, 
                            roc_curve, precision_recall_curve, auc, average_precision_score)
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

# Try importing advanced libraries
try:
    import xgboost as xgb
    XGBOOST_AVAILABLE = True
except ImportError:
    XGBOOST_AVAILABLE = False
    print("Warning: XGBoost not available. Install with: pip install xgboost")

try:
    import lightgbm as lgb
    LIGHTGBM_AVAILABLE = True
except ImportError:
    LIGHTGBM_AVAILABLE = False
    print("Warning: LightGBM not available. Install with: pip install lightgbm")

try:
    import shap
    SHAP_AVAILABLE = True
except ImportError:
    SHAP_AVAILABLE = False
    print("Warning: SHAP not available. Install with: pip install shap")

def load_preprocessed_data(file_path='../data/customer_churn_preprocessed.csv'):
    """
    Load preprocessed dataset
    
    Args:
        file_path (str): Path to preprocessed CSV file
        
    Returns:
        pd.DataFrame: Preprocessed dataset
    """
    try:
        df = pd.read_csv(file_path)
        return df
    except FileNotFoundError:
        print("Preprocessed data not found. Running preprocessing...")
        from data_preprocessing import main as preprocess_main
        X_train, X_test, y_train, y_test = preprocess_main()
        return None, X_train, X_test, y_train, y_test

def prepare_data(df):
    """
    Prepare data for training
    
    Args:
        df (pd.DataFrame): Preprocessed dataset
        
    Returns:
        tuple: X_train, X_test, y_train, y_test
    """
    exclude_cols = ['CustomerID', 'LastLogin', 'ChurnDate', 'Churn']
    feature_cols = [col for col in df.columns if col not in exclude_cols]
    
    X = df[feature_cols]
    y = df['Churn']
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    return X_train, X_test, y_train, y_test

def train_random_forest(X_train, y_train, X_test, y_test):
    """
    Train Random Forest classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        RandomForestClassifier: Trained model
        dict: Evaluation metrics
    """
    print("\nTraining Random Forest Classifier...")
    
    model = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1] if hasattr(model, 'predict_proba') else None
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred)
    }
    
    if y_pred_proba is not None:
        metrics['roc_auc'] = roc_auc_score(y_test, y_pred_proba)
        metrics['pr_auc'] = average_precision_score(y_test, y_pred_proba)
    
    return model, metrics

def train_logistic_regression(X_train, y_train, X_test, y_test):
    """
    Train Logistic Regression classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        LogisticRegression: Trained model
        dict: Evaluation metrics
    """
    print("\nTraining Logistic Regression Classifier...")
    
    model = LogisticRegression(random_state=42, max_iter=1000)
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred),
        'roc_auc': roc_auc_score(y_test, y_pred_proba),
        'pr_auc': average_precision_score(y_test, y_pred_proba)
    }
    
    return model, metrics

def train_gradient_boosting(X_train, y_train, X_test, y_test):
    """
    Train Gradient Boosting classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        GradientBoostingClassifier: Trained model
        dict: Evaluation metrics
    """
    print("\nTraining Gradient Boosting Classifier...")
    
    model = GradientBoostingClassifier(random_state=42, n_estimators=100)
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred),
        'roc_auc': roc_auc_score(y_test, y_pred_proba),
        'pr_auc': average_precision_score(y_test, y_pred_proba)
    }
    
    return model, metrics

def train_xgboost(X_train, y_train, X_test, y_test):
    """
    Train XGBoost classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        XGBClassifier: Trained model
        dict: Evaluation metrics
    """
    if not XGBOOST_AVAILABLE:
        return None, None
    
    print("\nTraining XGBoost Classifier...")
    
    model = xgb.XGBClassifier(random_state=42, n_estimators=100, eval_metric='logloss')
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred),
        'roc_auc': roc_auc_score(y_test, y_pred_proba),
        'pr_auc': average_precision_score(y_test, y_pred_proba)
    }
    
    return model, metrics

def train_lightgbm(X_train, y_train, X_test, y_test):
    """
    Train LightGBM classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        LGBMClassifier: Trained model
        dict: Evaluation metrics
    """
    if not LIGHTGBM_AVAILABLE:
        return None, None
    
    print("\nTraining LightGBM Classifier...")
    
    model = lgb.LGBMClassifier(random_state=42, n_estimators=100, verbose=-1)
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred),
        'roc_auc': roc_auc_score(y_test, y_pred_proba),
        'pr_auc': average_precision_score(y_test, y_pred_proba)
    }
    
    return model, metrics

def train_adaboost(X_train, y_train, X_test, y_test):
    """
    Train AdaBoost classifier
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
        
    Returns:
        AdaBoostClassifier: Trained model
        dict: Evaluation metrics
    """
    print("\nTraining AdaBoost Classifier...")
    
    model = AdaBoostClassifier(random_state=42, n_estimators=100)
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred),
        'roc_auc': roc_auc_score(y_test, y_pred_proba),
        'pr_auc': average_precision_score(y_test, y_pred_proba)
    }
    
    return model, metrics

def evaluate_model(model, X_test, y_test, model_name):
    """
    Evaluate model and print detailed metrics
    
    Args:
        model: Trained model
        X_test, y_test: Test data
        model_name (str): Name of the model
    """
    print(f"\n{'='*60}")
    print(f"Evaluation Results for {model_name}")
    print(f"{'='*60}")
    
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1] if hasattr(model, 'predict_proba') else None
    
    print("\nClassification Report:")
    print(classification_report(y_test, y_pred))
    
    print("\nConfusion Matrix:")
    cm = confusion_matrix(y_test, y_pred)
    print(cm)
    
    # Calculate advanced metrics
    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred)
    }
    
    if y_pred_proba is not None:
        metrics['roc_auc'] = roc_auc_score(y_test, y_pred_proba)
        metrics['pr_auc'] = average_precision_score(y_test, y_pred_proba)
        print(f"\nROC-AUC Score: {metrics['roc_auc']:.4f}")
        print(f"PR-AUC Score: {metrics['pr_auc']:.4f}")
    
    # Visualize confusion matrix
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title(f'Confusion Matrix - {model_name}')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.savefig(f'../output/confusion_matrix_{model_name.lower().replace(" ", "_")}.png')
    plt.close()
    
    return metrics

def plot_roc_curve(models_dict, X_test, y_test, model_names):
    """
    Plot ROC curves for multiple models
    
    Args:
        models_dict (dict): Dictionary of trained models
        X_test, y_test: Test data
        model_names (list): List of model names
    """
    plt.figure(figsize=(10, 8))
    
    for model_name in model_names:
        if model_name in models_dict:
            model = models_dict[model_name]
            if hasattr(model, 'predict_proba'):
                y_pred_proba = model.predict_proba(X_test)[:, 1]
                fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
                roc_auc = auc(fpr, tpr)
                plt.plot(fpr, tpr, label=f'{model_name} (AUC = {roc_auc:.3f})')
    
    plt.plot([0, 1], [0, 1], 'k--', label='Random Classifier')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve Comparison')
    plt.legend(loc="lower right")
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig('../output/roc_curves_comparison.png')
    plt.close()

def plot_pr_curve(models_dict, X_test, y_test, model_names):
    """
    Plot Precision-Recall curves for multiple models
    
    Args:
        models_dict (dict): Dictionary of trained models
        X_test, y_test: Test data
        model_names (list): List of model names
    """
    plt.figure(figsize=(10, 8))
    
    for model_name in model_names:
        if model_name in models_dict:
            model = models_dict[model_name]
            if hasattr(model, 'predict_proba'):
                y_pred_proba = model.predict_proba(X_test)[:, 1]
                precision, recall, _ = precision_recall_curve(y_test, y_pred_proba)
                pr_auc = auc(recall, precision)
                plt.plot(recall, precision, label=f'{model_name} (AUC = {pr_auc:.3f})')
    
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title('Precision-Recall Curve Comparison')
    plt.legend(loc="lower left")
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig('../output/pr_curves_comparison.png')
    plt.close()

def cross_validate_model(model, X, y, cv=5, scoring=['accuracy', 'precision', 'recall', 'f1', 'roc_auc']):
    """
    Perform cross-validation on a model
    
    Args:
        model: Model to validate
        X, y: Full dataset
        cv (int): Number of folds
        scoring (list): List of scoring metrics
        
    Returns:
        dict: Cross-validation results
    """
    print(f"\nPerforming {cv}-fold Cross-Validation...")
    cv_results = {}
    
    for metric in scoring:
        try:
            scores = cross_val_score(model, X, y, cv=cv, scoring=metric)
            cv_results[metric] = {
                'mean': scores.mean(),
                'std': scores.std(),
                'scores': scores
            }
            print(f"{metric}: {scores.mean():.4f} (+/- {scores.std() * 2:.4f})")
        except Exception as e:
            print(f"Warning: Could not compute {metric}: {str(e)}")
    
    return cv_results

def compare_models(models_dict):
    """
    Compare performance of multiple models
    
    Args:
        models_dict (dict): Dictionary of model names and their metrics
    """
    print("\n" + "="*60)
    print("MODEL COMPARISON")
    print("="*60)
    
    comparison_df = pd.DataFrame(models_dict).T
    print("\n" + comparison_df.to_string())
    
    # Visualize comparison
    comparison_df.plot(kind='bar', figsize=(12, 6))
    plt.title('Model Performance Comparison')
    plt.ylabel('Score')
    plt.xlabel('Model')
    plt.xticks(rotation=45)
    plt.legend(loc='best')
    plt.tight_layout()
    plt.savefig('../output/model_comparison.png')
    plt.close()

def feature_importance(model, feature_names, model_name):
    """
    Display and visualize feature importance
    
    Args:
        model: Trained model (must have feature_importances_ attribute)
        feature_names (list): List of feature names
        model_name (str): Name of the model
    """
    if hasattr(model, 'feature_importances_'):
        print(f"\n{'='*60}")
        print(f"Feature Importance - {model_name}")
        print(f"{'='*60}")
        
        importance_df = pd.DataFrame({
            'Feature': feature_names,
            'Importance': model.feature_importances_
        }).sort_values('Importance', ascending=False)
        
        print("\nTop 10 Most Important Features:")
        print(importance_df.head(10).to_string(index=False))
        
        # Visualize feature importance
        plt.figure(figsize=(10, 8))
        top_features = importance_df.head(15)
        plt.barh(range(len(top_features)), top_features['Importance'])
        plt.yticks(range(len(top_features)), top_features['Feature'])
        plt.xlabel('Importance')
        plt.title(f'Top 15 Feature Importance - {model_name}')
        plt.gca().invert_yaxis()
        plt.tight_layout()
        plt.savefig(f'../output/feature_importance_{model_name.lower().replace(" ", "_")}.png')
        plt.close()

def main():
    """
    Main function to run model training pipeline
    """
    print("Customer Churn Dataset - Model Training")
    print("Provided by: RSK World (https://rskworld.in/)")
    print("="*60)
    
    # Load data
    result = load_preprocessed_data()
    
    if isinstance(result, tuple) and result[0] is None:
        # Preprocessing script returned data directly
        X_train, X_test, y_train, y_test = result[1:]
        feature_names = X_train.columns.tolist()
    else:
        # Loaded from file
        df = result
        X_train, X_test, y_train, y_test = prepare_data(df)
        feature_names = X_train.columns.tolist()
    
    # Train multiple models
    models = {}
    trained_models = {}
    
    # Random Forest
    rf_model, rf_metrics = train_random_forest(X_train, y_train, X_test, y_test)
    models['Random Forest'] = rf_metrics
    trained_models['Random Forest'] = rf_model
    evaluate_model(rf_model, X_test, y_test, 'Random Forest')
    feature_importance(rf_model, feature_names, 'Random Forest')
    
    # Logistic Regression
    lr_model, lr_metrics = train_logistic_regression(X_train, y_train, X_test, y_test)
    models['Logistic Regression'] = lr_metrics
    trained_models['Logistic Regression'] = lr_model
    evaluate_model(lr_model, X_test, y_test, 'Logistic Regression')
    
    # Gradient Boosting
    gb_model, gb_metrics = train_gradient_boosting(X_train, y_train, X_test, y_test)
    if gb_metrics:
        models['Gradient Boosting'] = gb_metrics
        trained_models['Gradient Boosting'] = gb_model
        evaluate_model(gb_model, X_test, y_test, 'Gradient Boosting')
        feature_importance(gb_model, feature_names, 'Gradient Boosting')
    
    # XGBoost
    xgb_model, xgb_metrics = train_xgboost(X_train, y_train, X_test, y_test)
    if xgb_metrics:
        models['XGBoost'] = xgb_metrics
        trained_models['XGBoost'] = xgb_model
        evaluate_model(xgb_model, X_test, y_test, 'XGBoost')
        feature_importance(xgb_model, feature_names, 'XGBoost')
    
    # LightGBM
    lgb_model, lgb_metrics = train_lightgbm(X_train, y_train, X_test, y_test)
    if lgb_metrics:
        models['LightGBM'] = lgb_metrics
        trained_models['LightGBM'] = lgb_model
        evaluate_model(lgb_model, X_test, y_test, 'LightGBM')
        feature_importance(lgb_model, feature_names, 'LightGBM')
    
    # AdaBoost
    ada_model, ada_metrics = train_adaboost(X_train, y_train, X_test, y_test)
    if ada_metrics:
        models['AdaBoost'] = ada_metrics
        trained_models['AdaBoost'] = ada_model
        evaluate_model(ada_model, X_test, y_test, 'AdaBoost')
        feature_importance(ada_model, feature_names, 'AdaBoost')
    
    # Plot ROC and PR curves
    print("\n" + "="*60)
    print("GENERATING ADVANCED VISUALIZATIONS")
    print("="*60)
    plot_roc_curve(trained_models, X_test, y_test, list(trained_models.keys()))
    plot_pr_curve(trained_models, X_test, y_test, list(trained_models.keys()))
    
    # Compare models
    compare_models(models)
    
    # Cross-validation on best model
    print("\n" + "="*60)
    print("CROSS-VALIDATION ON BEST MODEL")
    print("="*60)
    best_model_name = max(models, key=lambda x: models[x].get('f1', 0))
    best_model = trained_models[best_model_name]
    
    # Prepare full dataset for CV
    if isinstance(result, tuple) and result[0] is None:
        # Use original data if available
        df_full = pd.concat([X_train, X_test])
        y_full = pd.concat([pd.Series(y_train), pd.Series(y_test)])
    else:
        df = result
        exclude_cols = ['CustomerID', 'LastLogin', 'ChurnDate', 'Churn']
        feature_cols = [col for col in df.columns if col not in exclude_cols]
        df_full = df[feature_cols]
        y_full = df['Churn']
    
    cv_results = cross_validate_model(best_model, df_full, y_full, cv=5)
    
    # Save best model
    joblib.dump(best_model, '../output/best_model.pkl')
    print(f"\nBest model ({best_model_name}) saved to '../output/best_model.pkl'")
    
    # Save model metrics
    metrics_df = pd.DataFrame(models).T
    metrics_df.to_csv('../output/model_metrics.csv')
    print("Model metrics saved to '../output/model_metrics.csv'")
    
    print("\n" + "="*60)
    print("TRAINING COMPLETE")
    print("="*60)
    print("\nFor more information, visit: https://rskworld.in/")
    print("Contact: help@rskworld.in | +91 93305 39277")

if __name__ == "__main__":
    main()