import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, GridSearchCV, cross_val_predict
from sklearn.metrics import roc_auc_score, accuracy_score, confusion_matrix
from tqdm import tqdm

# Load the dataset
# LAB
file_path = '/Users/rachitbisht/Desktop/research/data/BC_Machine_Learning/datasets/LAB_30.xlsx'

# TN
#file_path = '/Users/rachitbisht/Desktop/research/data/BC_Machine_Learning/datasets/TN_30.xlsx'

data = pd.read_excel(file_path)

def perform_grid_search(data, comparison_groups, target_group, feature_names):
    # Adjust binary labeling: '1' for target group (LAB), '0' for comparison groups (TN, Control)
    data['Binary_Label'] = data['Type'].apply(lambda x: 1 if x == target_group else 0 if x in comparison_groups else None)
    
    # Filter data to include only the relevant groups and features
    data_filtered = data[data['Type'].isin(comparison_groups + [target_group])].copy()
    
    # Use only the selected features
    X = data_filtered[feature_names]
    y = data_filtered['Binary_Label']
    
    # Initialize the Random Forest classifier with class_weight='balanced'
    rf_clf = RandomForestClassifier(random_state=42, class_weight='balanced')
    
    param_grid = {
        'n_estimators': [300, 350, 400, 450, 500],
        'min_samples_split': [2, 3, 4],               
        'min_samples_leaf': [1, 2, 3, 4],              
        'max_features': ['sqrt', 'log2'],         
        'bootstrap': [True],                      
        'max_depth': [4, 12, 20, None]            
    }

    # Setup stratified cross-validation
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    
    # Perform grid search with cross-validation
    grid_search = GridSearchCV(estimator=rf_clf, param_grid=param_grid, cv=cv, scoring='roc_auc', n_jobs=-1, return_train_score=False)
    grid_search.fit(X, y)

    # Extract results into a pandas DataFrame
    results = pd.DataFrame(grid_search.cv_results_)
    
    # Add accuracy, AUC scores, FPR, and False Positives for each combination with a progress bar
    all_results = []
    total_combinations = len(grid_search.cv_results_['params'])
    
    for idx, (params, mean_test_score) in enumerate(tqdm(zip(grid_search.cv_results_['params'], grid_search.cv_results_['mean_test_score']), total=total_combinations, desc="Grid Search Progress")):
        # Train the model for each parameter set
        best_rf_clf = RandomForestClassifier(random_state=42, class_weight='balanced', **params)
        best_rf_clf.fit(X, y)
        
        # Get cross-validated predictions for accuracy
        y_pred = cross_val_predict(best_rf_clf, X, y, cv=cv)
        accuracy = accuracy_score(y, y_pred)
        
        # Calculate the confusion matrix
        cm = confusion_matrix(y, y_pred)
        
        # Extract TN, FP, FN, TP from confusion matrix
        tn, fp, fn, tp = cm.ravel()
        
        # Compute False Positive Rate (FPR)
        fpr = fp / (fp + tn)
        
        # Store the result, including false positives for TN
        all_results.append({
            'params': params,
            'AUC Score': mean_test_score,
            'Accuracy': accuracy,
            'FPR (TN)': fpr,  # Minimizing this to optimize for TN
            'False Positives (TN)': fp  # Number of false positives for TN
        })
    
    # Convert the list of results to a DataFrame
    df_results = pd.DataFrame(all_results)
    
    # Save the DataFrame to an Excel file
    output_file = 'grid_search_results_with_fpr_and_fp_for_lab.xlsx'
    df_results.to_excel(output_file, index=False)

    print(f"Results saved to {output_file}")

    return df_results

# List of selected features without numeric suffixes

#TN features 
# selected_features = [
#     "Q06187|BTK_Phospho_STGDPQGVIR_1",
#     "Q06187|BTK_Phospho_STGDPQGVIR_2",
#     "Q99501|GA2L1_Phospho_VSPTTSPRPASPVPGSER_2",
#     "Q9Y3Q8|T22D4_Phospho_SPDPFGAVAAQK",
#     "P78324|SHPS1_Glyco_GTANLSETIR",
#     "P01133|EGF_Glyco_VFLNGSR",
#     "Q6WCQ1|MPRIP_Phospho_SNPDFLK"
# ]

#LAB features 
selected_features = [
    "Q14517|FAT1_Glyco_TGALTVQNTTQLR",
    "O95772|STR3N_Phospho_LLIVQDASER",
    "P50851|LRBA_Phospho_SPVDIVTGGISPVR",
    "Q15942|ZYX_Phospho_SPGAPGPLTLK",
    "O75326|SEM7A_Glyco_NPEAPLNVSR",
    "P01137|TGFB1_Glyco_YSNNSWR",
    "Q6WCQ1|MPRIP_Phospho_SNPDFLK"
]

# Perform the grid search for LAB vs TN and Control using the specified features
df_results = perform_grid_search(data, ['TN', 'Control'], 'LAB', selected_features)