import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.ensemble import RandomForestClassifier

# Load the dataset (make sure to update with your actual file path)
file_path = 'C:/Users/16693/OneDrive - purdue.edu/Desktop/Classes/research/data/Tao-Lab-Fall-2024/Adjusted_062223_MH_Data_For_Training.xlsx'
data = pd.read_excel(file_path)

# Task 1: LAB vs TN & Control
data['LAB_vs_Rest'] = data['Type'].map({'LAB': 1, 'TN': 0, 'Control': 0})

# Task 2: TN vs LAB & Control
data['TN_vs_Rest'] = data['Type'].map({'TN': 1, 'LAB': 0, 'Control': 0})

# LAB vs TN & Control features
lab_vs_rest_features = [
    "P01023|A2MG_Acetyl_GPTQEFKK_1", "P01023|A2MG_Acetyl_GPTQEFKK_2", "P02675|FIBB_Acetyl_YQISVNKYR",
    "O75326|SEM7A_Glyco_NPEAPLNVSR", "P01137|TGFB1_Glyco_YSNNSWR", "P10646|TFPI1_Glyco_YFYNNQTK_1",
    "P16109|LYAM3_Glyco_AYSWNISR", "Q14517|FAT1_Glyco_TGALTVQNTTQLR", "O95772|STR3N_Phospho_LLIVQDASER",
    "P13224|GP1BB_Phospho_LSLTDPLVAER_2", "P16949|STMN1_Phospho_ESVPEFPLSPPK", "P51149|RAB7A_Phospho_FQSLGVAFYR",
    "P61006|RAB8A_Phospho_LEGNSPQGSNQGVK_1", "P61006|RAB8A_Phospho_LEGNSPQGSNQGVK_2", "P98082|DAB2_Phospho_SSPNPFVGSPPK",
    "Q06187|BTK_Phospho_STGDPQGVIR_1", "Q15942|ZYX_Phospho_SPGAPGPLTLK"
]

# TN vs LAB & Control features
tn_vs_rest_features = [
    "P01023|A2MG_Acetyl_GPTQEFKK_1", "P01023|A2MG_Acetyl_GPTQEFKK_2", "P02675|FIBB_Acetyl_YQISVNKYR",
    "P62805|H4_Acetyl_GLGKGGAKR", "P16109|LYAM3_Glyco_AYSWNISR", "P31997|CEAM8_Glyco_LFIPNITTK",
    "P80188|NGAL_Glyco_SYNVTSVLFR", "Q9H330|TM245_Glyco_VNNTAVIEK", "P16949|STMN1_Phospho_ESVPEFPLSPPK",
    "P42331|RHG25_Phospho_DIPLSPPAQK", "Q06187|BTK_Phospho_STGDPQGVIR_1", "Q06187|BTK_Phospho_STGDPQGVIR_2",
    "Q15942|ZYX_Phospho_GPPASSPAPAPK_1", "Q15942|ZYX_Phospho_GPPASSPAPAPK_2", "Q15942|ZYX_Phospho_SPGAPGPLTLK",
    "Q6WCQ1|MPRIP_Phospho_SNPDFLK", "Q9UPY3|DICER_Phospho_TLLSESPGK", "Q9Y3Q8|T22D4_Phospho_SPDPFGAVAAQK"
]

# Function to get feature importance and plot VIP as a dot plot
def plot_vip_and_list_features(data, selected_features, target_column, title="VIP Plot"):
    # Subset the dataset to use only the selected features
    X_selected = data[selected_features]
    y = data[target_column]
    
    # Initialize the Random Forest classifier
    rf_clf = RandomForestClassifier(n_estimators=100, random_state=42, class_weight='balanced')
    
    # Fit the model on the selected features
    rf_clf.fit(X_selected, y)
    
    # Get the feature importances
    feature_importances = rf_clf.feature_importances_
    
    # List the features and their importance
    print(f"Selected Features and Their Importance ({title}):")
    for feature, importance in zip(selected_features, feature_importances):
        print(f"{feature}: {importance:.4f}")
    
    # Sort the features by importance
    sorted_idx = np.argsort(feature_importances)[::-1]
    sorted_features = np.array(selected_features)[sorted_idx]
    sorted_importances = feature_importances[sorted_idx]
    
    # Plot the VIP as a dot plot
    plt.figure(figsize=(10, 6))
    plt.scatter(sorted_importances, sorted_features, color='blue', s=100)  # Scatter plot (dot plot)
    plt.xlabel('Importance')
    plt.title(title)
    # Remove the background and display only grid lines
    plt.gca().patch.set_visible(False)  # Removes the background
    #plt.grid(True)  # Keeps the grid lines visible

    plt.gca().invert_yaxis()  # Invert y-axis to have the most important feature at the top
    # Add horizontal lines at each feature row
    for ytick in plt.gca().get_yticks():
        plt.axhline(y=ytick, color='gray', linestyle='-', linewidth=0.5)  # Horizontal lines only
    plt.tight_layout()  # Automatically adjusts the plot to make space for the labels

    plt.show()

# Plot VIP for LAB vs TN & Control as a dot plot
plot_vip_and_list_features(data, lab_vs_rest_features, target_column='LAB_vs_Rest', title="VIP LAB vs TN & Control")

# Plot VIP for TN vs LAB & Control as a dot plot
plot_vip_and_list_features(data, tn_vs_rest_features, target_column='TN_vs_Rest', title="VIP TN vs LAB & Control")