import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from itertools import combinations
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.metrics import roc_auc_score, roc_curve, accuracy_score
import matplotlib.pyplot as plt
import numpy as np

# Load the dataset LAB
file_path = '../datasets/LAB_30.xlsx'

# Load the dataset TN 
file_path = '../datasets/TN_30.xlsx'


data = pd.read_excel(file_path)

def exhaustive_feature_selection(data, comparison_groups, target_group, n_features_to_select):
    # Adjust binary labeling: '1' for target group, '0' for comparison groups
    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
    data_filtered = data[data['Type'].isin(comparison_groups + [target_group])].copy()
    
    # Prepare the data
    X = data_filtered.drop(['Type', 'Binary_Label'], axis=1)
    y = data_filtered['Binary_Label']
    
    # Initialize RandomForest classifier with class_weight='balanced'
    rf_clf = RandomForestClassifier(n_estimators=100, random_state=42, class_weight='balanced')
    
    # Initialize variables to store the best combination and its AUC score
    best_auc_score = 0
    best_accuracy = 0
    best_features = None
    best_fpr, best_tpr = None, None
    idx = 0

    # Loop over all combinations of features
    for feature_combo in combinations(X.columns, n_features_to_select):
        idx+=1

        if (idx % 100 == 0):
            print(idx)

        # Select the current combination of features
        X_selected = X[list(feature_combo)]
        
        # Setup stratified cross-validation
        cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
        
        # Perform cross-validated predictions
        y_prob = cross_val_predict(rf_clf, X_selected, y, cv=cv, method='predict_proba')
        y_pred = cross_val_predict(rf_clf, X_selected, y, cv=cv)
        
        # Calculate AUC and accuracy
        auc_score = roc_auc_score(y, y_prob[:, 1])
        accuracy = accuracy_score(y, y_pred)

        # If the current combination has the best AUC, save it
        if auc_score > best_auc_score:
            best_auc_score = auc_score
            best_accuracy = accuracy
            best_features = feature_combo
            best_fpr, best_tpr, _ = roc_curve(y, y_prob[:, 1])
    
    # Plot the ROC curve for the best feature combination
    plt.figure()
    plt.plot(best_fpr, best_tpr, color='blue', lw=2, label=f'ROC curve (AUC = {best_auc_score:.2f})')
    plt.plot([0, 1], [0, 1], color='gray', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.0])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve')
    plt.legend(loc='lower right')
    plt.show()
    
    print("Selected features:")
    for feature in best_features:
        print(f"{feature}")
    
    return best_auc_score, best_accuracy, best_features

# Perform the exhaustive feature selection analysis
auc_score, accuracy, selected_features = exhaustive_feature_selection(data, ['TN', 'Control'], 'LAB', 7)
print(f"AUC score: {auc_score:.2f}")
print(f"Accuracy: {accuracy:.2f}")