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

DEFAULT_MODEL_REPO_ID = "premswan/engine-predictive-maintenance-model"
MODEL_FILE = "best_engine_maintenance_model.joblib"
METADATA_FILE = "model_metadata.json"

st.set_page_config(page_title="Engine Predictive Maintenance", page_icon="W", layout="centered")

@st.cache_resource
def load_model_and_metadata(model_repo_id: str = DEFAULT_MODEL_REPO_ID):
    # Load model and metadata from Hugging Face Model Hub.
    model_path = hf_hub_download(repo_id=model_repo_id, filename=MODEL_FILE)
    metadata_path = hf_hub_download(repo_id=model_repo_id, filename=METADATA_FILE)
    model = joblib.load(model_path)
    with open(metadata_path, "r", encoding="utf-8") as file:
        metadata = json.load(file)
    return model, metadata

model, metadata = load_model_and_metadata()
feature_columns = metadata.get("feature_columns", ['Engine_RPM', 'Lub_Oil_Pressure', 'Fuel_Pressure', 'Coolant_Pressure', 'Lub_Oil_Temperature', 'Coolant_Temperature'])

st.title("Engine Predictive Maintenance App")
st.write("Enter engine sensor readings to predict whether the engine is normal or needs maintenance attention.")
st.subheader("Sensor Inputs")

default_values = {
    "Engine_RPM": 800.0,
    "Lub_Oil_Pressure": 3.2,
    "Fuel_Pressure": 6.5,
    "Coolant_Pressure": 2.4,
    "Lub_Oil_Temperature": 78.0,
    "Coolant_Temperature": 80.0,
}

user_inputs = {}
for feature in feature_columns:
    user_inputs[feature] = st.number_input(
        label=feature,
        value=float(default_values.get(feature, 0.0)),
        step=0.1,
        format="%.4f"
    )

# Save inputs into a dataframe as required by the deployment rubric.
input_df = pd.DataFrame([user_inputs], columns=feature_columns)
st.subheader("Input DataFrame")
st.dataframe(input_df, use_container_width=True)

if st.button("Predict Engine Condition"):
    prediction = int(model.predict(input_df)[0])
    probability_maintenance = None
    if hasattr(model, "predict_proba"):
        probability_maintenance = float(model.predict_proba(input_df)[0, 1])

    if prediction == 1:
        st.error("Prediction: Maintenance / Faulty condition")
        st.write("Recommended action: inspect engine health and schedule preventive maintenance.")
    else:
        st.success("Prediction: Normal / Healthy condition")
        st.write("Recommended action: continue normal monitoring.")

    if probability_maintenance is not None:
        st.metric("Maintenance Probability", f"{probability_maintenance:.2%}")

    st.write("Raw prediction output:", prediction)

st.caption("Model loaded from Hugging Face Model Hub: " + DEFAULT_MODEL_REPO_ID)