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

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

# Create the main function for server
# Create the main function for server

def main_func(CLIMATE_SCENARIO,SOVI_SCORE,SOCIAL,ECONOMY,HOUSING_INFRASTRUCTURE,COMMUNITY_CAPITAL,INSTITUTIONAL,ENVIRONMENT):
    
    new_row = pd.DataFrame.from_dict({'CLIMATE_SCENARIO': CLIMATE_SCENARIO,'SOVI_SCORE':SOVI_SCORE,
              'SOCIAL':SOCIAL,'ECONOMY':ECONOMY,'HOUSING_INFRASTRUCTURE':HOUSING_INFRASTRUCTURE,
              'COMMUNITY_CAPITAL':COMMUNITY_CAPITAL,'INSTITUTIONAL':INSTITUTIONAL,'ENVIRONMENT':ENVIRONMENT}, orient = 'index').transpose()
    
    prob = loaded_model.predict_proba(new_row)
    
    # compute SHAP values
    explainer = shap.TreeExplainer(loaded_model)
    shap_values = explainer.shap_values(new_row,check_additivity=False)
    plot = shap.summary_plot(shap_values[0], new_row.values, feature_names = new_row.columns)





    plt.tight_layout()
    local_plot = plt.gcf()
    plt.close()
    
    return {"Very Low Risk": float(prob[0][0]), "Moderately Low Risk": float(prob[0][1]), "Moderate Risk": float(prob[0][2]), "Moderately High Risk": float(prob[0][3]),  "Very High Risk": float(prob[0][4])}, local_plot 

# Create the UI
title = "**Climate Risk Model** 🌍"
description1 = """
The pre-selected counties are Miami-Dade County in Florida, Washington County in Minnesota, or Falls Church County in Virginia. To acquire the data for a county of their choosing, you can navigate to a team-built tableau dashboard (https://public.tableau.com/app/profile/michael.durst/viz/ClimateChangeHackathon/ExpectedRiskbyCounty?publish=yes). 
"""

description2 ="""
Either select one of three counties at bottom of page or enter custom data in the fields. Once the interface is pre-populated with data or you enter data from a county of your choosing, you can experiment with the data elements by clicking analyze and adjusting the variables available on the sliding scale to see how the outputs are affected.
"""

description3 = """
As an output, you can then see the effects these variables have on the risk rating.
""" 

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description1)
    gr.Markdown(description2)
    gr.Markdown(description3)
    with gr.Row():        
        with gr.Column():                 
            CLIMATE_SCENARIO = gr.Slider(label="Climate Scenario", minimum=0, maximum=2, value=0, step=1)        
            ##EAL_SCORE = gr.Slider(label="EAL Score", minimum=0, maximum=100, value=20, step=5)
            SOVI_SCORE = gr.Slider(label="SOVI Score", minimum=0, maximum=100, value=20, step=5)
            SOCIAL = gr.Slider(label="Social", minimum=0, maximum=1, value=.5, step=.1)
            ECONOMY = gr.Slider(label="Economy", minimum=0, maximum=1, value=.5, step=.1)
            HOUSING_INFRASTRUCTURE = gr.Slider(label="Housing Infrastructure", minimum=0, maximum=1, value=.5, step=.1)
            COMMUNITY_CAPITAL = gr.Slider(label="Community Capital", minimum=0, maximum=1, value=.5, step=.1)
            INSTITUTIONAL = gr.Slider(label="Institutional", minimum=0, maximum=1, value=.5, step=.1)
            ENVIRONMENT = gr.Slider(label="Environment", minimum=0, maximum=1, value=.5, step=.1)
            submit_btn = gr.Button("Analyze")
        with gr.Column(visible=True) as output_col:
            label = gr.Label(label = "RISK RATING")
            local_plot = gr.Plot(label = 'Shap:')
           
        
            submit_btn.click(
                main_func,
                [CLIMATE_SCENARIO,SOVI_SCORE,SOCIAL,ECONOMY,HOUSING_INFRASTRUCTURE,COMMUNITY_CAPITAL,INSTITUTIONAL,ENVIRONMENT],
                [label,local_plot], api_name="Climate Risk Model"
            )
    
    gr.Markdown("### Click on any of the examples below to see how it works:")
    gr.Examples([[0, 63.85, .564, .4703, .3068, .2161, .3623, .6264]], 
                [CLIMATE_SCENARIO,SOVI_SCORE,SOCIAL,ECONOMY,HOUSING_INFRASTRUCTURE,COMMUNITY_CAPITAL,INSTITUTIONAL,ENVIRONMENT], 
                [label,local_plot], main_func, cache_examples=True, label="Miami-Dade County, Florida")
    gr.Examples([[0, 15.37, .7231, .5359, .2884, .3828, .4070, .5015]], 
                [CLIMATE_SCENARIO,SOVI_SCORE,SOCIAL,ECONOMY,HOUSING_INFRASTRUCTURE,COMMUNITY_CAPITAL,INSTITUTIONAL,ENVIRONMENT], 
                [label,local_plot], main_func, cache_examples=True, label="Washington County, Minnesota")
    gr.Examples([[0, 4.178, .8181, .5221, .3878, .2463, .389,.3921]], 
                [CLIMATE_SCENARIO,SOVI_SCORE,SOCIAL,ECONOMY,HOUSING_INFRASTRUCTURE,COMMUNITY_CAPITAL,INSTITUTIONAL,ENVIRONMENT], 
                [label,local_plot], main_func, cache_examples=True, label="Falls Church, Virginia")

demo.launch()