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("h25_xgb_hyperopt.pkl", 'rb'))

#Setup SHAP
explainer = shap.Explainer(loaded_model) #DO NOT CHANGE PER INSTRUCTIONS

#Main Function for Server
def main_func(PassionateAtWork,Workload,SupportiveGM,WorkEnvironment,Informed,LearningDevelopment,JobSecurity):
    new_row = pd.DataFrame.from_dict({'PassionateAtWork':PassionateAtWork, 'Workload':Workload, 
                'SupportiveGM':SupportiveGM,'WorkEnvironment':WorkEnvironment, 'Informed':Informed,
                'LearningDevelopment': LearningDevelopment, 'JobSecurity':JobSecurity}, 
                orient = 'index').transpose()

    prob = loaded_model.predict_proba(new_row)

    shap_values = explainer(new_row)
    # plot = shap.force(shap_values[0], matplotlib=True, figsize(30,30), show = False)
    # plot = shap.plots.waterfall(shap_values[0], max_display = 7, show = False)
    plot = shap.plots.bar(shap_values[0], max_display = 7, order=shap.Explanation.abs, show_data = 'auto', show = False)

    plt.tight_layout()
    local_plot = plt.gcf()
    plt.rcParams['figure.figsize'] = 6,4
    plt.close()

    return {"Leave": float(prob[0][0]), "Stay": 1-float(prob[0][0])}, local_plot

#Create UI
title = "**Employee Turnover Predictor & Intrepreter**"
description1 = """
This app takes seven inputs about employees' satisfaction with different aspects of their work (such as Passionate at Work, ...) and predicts whether the employee intends to stay with the employer or leave. There are two outputs from the app: 1- the predicted probability of stay or leave, 2- Shapley's force-plot which visualizes the extent to which each factor impacts the stay/ leave prediction.
"""

description2 = """
To use the app, click on one of the examples, or adjust the values of the seven employee satisfaction factors, and click Analyze.
"""

with gr.Blocks(title = title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description1)
    gr.Markdown("""---""")
    gr.Markdown(description2)
    gr.Markdown("""---""")
    with gr.Row():
        with gr.Column():
            PassionateAtWork = gr.Slider(label= "Passionate At Work", minimum = 1, maximum = 5, value = 4, step = .1)
            Workload = gr.Slider(label= "Workload", minimum = 1, maximum = 5, value = 4, step = .1)
            SupportiveGM = gr.Slider(label= "Supportive GM", minimum = 1, maximum = 5, value = 4, step = .1)
            WorkEnvironment = gr.Slider(label= "Work Environment", minimum = 1, maximum = 5, value = 4, step = .1)
            Informed = gr.Slider(label= "Informed", minimum = 1, maximum = 5, value = 4, step = .1)
            LearningDevelopment = gr.Slider(label= "Learning Development", minimum = 1, maximum = 5, value = 4, step = .1)
            JobSecurity = gr.Slider(label= "Job Security", minimum = 1, maximum = 5, value = 4, step = .1)
            submit_btn = gr.Button("Analyze")
        with gr.Column(visible=True,scale=1, min_width=600) as output_col:
            label = gr.Label(label = "Predicted Label")
            local_plot = gr.Plot(label = 'Shap:')

            submit_btn.click(
                main_func,
                [PassionateAtWork,Workload,SupportiveGM,WorkEnvironment,Informed,LearningDevelopment,JobSecurity],
                [label,local_plot], api_name="Employee_Turnover"
            )

    gr.Markdown('### Click on any of the examples below to see how it works:')
    gr.Examples([[3.7,3.1,3.1,3.3,3.1,2.9,3.1], [3.7,3.1,5.0,5.0,3.1,5.0,3.1], [4.2,3.8,4.0,4.1,3.8,4.0,4.0], [4.2,3.8,5.0,5.0,3.8,5.0,4.0]], 
                [PassionateAtWork,Workload,SupportiveGM,WorkEnvironment,Informed,LearningDevelopment,JobSecurity], 
                [label,local_plot], main_func, cache_examples=True)

demo.launch()