import os
import pandas as pd
import joblib
import streamlit as st
from huggingface_hub import snapshot_download

# Model repository configuration
repo_identifier = "rakesh1248/random_forest_engine_condition_classifier"
model_file = "random_forest_model.joblib"

cache_dir = "./model_cache"
os.makedirs(cache_dir, exist_ok=True)

@st.cache_resource
def initialize_model():
    try:
        local_repo = snapshot_download(repo_id=repo_identifier, local_dir=cache_dir)
        model_location = os.path.join(local_repo, model_file)
        return joblib.load(model_location)
    except Exception as err:
        st.error(f"Model loading failed: {err}")
        st.stop()

model_instance = initialize_model()

# UI configuration
st.set_page_config(layout="wide")
st.title("Predictive Maintenance Solution for Engine Systems")

st.sidebar.header("Input Parameters")

rpm_val = st.sidebar.slider("Engine RPM", 60, 2300, 750)
oil_pressure = st.sidebar.slider("Lub Oil Pressure", 0.0, 8.0, 3.5, 0.1)
fuel_press = st.sidebar.slider("Fuel Pressure", 0.0, 22.0, 6.0, 0.1)
cool_press = st.sidebar.slider("Coolant Pressure", 0.0, 8.0, 2.0, 0.1)
oil_temp = st.sidebar.slider("Lub Oil Temperature", 70.0, 90.0, 78.0, 0.1)
cool_temp = st.sidebar.slider("Coolant Temperature", 60.0, 200.0, 80.0, 0.1)

input_frame = pd.DataFrame([{
    'Engine rpm': rpm_val,
    'Lub oil pressure': oil_pressure,
    'Fuel pressure': fuel_press,
    'Coolant pressure': cool_press,
    'lub oil temp': oil_temp,
    'Coolant temp': cool_temp
}])

st.write(input_frame)

if st.button("Predict"):
    pred = model_instance.predict(input_frame)
    prob = model_instance.predict_proba(input_frame)

    if pred[0] == 1:
        st.error("Faulty Engine")
    else:
        st.success("Normal Engine")

    st.write(f"Normal: {prob[0][0]:.2f}")
    st.write(f"Faulty: {prob[0][1]:.2f}")