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

st.set_page_config(page_title="Engine Maintenance Predictor", layout="wide")
st.title("🚢 Engine Predictive Maintenance")

# 1. Load the saved model from the Hugging Face model hub
@st.cache_resource
def load_my_model():
    # Update this to match your registration details: f"{HF_USERNAME}/engine-condition-predictor"
    # Example: "ranasinghnaveen/engine-condition-predictor"
    REPO_ID = "ranasinghnaveen/engine-condition-predictor" 
    FILENAME = "engine_model.joblib"
    
    try:
        model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
        return joblib.load(model_path)
    except Exception as e:
        st.error(f"Error loading model from Hugging Face: {e}")
        return None

model = load_my_model()

# 2. Get the inputs from user
st.sidebar.header("User Input Features")
def user_input_features():
    # These inputs correspond to the dataframe columns expected by your model
    rpm = st.sidebar.number_input("Engine RPM", value=800.0)
    oil_p = st.sidebar.number_input("Lub Oil Pressure", value=3.5)
    fuel_p = st.sidebar.number_input("Fuel Pressure", value=6.5)
    cool_p = st.sidebar.number_input("Coolant Pressure", value=2.5)
    oil_t = st.sidebar.number_input("Lub Oil Temp", value=78.0)
    cool_t = st.sidebar.number_input("Coolant Temp", value=80.0)
    
    # 3. Save them into a dataframe
    data = {
        'Engine_RPM': rpm,
        'Lub_Oil_Pressure': oil_p,
        'Fuel_Pressure': fuel_p,
        'Coolant_Pressure': cool_p,
        'Lub_Oil_Temp': oil_t,
        'Coolant_Temp': cool_t
    }
    return pd.DataFrame(data, index=[0])

df = user_input_features()

# Display input parameters
st.subheader('User Input Parameters')
st.write(df)

# Prediction Logic
if st.button("Predict"):
    if model is not None:
        prediction = model.predict(df)
        # Based on your class balance: 0 = Maintenance Required, 1 = Normal
        result = "Maintenance Required" if prediction[0] == 0 else "Normal"
        
        if prediction[0] == 0:
            st.error(f"Prediction: {result}")
        else:
            st.success(f"Prediction: {result}")
    else:
        st.warning("Model not loaded. Please check the repository ID.")