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

def load_components():
    rf_model = joblib.load('rf_model.pkl')
    gb_model = joblib.load('gb_model.pkl')
    stacked_model = joblib.load('stacked_es_model.pkl')
    scaler = joblib.load('model_scaler.pkl')
    numeric_columns = ['age', 'creatinine_phosphokinase', 'ejection_fraction', 'platelets', 'serum_creatinine', 'serum_sodium', 'time']
    categorical_columns = ['anaemia', 'diabetes', 'high_blood_pressure', 'sex', 'smoking']
    
    trained_features = rf_model.feature_names_in_  # Extract feature names from the model

    return rf_model, gb_model, stacked_model, scaler, numeric_columns, categorical_columns, trained_features

# Function to get user input
def get_user_input(numeric_columns, categorical_columns):
    inputs = {}
    for col in numeric_columns + categorical_columns:
        if col in categorical_columns:
            inputs[col] = st.selectbox(col, [0, 1])
        else:
            default_val = np.expm1(50.0) if col != 'time' else 50.0
            inputs[col] = np.log1p(st.number_input(col, value=default_val))
    
    user_data = pd.DataFrame([inputs])
    user_data = user_data[numeric_columns + categorical_columns]  # Ensure correct column order
    return user_data

def main():
    st.title("Prediksi Gagal Jantung")

    # Load models and preprocessing components
    rf_model, gb_model, stacked_model, scaler, numeric_columns, categorical_columns, trained_features = load_components()

    # Get user input
    user_input = get_user_input(numeric_columns, categorical_columns)

    # Reorder columns to match trained features and convert to numpy array
    user_input_ordered = user_input[trained_features].to_numpy()

    # Preprocess input
    user_input_ordered[:, :len(numeric_columns)] = scaler.transform(user_input_ordered[:, :len(numeric_columns)])

    # Make predictions using both models
    rf_prediction = rf_model.predict(user_input_ordered)
    gb_prediction = gb_model.predict(user_input_ordered)
    stacked_prediction = stacked_model.predict(user_input_ordered)

    st.write(f"Random Forest Prediction: {'Gagal Jantung' if rf_prediction[0] == 1 else 'Tidak Ada Gagal Jantung'}")
    st.write(f"Gradient Boosting Prediction: {'Gagal Jantung' if gb_prediction[0] == 1 else 'Tidak Ada Gagal Jantung'}")
    st.write(f"stacked Prediction: {'Gagal Jantung' if stacked_prediction [0] == 1 else 'Tidak Ada Gagal Jantung'}")

    # Option: Display user input data (for verification)
    st.subheader("Data Input Pengguna")
    st.write(user_input)

if __name__ == "__main__":
    main()