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("heart_xgb.pkl", 'rb'))
# Setup SHAP
explainer = shap.Explainer(loaded_model) # PLEASE DO NOT CHANGE THIS.

# Create the main function for server
#def main_func(ValueDiversity,AdequateResources,Voice,GrowthAdvancement,Workload,WorkLifeBalance):
#    new_row = pd.DataFrame.from_dict({'ValueDiversity':ValueDiversity,'AdequateResources':AdequateResources,
#              'Voice':Voice,'GrowthAdvancement':GrowthAdvancement,'Workload':Workload,
#              'WorkLifeBalance':WorkLifeBalance}, 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 {"Leave": float(prob[0][0]), "Stay": 1-float(prob[0][0])}, local_plot

def main_func(age, sex, exang, ca, cp, trtbps, chol, fbs, rest_ecg, thalach):
    new_row = pd.DataFrame.from_dict({'age': age, 'sex': sex, 'cp': cp, 'trtbps': trtbps,
                                      'chol': chol, 'fbs': fbs, 'restecg': rest_ecg, 
                                      'thalachh': thalach, 'exng': exang, 'caa': ca}, orient = 'index').transpose()
    #new_row.info()
    #print(new_row)
    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 {"Less Chance of Heart Attack": float(prob[0][0]), "More Chance of Heart Attack": 1-float(prob[0][0])}, local_plot

# Create the UI
title = "Heart Attack Predictor Application ❤️"
description1 = "This app takes information gathered from patients and healthcare providers to predict the likelihood of a cardiac event. NOTE- this application should not be used for medical or diagnostic purposes👍"  

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

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
#    gr.Markdown("""![marketing](file/marketing.jpg)""")
    gr.Markdown(description1)
    gr.Markdown("""---""")
    gr.Markdown(description2)
    gr.Markdown("""---""")
    #ValueDiversity = gr.Slider(label="ValueDiversity Score", minimum=1, maximum=5, value=4, step=1)
    #AdequateResources = gr.Slider(label="AdequateResources Score", minimum=1, maximum=5, value=4, step=1)
    #Voice = gr.Slider(label="Voice Score", minimum=1, maximum=5, value=4, step=1)
    #GrowthAdvancement = gr.Slider(label="GrowthAdvancement Score", minimum=1, maximum=5, value=4, step=1)
    #Workload = gr.Slider(label="Workload Score", minimum=1, maximum=5, value=4, step=1)
    #WorkLifeBalance = gr.Slider(label="WorkLifeBalance Score", minimum=1, maximum=5, value=4, step=1)
    with gr.Row():
        with gr.Column():
            age = gr.Slider(label = "Age", minimum = 1, maximum = 100, value = 30, step = 1)
        with gr.Column():
            sex = gr.Radio(["Male", "Female"],label = "Sex", type="index")
    with gr.Row():
        with gr.Column():
            ca = gr.Slider(label = "Number of Major Blood Vessels", minimum = 0, maximum = 3, value = 3, step = 1)
        with gr.Column():
            cp = gr.Dropdown(["Typical Angina", "Atypical Angina", "Non-anginal Pain", "Asymptomatic"], label = "Chest Pain Type", type = "index")
    with gr.Row():
        with gr.Column():
            exang = gr.Radio(["No", "Yes"], label = "Do you have Exercise Induced Angina", type = "index")
        with gr.Column():
            fbs = gr.Radio(["Yes", "No"], label = "Is the fasting Blood Sugar >120 mg/dl", type = "index")
    with gr.Row():
        rest_ecg = gr.Dropdown(["Normal", "Having ST-T abnormality", "Showing probable or definite left ventricular hypertrophy by Estes' Criteria"], label = "Resting ECG Results", type = "index")
    with gr.Row():
        thalach = gr.Slider(label = "Maximum Heart Rate Achieved", minimum = 50, maximum = 250, value = 80, step = 1)
    with gr.Row():
        trtbps = gr.Slider(label = "Resting Blood Pressure (in mm Hg)", minimum = 75, maximum = 250, value = 100, step = 1)
    with gr.Row():
        chol = gr.Slider(label="Cholesterol in mg/dl", minimum = 100, maximum = 600, value = 200, step = 1)
    submit_btn = gr.Button("Analyze")

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

    submit_btn.click(
        main_func,
        [age, sex, exang, ca, cp, trtbps, chol, fbs, rest_ecg, thalach],
        [label,local_plot], api_name="Heart Attack Probability"
    )
    gr.Markdown("### Click on any of the examples below to see how it works:")
    gr.Examples([[36, "Male", "Yes", 2, "Typical Angina", 100, 200, "No", "Normal", 120], [80, "Female", "No", 1, "Atypical Angina", 200, 400, "Yes", "Having ST-T abnormality", 100]], [age, sex, exang, ca, cp, trtbps, chol, fbs, rest_ecg, thalach], [label,local_plot], main_func, cache_examples=True)

    

demo.launch()