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

# Load the model
loaded_model = pickle.load(open("salar_xgb_team.pkl", 'rb'))

# SHAP setup
explainer = shap.Explainer(loaded_model)  # DO NOT CHANGE

# Education mapping
education_map = {
    "Less than 1st grade": 1,
    "1st–4th grade": 2,
    "5th–6th grade": 3,
    "7th–8th grade": 4,
    "9th grade": 5,
    "10th grade": 6,
    "11th grade": 7,
    "12th grade (no diploma)": 8,
    "High School Grad": 9,
    "Some College": 10,
    "Associate's Degree (Voc)": 11,
    "Associate's Degree (Acad)": 12,
    "Bachelor's Degree": 13,
    "Master's Degree": 14,
    "Professional School": 15,
    "Doctorate": 16
}

# Main model logic
def main_func(age, education_level, sex, capital_gain, capital_loss, hours_per_week):
    education_num = education_map[education_level]
    sex = 1 if sex == "Female" else 0

    new_row = pd.DataFrame.from_dict({
        'age': age,
        'education-num': education_num,
        'sex': sex,
        'capital-gain': capital_gain,
        'capital-loss': capital_loss,
        'hours-per-week': hours_per_week
    }, orient='index').transpose()

    prob = loaded_model.predict_proba(new_row)

    shap_values = explainer(new_row)
    plot = shap.plots.bar(shap_values[0], max_display=6, order=shap.Explanation.abs, show_data='auto', show=False)

    plt.tight_layout()
    local_plot = plt.gcf()
    plt.close()

    return {
        "Chance of Earning > $50K": float(prob[0][1]),
        "Chance of Earning ≤ $50K": float(prob[0][0])
    }, local_plot

# Gradio UI
title = "**Household Income Predictor** 💰"
description1 = """This app uses your input to predict whether a household earns more or less than $50K per year."""
description2 = """Adjust the values below or select a sample profile, then click 'Analyze' to see the prediction and feature impact."""

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description1)
    gr.Markdown("---")
    gr.Markdown(description2)
    gr.Markdown("---")

    # Sample profile dropdown
    scenario = gr.Dropdown(
        ["Select a Sample", 
         "👨‍💻 Young Tech Worker: 28 yrs, Bachelor's, 45 hrs/week", 
         "👵 Retired Part-Timer: 65 yrs, no college, 20 hrs/week", 
         "👩‍🏫 Mid-Career Teacher: 42 yrs, Master's, 38 hrs/week", 
         "👨‍🔧 Manual Laborer: 50 yrs, High School Grad, 60 hrs/week"],
        label="📋 Choose a Sample Profile (optional — autofills values to explore common cases)"
    )

    # Inputs
    with gr.Row():
        age = gr.Number(label="🧓 Age", value=35)
        education_level = gr.Dropdown(
            list(education_map.keys()),
            label="🎓 Education Level",
            value="Some College"
        )
    with gr.Row():
        sex = gr.Radio(["Male", "Female"], label="🧍 Sex")
        capital_gain = gr.Number(label="📈 Capital Gain", value=0)
        capital_loss = gr.Number(label="📉 Capital Loss", value=0)
        hours_per_week = gr.Number(label="⏱ Hours per Week", value=40)

    # Handle preset scenario changes
    def fill_scenario(scenario_choice):
        if scenario_choice == "👨‍💻 Young Tech Worker: 28 yrs, Bachelor's, 45 hrs/week":
            return [28, "Bachelor's Degree", "Male", 0, 0, 45]
        elif scenario_choice == "👵 Retired Part-Timer: 65 yrs, no college, 20 hrs/week":
            return [65, "9th grade", "Female", 0, 0, 20]
        elif scenario_choice == "👩‍🏫 Mid-Career Teacher: 42 yrs, Master's, 38 hrs/week":
            return [42, "Master's Degree", "Female", 0, 0, 38]
        elif scenario_choice == "👨‍🔧 Manual Laborer: 50 yrs, High School Grad, 60 hrs/week":
            return [50, "High School Grad", "Male", 0, 0, 60]
        else:
            return [35, "Some College", "Male", 0, 0, 40]

    scenario.change(
        fn=fill_scenario,
        inputs=[scenario],
        outputs=[age, education_level, sex, capital_gain, capital_loss, hours_per_week]
    )

    # Outputs
    with gr.Column(visible=True) as output_col:
        label = gr.Label(label="🧠 Predicted Income")
        confidence = gr.Slider(0, 100, value=50, label="📊 Confidence in > $50K", interactive=False)
        local_plot = gr.Plot(label="🔍 Top SHAP Features")

    # Wrapped function for UI
    def wrapped_main(age, education_level, sex, capital_gain, capital_loss, hours_per_week):
        result, shap_plot = main_func(age, education_level, sex, capital_gain, capital_loss, hours_per_week)
        return result, float(result["Chance of Earning > $50K"]) * 100, shap_plot

    # Button
    submit_btn = gr.Button("🔎 Analyze")
    submit_btn.click(
        wrapped_main,
        [age, education_level, sex, capital_gain, capital_loss, hours_per_week],
        [label, confidence, local_plot],
        api_name="Salary_Predictor"
    )

    # Examples
    gr.Markdown("### 🧪 Try Some Examples:")
    gr.Examples(
        [
            [28, "Bachelor's Degree", "Male", 0, 0, 45],
            [60, "9th grade", "Female", 0, 0, 25]
        ],
        [age, education_level, sex, capital_gain, capital_loss, hours_per_week],
        [label, confidence, local_plot],
        wrapped_main,
        cache_examples=True
    )

demo.launch()