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

st.set_page_config(page_title="Predictive Maintenance – Engine Health", layout="centered")

st.title("Predictive Maintenance – Engine Health")
st.write("Enter engine sensor readings to predict whether maintenance is needed.")

MODEL_REPO = "SabarnaDeb/Capstone_PredictiveMaintenance_Model"
MODEL_FILE = "model.joblib"

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

model = load_model()

FEATURES = [
    "engine_rpm",
    "lub_oil_pressure",
    "fuel_pressure",
    "coolant_pressure",
    "lub_oil_temp",
    "coolant_temp"
]


engine_rpm = st.number_input("Engine RPM", value=800.0)
lub_oil_pressure = st.number_input("Lub Oil Pressure", value=4.0)
fuel_pressure = st.number_input("Fuel Pressure", value=6.5)
coolant_pressure = st.number_input("Coolant Pressure", value=3.5)
lub_oil_temperature = st.number_input("Lub Oil Temperature", value=80.0)
coolant_temperature = st.number_input("Coolant Temperature", value=85.0)

if st.button("Predict"):
    input_df = 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_temperature,
        "coolant_temp": coolant_temperature,
    }])

    pred = int(model.predict(input_df[FEATURES])[0])

    prob = None
    if hasattr(model, "predict_proba"):
        prob = float(model.predict_proba(input_df[FEATURES])[:, 1][0])

    st.subheader("Prediction Result")
    if pred == 1:
        st.error("⚠️ Maintenance Needed")
    else:
        st.success("✅ Normal Operation")

    if prob is not None:
        st.write(f"Confidence (maintenance probability): **{prob:.2f}**")

    st.subheader("Input Data (saved as DataFrame)")
    st.dataframe(input_df)