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

This script performs hyperparameter tuning for machine learning models.
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, StratifiedKFold
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
import joblib
import warnings
warnings.filterwarnings('ignore')

# Try importing advanced libraries
try:
    import xgboost as xgb
    XGBOOST_AVAILABLE = True
except ImportError:
    XGBOOST_AVAILABLE = False

try:
    import lightgbm as lgb
    LIGHTGBM_AVAILABLE = True
except ImportError:
    LIGHTGBM_AVAILABLE = False

def load_data(file_path='../data/customer_churn_preprocessed.csv'):
    """
    Load preprocessed dataset
    """
    try:
        df = pd.read_csv(file_path)
        return df
    except FileNotFoundError:
        print("Preprocessed data not found. Please run data_preprocessing.py first.")
        return None

def prepare_data(df, test_size=0.2):
    """
    Prepare data for hyperparameter tuning
    """
    from sklearn.model_selection import train_test_split
    
    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=test_size, random_state=42, stratify=y
    )
    
    return X_train, X_test, y_train, y_test, feature_cols

def tune_random_forest(X_train, y_train, cv=5, method='randomized'):
    """
    Tune Random Forest hyperparameters
    """
    print("\n" + "="*60)
    print("TUNING RANDOM FOREST")
    print("="*60)
    
    param_grid = {
        'n_estimators': [100, 200, 300],
        'max_depth': [10, 20, 30, None],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4],
        'max_features': ['sqrt', 'log2', None]
    }
    
    model = RandomForestClassifier(random_state=42, n_jobs=-1)
    cv_fold = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
    
    if method == 'grid':
        search = GridSearchCV(model, param_grid, cv=cv_fold, scoring='f1', 
                            n_jobs=-1, verbose=1)
    else:
        search = RandomizedSearchCV(model, param_grid, n_iter=50, cv=cv_fold,
                                   scoring='f1', n_jobs=-1, random_state=42, verbose=1)
    
    search.fit(X_train, y_train)
    
    print(f"\nBest parameters: {search.best_params_}")
    print(f"Best cross-validation score: {search.best_score_:.4f}")
    
    return search.best_estimator_, search.best_params_, search.best_score_

def tune_gradient_boosting(X_train, y_train, cv=5, method='randomized'):
    """
    Tune Gradient Boosting hyperparameters
    """
    print("\n" + "="*60)
    print("TUNING GRADIENT BOOSTING")
    print("="*60)
    
    param_grid = {
        'n_estimators': [100, 200, 300],
        'learning_rate': [0.01, 0.1, 0.2],
        'max_depth': [3, 5, 7],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4]
    }
    
    model = GradientBoostingClassifier(random_state=42)
    cv_fold = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
    
    if method == 'grid':
        search = GridSearchCV(model, param_grid, cv=cv_fold, scoring='f1',
                            n_jobs=-1, verbose=1)
    else:
        search = RandomizedSearchCV(model, param_grid, n_iter=50, cv=cv_fold,
                                   scoring='f1', n_jobs=-1, random_state=42, verbose=1)
    
    search.fit(X_train, y_train)
    
    print(f"\nBest parameters: {search.best_params_}")
    print(f"Best cross-validation score: {search.best_score_:.4f}")
    
    return search.best_estimator_, search.best_params_, search.best_score_

def tune_xgboost(X_train, y_train, cv=5, method='randomized'):
    """
    Tune XGBoost hyperparameters
    """
    if not XGBOOST_AVAILABLE:
        print("XGBoost not available. Skipping...")
        return None, None, None
    
    print("\n" + "="*60)
    print("TUNING XGBOOST")
    print("="*60)
    
    param_grid = {
        'n_estimators': [100, 200, 300],
        'learning_rate': [0.01, 0.1, 0.2],
        'max_depth': [3, 5, 7],
        'min_child_weight': [1, 3, 5],
        'subsample': [0.8, 0.9, 1.0],
        'colsample_bytree': [0.8, 0.9, 1.0]
    }
    
    model = xgb.XGBClassifier(random_state=42, eval_metric='logloss', n_jobs=-1)
    cv_fold = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
    
    if method == 'grid':
        search = GridSearchCV(model, param_grid, cv=cv_fold, scoring='f1',
                            n_jobs=-1, verbose=1)
    else:
        search = RandomizedSearchCV(model, param_grid, n_iter=50, cv=cv_fold,
                                   scoring='f1', n_jobs=-1, random_state=42, verbose=1)
    
    search.fit(X_train, y_train)
    
    print(f"\nBest parameters: {search.best_params_}")
    print(f"Best cross-validation score: {search.best_score_:.4f}")
    
    return search.best_estimator_, search.best_params_, search.best_score_

def evaluate_tuned_model(model, X_test, y_test, model_name):
    """
    Evaluate tuned model on test set
    """
    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)
    
    print(f"\nTest set performance for {model_name}:")
    for metric, value in metrics.items():
        print(f"  {metric}: {value:.4f}")
    
    return metrics

def main():
    """
    Main function to run hyperparameter tuning pipeline
    """
    print("Customer Churn Dataset - Hyperparameter Tuning")
    print("Provided by: RSK World (https://rskworld.in/)")
    print("="*60)
    
    # Load data
    df = load_data()
    if df is None:
        return
    
    # Prepare data
    X_train, X_test, y_train, y_test, feature_names = prepare_data(df)
    
    results = {}
    
    # Tune Random Forest
    rf_model, rf_params, rf_cv_score = tune_random_forest(X_train, y_train, method='randomized')
    if rf_model:
        rf_metrics = evaluate_tuned_model(rf_model, X_test, y_test, 'Random Forest')
        results['Random Forest'] = {'params': rf_params, 'cv_score': rf_cv_score, 'test_metrics': rf_metrics}
        joblib.dump(rf_model, '../output/rf_tuned_model.pkl')
    
    # Tune Gradient Boosting
    gb_model, gb_params, gb_cv_score = tune_gradient_boosting(X_train, y_train, method='randomized')
    if gb_model:
        gb_metrics = evaluate_tuned_model(gb_model, X_test, y_test, 'Gradient Boosting')
        results['Gradient Boosting'] = {'params': gb_params, 'cv_score': gb_cv_score, 'test_metrics': gb_metrics}
        joblib.dump(gb_model, '../output/gb_tuned_model.pkl')
    
    # Tune XGBoost
    if XGBOOST_AVAILABLE:
        xgb_model, xgb_params, xgb_cv_score = tune_xgboost(X_train, y_train, method='randomized')
        if xgb_model:
            xgb_metrics = evaluate_tuned_model(xgb_model, X_test, y_test, 'XGBoost')
            results['XGBoost'] = {'params': xgb_params, 'cv_score': xgb_cv_score, 'test_metrics': xgb_metrics}
            joblib.dump(xgb_model, '../output/xgb_tuned_model.pkl')
    
    # Save results
    import json
    with open('../output/hyperparameter_tuning_results.json', 'w') as f:
        # Convert numpy types to Python types for JSON serialization
        def convert_types(obj):
            if isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, dict):
                return {k: convert_types(v) for k, v in obj.items()}
            return obj
        
        json.dump(convert_types(results), f, indent=2)
    
    print("\nHyperparameter tuning results saved to '../output/hyperparameter_tuning_results.json'")
    
    print("\n" + "="*60)
    print("HYPERPARAMETER TUNING COMPLETE")
    print("="*60)
    print("\nFor more information, visit: https://rskworld.in/")
    print("Contact: help@rskworld.in | +91 93305 39277")

if __name__ == "__main__":
    main()