import pickle 
import pandas as pd
import shap
from shap.plots._force_matplotlib import draw_additive_plot
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt

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

# Setup SHAP
explainer = shap.Explainer(loaded_model) # PLEASE DO NOT CHANGE THIS.

# Create the main function for server
def main_func(age, education_num, sex, capital_gain, capital_loss, hours_per_week):
    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.force_plot(shap_values[0], matplotlib=True, figsize=(30,30), show=False)
    # plot = shap.plots.waterfall(shap_values[0], max_display=6, show=False)
    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

# Create the 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 and click 'Analyze' to see the prediction and explanation."""

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

    # 🎛 Preset scenario dropdown
    scenario = gr.Dropdown(
        ["Select a Sample", 
     "👨‍💻 Young Tech Worker: 28 yrs, college degree, 45 hrs/week", 
     "👵 Retired Part-Timer: 65 yrs, no college, 20 hrs/week", 
     "👩‍🏫 Mid-Career Teacher: 42 yrs, 14 education years, 38 hrs/week", 
     "👨‍🔧 Manual Laborer: 50 yrs, 9 education years, 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_num = gr.Number(label="🎓 Education Level (numeric)", value=10)
    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)

    submit_btn = gr.Button("🔎 Analyze")

    # 🔁 Handle preset scenario changes
    def fill_scenario(scenario_choice):
        if scenario_choice == "👨‍💻 Young Tech Worker: 28 yrs, college degree, 45 hrs/week":
            return [28, 16, "Male", 0, 0, 45]
        elif scenario_choice == "👵 Retired Part-Timer: 65 yrs, no college, 20 hrs/week":
            return [65, 8, "Female", 0, 0, 20]
        elif scenario_choice == "👩‍🏫 Mid-Career Teacher: 42 yrs, 14 education years, 38 hrs/week":
            return [42, 14, "Female", 0, 0, 38]
        elif scenario_choice == "👨‍🔧 Manual Laborer: 50 yrs, 9 education years, 60 hrs/week":
            return [50, 9, "Male", 0, 0, 60]
        else:
            return [35, 10, "Male", 0, 0, 40]  # Default values

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

    # 🧠 Prediction output
    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")

    # 🧠 Wrap predict + confidence slider logic
    def wrapped_main(age, education_num, sex, capital_gain, capital_loss, hours_per_week):
        result, shap_plot = main_func(age, education_num, sex, capital_gain, capital_loss, hours_per_week)
        return result, float(result["Chance of Earning > $50K"]) * 100, shap_plot

    submit_btn.click(
        wrapped_main,
        [age, education_num, sex, capital_gain, capital_loss, hours_per_week],
        [label, confidence, local_plot],
        api_name="Salary_Predictor"
    )

    gr.Markdown("### 🧪 Try Some Examples:")
    gr.Examples(
        [
            [28, 16, "Male", 0, 0, 45],
            [60, 8, "Female", 0, 0, 25]
        ],
        [age, education_num, sex, capital_gain, capital_loss, hours_per_week],
        [label, confidence, local_plot],
        wrapped_main,
        cache_examples=True
    )

demo.launch()