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

HF_USERNAME = "Vbhadiar"
MODEL_ID = f"{HF_USERNAME}/engine-failure-classifier"

@st.cache_resource
def load_model():
    model_path = hf_hub_download(
        repo_id=MODEL_ID,
        filename="best_model.pkl",
        repo_type="model"
    )
    return joblib.load(model_path)

model = load_model()

st.title("🔧 Engine Predictive Maintenance")
st.markdown("Predict whether an engine requires maintenance based on sensor readings.")

col1, col2 = st.columns(2)

with col1:
    engine_rpm = st.number_input("Engine RPM", 0, 3000, 800)
    lub_oil_pressure = st.number_input("Lub Oil Pressure (bar)", 0.0, 10.0, 3.5, 0.1)
    fuel_pressure = st.number_input("Fuel Pressure (bar)", 0.0, 25.0, 7.0, 0.1)

with col2:
    coolant_pressure = st.number_input("Coolant Pressure (bar)", 0.0, 10.0, 2.5, 0.1)
    lub_oil_temp = st.number_input("Lub Oil Temp (°C)", 50.0, 120.0, 78.0, 0.5)
    coolant_temp = st.number_input("Coolant Temp (°C)", 50.0, 200.0, 80.0, 0.5)

if st.button("Predict"):
    df_input = 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],
    })
    pred = model.predict(df_input)[0]
    prob = model.predict_proba(df_input)[0][1]

    if pred == 1:
        st.error(f"⚠️ Maintenance required (failure probability: {prob*100:.1f}%)")
    else:
        st.success(f"✅ Engine normal (failure probability: {prob*100:.1f}% ")

    st.write("Input data:")
    st.dataframe(df_input)