import streamlit as st
import pandas as pd
import numpy as np
import joblib
import ast

# Set page title and favicon
st.set_page_config(page_title='Heart Failure Prediction', page_icon=':heart:')

# Load models and data
with open('model_scaler.pkl', 'rb') as file_1:
  model_scaler = joblib.load(file_1)

with open('rf_best_model.pkl', 'rb') as file_2:
  rf_random_best = joblib.load(file_2)

with open('list_num_col.txt', 'r') as file_3:
  num_col = file_3.read()

with open('list_cat_col.txt', 'r') as file_4:
  cat_col = file_4.read()

with open('boxcox.pkl', 'rb') as file_5:
  box = joblib.load(file_5)

num_col = ast.literal_eval(num_col)
cat_col = ast.literal_eval(cat_col)

# Set page header
st.title('Heart Failure Prediction')

# Add page description
st.markdown('This web app predicts the likelihood of heart failure based on several factors. Fill in the form below and click the "Predict" button to get your results.')

# Add form inputs
time = st.slider('Follow-up period (days)', min_value=1, max_value=365, step=1)
age = st.slider('Age', 30, 120, step=1)
ejection_fraction = st.number_input('Ejection fraction', min_value=1, max_value=100, step=1)
serum_creatinine = st.number_input('Serum creatinine', min_value=0.0, max_value=10.0, step=0.01)
serum_sodium = st.number_input('Serum sodium', min_value=100, max_value=150, step=1)
anaemia = st.radio('Anemia', ('Yes', 'No'))
diabetes = st.radio('Diabetes', ('Yes', 'No'))
smoking = st.radio('Smoking', ('Yes', 'No'))

# Add prediction button
if st.button('Predict'):
    # Process form inputs
    data_inf = pd.DataFrame({'time': time,
                             'ejection_fraction': ejection_fraction,
                             'age': age,
                             'serum_creatinine': serum_creatinine,
                             'serum_sodium': serum_sodium,
                             'anaemia': 1 if anaemia == 'Yes' else 0,
                             'diabetes': 1 if diabetes == 'Yes' else 0,
                             'smoking': 1 if smoking == 'Yes' else 0
                            }, index=[0])

    data_inf_t = box.transform(data_inf)
    data_inf_num = data_inf_t[num_col]
    data_inf_cat = data_inf_t[cat_col]

    data_inf_num_scaled = data_inf_num.copy()
    data_inf_num_scaled[num_col] = model_scaler.transform(data_inf_num[num_col])
    data_inf_cat_encoded = data_inf_cat.copy()

    data_inf_final = np.concatenate([data_inf_num_scaled, data_inf_cat_encoded], axis=1)

    result = rf_random_best.predict(data_inf_final)[0]

    # Display prediction result
    st.markdown('## Prediction result')
    if result == 1:
        st.markdown(':heart: **High risk of heart failure!** :heart:')
    else:
        st.markdown(':thumbsup: **Low risk of heart failure.** :thumbsup:')