import streamlit as st
import streamlit.components.v1 as components
import numpy as np
import pickle 
import pandas as pd

pd.set_option('display.max_columns', 500)
def impute_missing_values(df_,dict_impute):
    return df_.fillna(dict_impute)

to_rescale_ = ['diff_MS','MS_d','pr_pre_tavi',
               'surface_systole','MS_s','age',
               'ncc_calcif_n','calc_risque_n']

to_encode_ = ['syncope_Oui', 'lcc_calc_1.0', 'bloc_branche_pre_bbd']

st.title("Pacemaker Implantation (IPM) risk prediction")
st.subheader(f':red[Pre procedural] Tavi implantation model')
         
cols =  to_rescale_+to_encode_



with st.form(key='cols_form'):
    c1,c2= st.columns(2)
    c3,c4 = st.columns(2)
    c5,c6,c7,c8 = st.columns(4)
    c9,c10,c11  = st.columns(3)

    with c1:
        age = st.number_input('Age (year)',value=83)
    with c2:
        syncope_Oui = st.selectbox("Syncope",key="syncope",options=["Yes", "No"])
    with c3:
        pr_pre_tavi = st.number_input('Pre-TAVI PR time (ms)',value=186)
    with c4:
        bloc_branche_pre_bbd = st.selectbox("Right bundle branch block",key="prebbd",options=["Yes", "No"])
    with c5:
        lcc_calc_1 = st.selectbox("LCC calcification grade",key='lcccalc1',options=[0,1,2,3])
    with c6:
        ncc_calcif_n = st.selectbox("NCC calcification grade",key='nccalcn',options=[0,1,2,3])
    with c7:
        rcc_calcif_n = st.selectbox("RCC calcification grade",key='rccalcn',options=[0,1,2,3])
    with c8:
        isc_calcif_n = st.selectbox("IS calcification",key='isccalcn',options=["Yes", "No"])
    with c9:
        MS_s = st.number_input('Systolic membranous septal length(mm)',value=8.0)
    with c10:
        MS_d = st.number_input('Diastolic membranous septal length (mm)',value=8.9)
    with c11:
        surface_systole = st.number_input('Systolic aortic annular surface (mm²)',value=483)

        
    submitButton = st.form_submit_button(label = 'Predict')


#load model, set cache to prevent reloading
@st.cache_resource()
def load_model():
    with open(r"models/ensemble_models_pre.pkl", "rb") as input_file:
        models = pickle.load(input_file)
    return models


#@st.cache_resource()
def load_scaler():
    with open(r"models/ensemble_scaler_pre.pkl", "rb") as input_file:
        scalers = pickle.load(input_file)
    return scalers


@st.cache_resource()
def load_impute():
    with open(r"models/ensemble_dict_impute_pre.pkl", "rb") as input_file:
        dicts_impute = pickle.load(input_file)
    return dicts_impute

with st.spinner("Loading Model...."):
    models = load_model()
    dicts_impute = load_impute()
    scalers = load_scaler()

    


#st.write("Predicting Class...") 
with st.spinner("Prediction..."):
    
    # preprocessing
    diff_MS = MS_d - MS_s
    lcc_calc_1 = 1 if lcc_calc_1 == 1 else 0
    syncope_Oui = 1 if syncope_Oui == "Yes" else 0
    isc_calcif_n = 1 if isc_calcif_n == "Yes" else 0
    bloc_branche_pre_bbd = 1 if bloc_branche_pre_bbd == "Yes" else 0
    calc_risque_n = ncc_calcif_n + rcc_calcif_n + isc_calcif_n
    age = int(age)*365

       
    bool_warning = False 
    if MS_d<1.5 or MS_d>15.30:
        bool_warning = True
    if MS_s<2.0 or MS_s> 14.61:
        bool_warning = True
    if pr_pre_tavi<90 or pr_pre_tavi>344: 
        bool_warning = True
    if surface_systole<300 or surface_systole>724:
        bool_warning = True
    if age<21156 or age>34939:
        bool_warning = True
    if bool_warning:
        st.write(":warning: Warning Results might not be reliable because:")
    if MS_d<1.5 or MS_d>15.30:
        st.write("Our study population had a dMS (mm) between 1.5 and 15.30 (mm)")
    if MS_s<2.0 or MS_s> 14.61:
        st.write("Our study population had a sMS (mm) between 2 and 14.61 (mm) ")
    if pr_pre_tavi<90 or pr_pre_tavi>344: 
        st.write("Our study population had a pre-TAVI PR surface between 90 and 344 (ms)")
    if surface_systole<300 or surface_systole>724:
        st.write("Our study population had a systolic aortic annular surface between 300 and 724 (mm²) ")
    if age<21156 or age>34939:
        st.write("Our study population had an age between 58 and 95 years old ")
        
    
    
    pred_arr = np.array([[diff_MS,MS_d,pr_pre_tavi,
               surface_systole,MS_s,age,
               ncc_calcif_n,calc_risque_n,
               syncope_Oui, lcc_calc_1, bloc_branche_pre_bbd]])
    
    pred_df = pd.DataFrame(pred_arr,columns=cols)
    print(pred_df.head())
    pred_dfs = [impute_missing_values(pred_df,dict_impute) for dict_impute in dicts_impute]
    dfs_scaled_ = [scaler.transform(df_[to_rescale_]) for df_,scaler in zip(pred_dfs,scalers)]
    dfs_scaled = [pd.DataFrame(columns=to_rescale_,data=df_scaled_) for df_scaled_ in dfs_scaled_]
    dfs_cat = [pd.concat([df_scaled,pred_df[to_encode_]],axis=1) for df_scaled in dfs_scaled]
    pred= round( np.array([model.predict_proba(pred_df_)[0][1] for pred_df_,model in zip(dfs_cat,models)]).mean()*100,2)
    
    if pred>90:
        st.write("PMI risk probability >90 %")
    elif pred<10:
        st.write("PMI risk probability <10 %")
    else:
        st.write("PMI risk probability :",pred,' %')