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

# --- Configuration ---
REPO_ID_MODEL = "RajendrakumarPachaiappan/engine-predictive-model"
MODEL_FILE = "final_random_forest_model.joblib"
SCALER_FILE = "standard_scaler.joblib"

# The feature columns must match the order expected by the scaler (validated against joblib file)
FEATURE_COLS = ['Engine_RPM', 'Lub_Oil_Pressure', 'Fuel_Pressure',
                'Coolant_Pressure', 'Lub_Oil_Temperature', 'Coolant_Temperature']

# --- Resource Loading Function ---
@st.cache_resource(show_spinner=False) # Suppress default spinner since we use custom status messages
def load_model_and_scaler():
    """
    Downloads and loads the model and scaler from Hugging Face Hub.
    Does NOT use st. commands inside to avoid initial warnings.
    """
    try:
        # Download files from Hugging Face
        model_path = hf_hub_download(repo_id=REPO_ID_MODEL, filename=MODEL_FILE)
        model = joblib.load(model_path)
        scaler_path = hf_hub_download(repo_id=REPO_ID_MODEL, filename=SCALER_FILE)
        scaler = joblib.load(scaler_path)
        return model, scaler
    except Exception as e:
        # Re-raise a descriptive exception for the main script to catch
        raise Exception(f"Failed to load required artifacts: {e}")

# --- Streamlit UI and Prediction Logic ---

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

# 1. Load Resources and Display Status
st.info("Loading predictive model and scaler from Hugging Face Hub...")
try:
    # This call triggers the download/caching
    model, scaler = load_model_and_scaler()
    st.success("Artifacts loaded successfully! Ready for prediction.")

    st.markdown("Use the sliders to simulate real-time sensor data and predict the **Engine Condition** (0=Healthy, 1=Faulty).")

except Exception as e:
    # Display error and halt execution if resources fail to load
    st.error(f"🔴 Error loading resources: {e}")
    st.stop()


# 2. Input Sliders
# Use columns for a cleaner layout
col1, col2, col3 = st.columns(3)

with col1:
    Engine_RPM = st.slider("Engine RPM (rev/min)", min_value=60, max_value=2300, value=791, step=10)
    Lub_Oil_Pressure = st.slider("Lub Oil Pressure (bar)", min_value=0.0, max_value=7.3, value=3.3, step=0.1)
    Fuel_Pressure = st.slider("Fuel Pressure (bar)", min_value=0.0, max_value=22.0, value=6.7, step=0.1)

with col2:
    Coolant_Pressure = st.slider("Coolant Pressure (bar)", min_value=0.0, max_value=7.5, value=2.3, step=0.1)
    Lub_Oil_Temperature = st.slider("Lub Oil Temp (°C)", min_value=71.0, max_value=90.0, value=78.0, step=0.1)
    Coolant_Temperature = st.slider("Coolant Temp (°C)", min_value=60.0, max_value=200.0, value=78.5, step=0.5)

# 3. Prediction
if st.button("Predict Engine Condition", type="primary"):

    # a. Prepare Input
    input_data = pd.DataFrame({
        'Engine_RPM': [Engine_RPM],
        'Lub_Oil_Pressure': [Lub_Oil_Pressure],
        'Fuel_Pressure': [Fuel_Pressure],
        'Coolant_Pressure': [Coolant_Pressure],
        'Lub_Oil_Temperature': [Lub_Oil_Temperature],
        'Coolant_Temperature': [Coolant_Temperature]
    })

    # b. Scale Input
    # Important: Use FEATURE_COLS to ensure correct order for the scaler
    input_scaled = scaler.transform(input_data[FEATURE_COLS])

    # c. Make Prediction
    prediction = model.predict(input_scaled)[0]
    # Get the probability of the *faulty* class (1)
    prediction_proba = model.predict_proba(input_scaled)[:, 1][0] 

    # d. Display Results
    st.divider()
    st.subheader("Prediction Result:")
    
    if prediction == 1:
        st.error(f"Engine Condition: **FAULTY** (Probability of Fault: {prediction_proba:.2f}) ⚠️")
        st.write("Immediate maintenance is recommended to prevent breakdown.")
    else:
        st.success(f"Engine Condition: **HEALTHY** (Probability of Fault: {prediction_proba:.2f}) ✅")
        st.write("Engine is operating normally. Continue regular monitoring.")

    st.caption("Note: Probability of fault close to 0.5 indicates uncertainty.")