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

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

# Create the main function for server
def main_func(age, sex, cp, trestbps, chol, fbs, restecg, thalach,
       exang, oldpeak, slope, ca, thal):
    new_row = pd.DataFrame.from_dict({'age': age,
                                      'sex':sex,
                                      'cp':cp,
                                      'trestbps':trestbps,
                                      'chol':chol,
                                      'fbs':fbs,
                                      'restecg':restecg,
                                      'thalach':thalach,
                                      'exang':exang,
                                      'oldpeak':oldpeak,
                                      'slope':slope,
                                      'ca':ca,
                                      'thal':thal
                                     }, 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=7, order=shap.Explanation.abs, show_data='auto', show=False)

    plt.tight_layout()
    local_plot = plt.gcf()
    plt.rcParams['figure.figsize'] = 7,4
    plt.close()
    
    return {"Normal Heart Condition": float(prob[0][0]), "Critical Heart Condition": 1-float(prob[0][0])}, local_plot

# Create the UI
title = "**Heart Condition Predictor & Interpreter** 🪐"
description1 = """
This app takes inputs about patients' demographics and medical history to predict whether the patient has heart condition. There are two outputs from the app: 1- the predicted probability of normal condition or heart condition, 2- Shapley's force-plot which visualizes the extent to which each factor impacts the prediction.   
"""

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

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    # gr.Markdown("""![marketing](types-of-employee-turnover.jpg)""")
    gr.Markdown(description1)
    gr.Markdown("""---""")
    gr.Markdown(description2)
    gr.Markdown("""---""")
    with gr.Row():        
        with gr.Column():
            age = gr.Slider(label="age", minimum=17, maximum=74, value=24, step=1)
            sex = gr.Slider(label="sex", minimum=0, maximum=1, value=1, step=1)
            cp = gr.Slider(label="cp Score", minimum=1, maximum=4, value=3, step=.1)
            trestbps = gr.Slider(label="trestbps Score", minimum=94, maximum=200, value=150, step=.1)
            chol = gr.Slider(label="chol Score", minimum=126, maximum=564, value=400, step=.1)
            fbs = gr.Slider(label="fbs Score", minimum=0, maximum=1, value=0, step=.1)
            restecg = gr.Slider(label="restecg Score", minimum=0, maximum=2, value=1, step=.1)
            thalach = gr.Slider(label="thalach Score", minimum=71, maximum=202, value=90, step=.1) 
            exang = gr.Slider(label="exang Score", minimum=0, maximum=1, value=1, step=.1)
            oldpeak = gr.Slider(label="oldpeak Score", minimum=0, maximum=6, value=4, step=.1)
            slope = gr.Slider(label="slope Score", minimum=1, maximum=3, value=2, step=.1)
            ca = gr.Slider(label="ca Score", minimum=0, maximum=3, value=2, step=.1)
            thal = gr.Slider(label="thal Score", minimum=3, maximum=7, 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,
                [age, sex,cp,trestbps,chol,fbs,restecg,thalach,exang,oldpeak,slope,ca,thal],
                [label,local_plot], api_name="Heart_Condition"
            )
    
    gr.Markdown("### Click on any of the examples below to see how it works:")
    gr.Examples([[33,0,1,100,230,1,1,150,0,.9,2,1,6], [39,1,0,170,200,1,1,150,0,1.4,2,1,6]], 
                [age,sex,cp,trestbps,chol,fbs,restecg,thalach,exang,oldpeak,slope,ca,thal], 
                [label,local_plot], main_func, cache_examples=True)

demo.launch()