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.

sex_dictionary = {"Male":0,"Female":1}

# Create the main function for server


def main_func(age,sex,cp,trtbps,chol,fbs,restecg,thalachh,exng,oldpeak,slp,caa,thall):


   new_row = pd.DataFrame.from_dict({'age':age,'sex':sex_dictionary[sex],


             'cp':cp,'trtbps':trtbps,'chol':chol,'fbs':fbs,


             'restecg':restecg,'thalachh':thalachh,'exng':exng,


             'oldpeak':oldpeak,'slp':slp,'caa':caa,'thall':thall}


                                   , 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 {"High Risk": float(prob[0][1]), "Low Risk": 1-float(prob[0][1])}, local_plot



# Create the UI

title = "**Heart Attack Demo App** 🫀"


description1 = """


This app takes thirteen inputs from patients including age, sex, and many other questions that would be asked by a physician when analyzing a patients symptoms and possibility of a heart attack. There are two outputs from the app: 1- the predicted risk percentage of having a heart attack , 2- Shapley's force-plot which visualizes the extent to which each factor impacts the heart attack risk prediction. 🕺


"""


description2 = """


To use the app, click 🫵 on one of the examples, or adjust the values of the 13 factors, and click on Process.


"""


with gr.Blocks(title=title) as demo:


   with gr.Row():


       with gr.Column():


           gr.Markdown(f"## {title}")


           gr.Markdown(description1)


           gr.Markdown("""---""")


           gr.Markdown(description2)


           gr.Markdown("""---""")


       with gr.Column():


           gr.Markdown("""![Take care!](file/WechatIMG343.jpeg)""")
   


   with gr.Row():


       with gr.Column():


           age = gr.Number(label="Age", value=40)


           sex = gr.Dropdown(["Male", "Female"], label="Sex")


           cp = gr.Slider(minimum=1, maximum=4, default=1, step=1, label="Chest Pain Type")


           trtbps = gr.Slider(minimum=50, maximum=200, default=120, step=1, label="Resting Blood Pressure (in mm Hg)")


           chol = gr.Slider(minimum=80, maximum=500, default=190, step=1, label="Cholesterol Level (mg/dL)")


           fbs = gr.Slider(minimum=0, maximum=1, default=0, step=.1, label="(fasting blood sugar > 120 mg/dl) (1 = true; 0 = false)")


           restecg = gr.Slider(minimum=0, maximum=200, step=1, default=80, label="resting electrocardiographic results")


       with gr.Column():


           thalachh = gr.Slider(minimum=80, maximum=400, step=1, default=200, label="maximum heart rate achieved")


           exng = gr.Slider(minimum=80, maximum=400, step=1, default=200, label="maximum heart rate achieved")


           oldpeak = gr.Slider(minimum=0, maximum=10, step=.1, default=1, label="ST depression induced by exercise relative to rest")


           slp = gr.Slider(minimum=0, maximum=2, step=.1, default=1, label="speech-language pathologist")


           caa = gr.Slider(minimum=0, maximum=4, step=.1, default=2, label="cerebral amyloid angiopathy")


           thall = gr.Slider(minimum=0, maximum=3, default=2, step=.1, label="thallium stress test")


           submit_btn = gr.Button("Process")


   with gr.Row(visible=True) as output_col:


       label = gr.Label(label = "Predicted Label")


       local_plot = gr.Plot(label = 'Shap:')


   submit_btn.click(


       main_func,


       [age,sex,cp,trtbps,chol,fbs,restecg,thalachh,exng,oldpeak,slp,caa,thall],


       [label,local_plot], api_name="Heart Attack Rate"


   )


   gr.Markdown("### Click on any of the examples below to see how it works:")


   gr.Examples([[20,"Male",1,120,190,0,80,200,200,1,1,2,2], [30,"Female",1,120,190,0,80,200,200,1,1,2,2]],


               [age,sex,cp,trtbps,chol,fbs,restecg,thalachh,exng,oldpeak,slp,caa,thall],


               [label,local_plot], main_func, cache_examples=True)


demo.launch()