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

# ------------------------------------------------------------
# Deployment configuration
# ------------------------------------------------------------
# MODEL_REPO_ID can be overridden in Hugging Face Space secrets/variables.
MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "premswan/engine-predictive-maintenance-model")
MODEL_FILENAME = "best_engine_maintenance_model.joblib"
FEATURE_COLUMNS = [
    "Engine_RPM",
    "Lub_Oil_Pressure",
    "Fuel_Pressure",
    "Coolant_Pressure",
    "Lub_Oil_Temperature",
    "Coolant_Temperature"
]
FEATURE_DEFAULTS = {
    "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
}
FEATURE_HELP = {
    "Engine_RPM": "Engine speed in revolutions per minute.",
    "Lub_Oil_Pressure": "Lubricating oil pressure reading.",
    "Fuel_Pressure": "Fuel pressure reading.",
    "Coolant_Pressure": "Coolant pressure reading.",
    "Lub_Oil_Temperature": "Lubricating oil temperature in Celsius.",
    "Coolant_Temperature": "Coolant temperature in Celsius."
}

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

@st.cache_resource
def load_model():
    # Load the registered model from Hugging Face Model Hub.
    token = os.getenv("HF_TOKEN")
    model_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=MODEL_FILENAME,
        token=token
    )
    return joblib.load(model_path)

st.title("Engine Predictive Maintenance")
st.write(
    "Enter engine sensor readings to predict whether the engine is operating normally "
    "or may require maintenance."
)

model = load_model()

# ------------------------------------------------------------
# Capture input readings and save them into a dataframe
# ------------------------------------------------------------
input_values = {}
for feature in FEATURE_COLUMNS:
    input_values[feature] = st.number_input(
        label=feature,
        value=float(FEATURE_DEFAULTS.get(feature, 0.0)),
        help=FEATURE_HELP.get(feature, "Enter sensor value")
    )

input_df = pd.DataFrame([input_values], columns=FEATURE_COLUMNS)
input_df.to_csv("latest_input.csv", index=False)

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 the engine before continuing heavy operation.")
    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%}")