import streamlit as st
import pandas as pd
import joblib
from huggingface_hub import hf_hub_download

# Configuration
MODEL_REPO_ID = "vnsonly05/engine-condition-rf-model"
MODEL_FILENAME = "best_xgboost_engine_model.joblib"

@st.cache_resource
def load_model():
    # Load the saved model directly from the Hugging Face model hub
    model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME)
    return joblib.load(model_path)

try:
    model = load_model()
    st.success("Model loaded successfully from Hugging Face Hub!")
except Exception as e:
    st.error(f"Error loading model: {e}")

st.title("Engine Failure Prediction")
st.write("Enter the real-time sensor data below to predict if the engine is Healthy or Faulty.")

# Get inputs from the user
col1, col2 = st.columns(2)
with col1:
    engine_rpm = st.number_input("Engine RPM", min_value=0, value=700)
    lub_oil_pressure = st.number_input("Lubrication Oil Pressure", value=2.5)
    fuel_pressure = st.number_input("Fuel Pressure", value=11.5)
with col2:
    coolant_pressure = st.number_input("Coolant Pressure", value=3.0)
    lub_oil_temp = st.number_input("Lubrication Oil Temperature", value=84.0)
    coolant_temp = st.number_input("Coolant Temperature", value=81.0)

if st.button("Predict Engine Condition"):
    # Save inputs into a pandas dataframe
    input_data = pd.DataFrame({
        'Engine rpm': [engine_rpm],
        'Lub oil pressure': [lub_oil_pressure],
        'Fuel pressure': [fuel_pressure],
        'Coolant pressure': [coolant_pressure],
        'lub oil temp': [lub_oil_temp],
        'Coolant temp': [coolant_temp]
    })
    
    st.write("### Captured Input DataFrame:")
    st.dataframe(input_data)
    
    # Predict
    prediction = model.predict(input_data)
    
    if prediction[0] == 1:
        st.success("✅ Prediction: Engine is Healthy (1)")
    else:
        st.error("🚨 Prediction: Engine is Faulty (0) - Maintenance Required!")