# -*- coding: utf-8 -*-
"""app

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1B_g2XLYu46kFDIFzNnnJzBQ0GBPssCQw
"""

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
loaded_model = pickle.load(open("salar_xgb_team.pkl", 'rb'))

# Setup SHAP (do not change)
explainer = shap.Explainer(loaded_model)

# Define main prediction function
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.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 {"≤ $50K": float(prob[0][0]), "> $50K": float(prob[0][1])}, local_plot

# Gradio UI
title = "**Household Income Predictor & Interpreter** 💰"
description1 = """This app takes demographic and economic features to predict whether a household earns ≤ $50K or > $50K annually.🚀"""
description2 = """Adjust the values and click Analyze to get predictions and feature importance."""

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

    gr.Image("Household.png")

    age = gr.Number(label="Age", value=35)
    education_num = gr.Number(label="Education Level (numeric)", value=10)
    sex = gr.Radio(choices=["Male", "Female"], label="Sex", value="Female")
    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)
    # salary_class = gr.Number(label="(Optional) Salary Class for SHAP Context", value=0)  # Can remove if not needed

    submit_btn = gr.Button("Analyze")

    with gr.Column(visible=True) as output_col:
        label = gr.Label(label="Predicted Income")
        local_plot = gr.Plot(label='SHAP Interpretation:')

    submit_btn.click(
        main_func,
        [age, education_num, sex, capital_gain, capital_loss, hours_per_week],
        [label, local_plot], api_name="Income_Predictor"
    )

    gr.Markdown("### Try these examples:")
    gr.Examples(
        [[39,13, "Male", 0, 0, 40], [52, 9, "Female", 0, 1876, 45]],
        [age, education_num, sex, capital_gain, capital_loss, hours_per_week],
        [label, local_plot],
        main_func,
        cache_examples=True
    )

demo.launch()