# with cluster profiles 

import gradio as gr
import pickle
import pandas as pd
import shap
import matplotlib.pyplot as plt

# Load model
filename = 'xgb_h_new.pkl'
with open(filename, 'rb') as f:
    loaded_model = pickle.load(f)

# Setup SHAP
explainer = shap.Explainer(loaded_model)

# Employee Profiles (From SPSS 3-Cluster Solution)
employee_profiles = {
    "🏆 Leslie Knope": [4.716, 4.792, 4.864, 4.588, 4.849, 4.601],  # Cluster group 1 averages - high engagement, strong support, high workload
    "⚠️ Kevin Malone": [3.045, 3.122, 3.129, 2.886, 3.113, 2.197],  # Cluster group 2 averages - disengaged, low recognition, weak support
    "🌱 Jim Halpert": [3.885, 3.992, 4.119, 3.704, 4.090, 3.377] # Cluster group 3 averages - Moderately engaged, could be more recognized - room to grow
}

# Define the prediction function
def main_func(WellBeing, SupportiveGM, Engagement, Workload, WorkEnvironment, Merit):
    new_row = pd.DataFrame({
        'WellBeing': [WellBeing],
        'SupportiveGM': [SupportiveGM],
        'Engagement': [Engagement],
        'Workload': [Workload],
        'WorkEnvironment': [WorkEnvironment],
        'Merit': [Merit]
    })

    # Predict probability
    prob = loaded_model.predict_proba(new_row)
    shap_values = explainer(new_row)

    # Calculate probability values
    stay_prob = round((1 - float(prob[0][0])) * 100, 2)
    leave_prob = round(float(prob[0][0]) * 100, 2)

    # Dynamic risk label: Changes color & text based on probability
    risk_label = "🔴 High Risk of Turnover" if leave_prob > 50 else "🟢 Low Risk of Turnover"
    risk_color = "red" if leave_prob > 50 else "green"

    risk_html = f"""
    <div style='padding: 15px; border-radius: 8px;'>
        <span style='color: {risk_color}; font-size: 26px; font-weight: bold;'>{risk_label}</span>
        <ul style='list-style-type: none; padding-left: 0; font-size: 20px; font-weight: bold; color: #0057B8;'>
            <li>🧲 Likelihood of Staying: {stay_prob}%</li>
            <li>🚪 Likelihood of Leaving: {leave_prob}%</li>
        </ul>
    </div>
    """

    # Key Insights (Updated for 0.1-point increments)
    insights_html = "<div style='font-size: 18px;'>"
    for feature, shap_val in dict(zip(new_row.columns, shap_values.values[0])).items():
        impact = round(shap_val * 10, 2)  # Scaling impact for 0.1 changes
        icon = "📈" if shap_val > 0 else "📉"
        effect = "raises turnover risk" if shap_val > 0 else "improves retention"
        insights_html += f"<p style='margin: 5px 0;'> {icon} <b>Each 0.1-point increase in {feature} {effect} by {abs(impact)}%.</b></p>"
    insights_html += "</div>"

    # Final Layout (Risk + Key Insights)
    final_layout = f"""
    <table style='width:100%; border-collapse: collapse; margin-top: 10px; background-color: #FFFFFF;'>
        <tr>
            <td style='width: 33%; padding: 15px; background-color: #FFFFFF; border-radius: 8px; vertical-align: top;'>
                {risk_html}
            </td>
            <td style='width: 67%; padding: 15px; background-color: #FFFFFF; border-radius: 8px; vertical-align: top;'>
                <b style='color: #0057B8; font-size: 22px;'>Key Insights:</b>
                {insights_html}
            </td>
        </tr>
    </table>
    """

    # Retention vs. Turnover Chart
    fig, ax = plt.subplots()
    categories = ["Stay", "Leave"]
    values = [stay_prob, leave_prob]
    colors = ["#0057B8", "#D43F00"]
    ax.barh(categories, values, color=colors)
    for i, v in enumerate(values):
        ax.text(v + 2, i, f"{v:.2f}%", va='center', fontweight='bold', fontsize=12)
    ax.set_xlabel("Probability (%)")
    ax.set_title("Retention vs. Turnover Probability")
    plt.tight_layout()
    prob_chart_path = "prob_chart.png"
    plt.savefig(prob_chart_path, transparent=True)
    plt.close()

    # SHAP Chart
    fig, ax = plt.subplots()
    shap.plots.bar(shap_values[0], max_display=6, show=False)
    ax.set_title("Key Drivers of Turnover Risk")
    plt.tight_layout()
    shap_plot_path = "shap_plot.png"
    plt.savefig(shap_plot_path, transparent=True)
    plt.close()

    return final_layout, prob_chart_path, shap_plot_path

# UI Setup
with gr.Blocks() as demo:
    gr.HTML("""
    <div style="text-align: center;">
        <img src="https://upload.wikimedia.org/wikipedia/commons/thumb/5/5f/Hilton_Worldwide_Logo_2010.svg/2560px-Hilton_Worldwide_Logo_2010.svg.png" width="250px">
    </div>
""")
    gr.Markdown("<h1 style='color: #0057B8;'>Hilton Team Member Retention Predictor</h1>")
    gr.Markdown("""
    <div style='font-size: 20px; color: #0057B8;'>
        ✨ <b>Welcome to Hilton’s Employee Retention Predictor</b><br>
        This tool helps <b>HR leaders & managers</b> assess <b>team member engagement</b>  
        and predict <b>turnover risk</b> using AI-powered insights.<br>  
        🔍 <b>See what factors drive retention & make data-driven decisions.</b>  
    </div>
    """)

    # Dropdown for Employee Profiles
    profile_dropdown = gr.Dropdown(choices=list(employee_profiles.keys()), label="Select Employee Profile")

    # Sliders for input features
    with gr.Row():
        WellBeing = gr.Slider(label="WellBeing Score", minimum=1, maximum=5, value=4, step=0.1)
        SupportiveGM = gr.Slider(label="Supportive GM Score", minimum=1, maximum=5, value=4, step=0.1)
        Engagement = gr.Slider(label="Engagement Score", minimum=1, maximum=5, value=4, step=0.1)
    with gr.Row():
        Workload = gr.Slider(label="Workload Score", minimum=1, maximum=5, value=4, step=0.1)
        WorkEnvironment = gr.Slider(label="Work Environment Score", minimum=1, maximum=5, value=4, step=0.1)
        Merit = gr.Slider(label="Merit Score", minimum=1, maximum=5, value=4, step=0.1)

    submit_btn = gr.Button("🔎 Click Here to Analyze Retention")

    # Output elements
    prediction = gr.HTML()
    with gr.Row():
        prob_chart = gr.Image(label="Retention vs. Turnover Probability", type="filepath")
        shap_plot = gr.Image(label="Key Drivers of Turnover Risk", type="filepath")

    # Allow profile selection to update sliders
    def update_sliders(profile):
        if profile in employee_profiles:
            return employee_profiles[profile]
        return [4, 4, 4, 4, 4, 4]

    profile_dropdown.change(update_sliders, inputs=[profile_dropdown], outputs=[WellBeing, SupportiveGM, Engagement, Workload, WorkEnvironment, Merit])

    submit_btn.click(main_func, [WellBeing, SupportiveGM, Engagement, Workload, WorkEnvironment, Merit], [prediction, prob_chart, shap_plot])

demo.launch()