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

# --- Hugging Face Model Repository Details ---
REPO_ID_MODEL = "RajendrakumarPachaiappan/engine-predictive-model"
# Assuming the best model found in the notebook was Random Forest based on your script
MODEL_FILENAME = "final_random_forest_model.joblib" 
SCALER_FILENAME = "standard_scaler.joblib"

# The six feature columns for input
FEATURE_COLS = [
    'Engine_RPM', 'Lub_Oil_Pressure', 'Fuel_Pressure',
    'Coolant_Pressure', 'Lub_Oil_Temperature', 'Coolant_Temperature'
]

# --- Caching Function to Load Model and Scaler ---
@st.cache_resource
def load_artifacts():
    """
    Downloads and loads the model and scaler from the Hugging Face Hub.
    The @st.cache_resource decorator ensures this function only runs once.
    """
    try:
        with st.spinner("Downloading and loading model artifacts..."):
            # Download and load the model
            model_path = hf_hub_download(repo_id=REPO_ID_MODEL, filename=MODEL_FILENAME, repo_type="model")
            model = joblib.load(model_path)

            # Download and load the scaler
            scaler_path = hf_hub_download(repo_id=REPO_ID_MODEL, filename=SCALER_FILENAME, repo_type="model")
            scaler = joblib.load(scaler_path)

        return model, scaler
    except Exception as e:
        st.error(f"CRITICAL ERROR: Failed to load model or scaler. Check repository ID and filenames. Error: {e}")
        return None, None

# Load the model and scaler
model, scaler = load_artifacts()

# --- Streamlit UI Setup ---
st.set_page_config(
    page_title="Engine Predictive Maintenance",
    layout="centered",
    initial_sidebar_state="expanded"
)

st.title("🚗 Engine Predictive Maintenance Model")
st.markdown(
    "Use the sliders below to input the current engine sensor readings and predict "
    "if the engine is running **NORMAL (0)** or requires **IMMEDIATE MAINTENANCE (1)**. "
    "This model uses a pre-trained Random Forest classifier."
)
st.markdown(
    "**Note:** The slider ranges are based on the minimum and maximum values observed in the training dataset (19,535 records)."
)

if model is None or scaler is None:
    st.stop() # Stop the app if artifacts failed to load

# --- Input Components ---
st.header("Sensor Readings")


# Define ranges based on the provided dataset statistics
# Format: (min_value, max_value, mean_value/default) - all values are floats.
ranges = {
    'Engine_RPM': (61.0, 2239.0, 791.0), 
    'Lub_Oil_Pressure': (0.0, 7.3, 3.3),
    'Fuel_Pressure': (0.0, 21.1, 6.7),
    'Coolant_Pressure': (0.0, 7.5, 2.3),
    'Lub_Oil_Temperature': (71.3, 89.6, 77.6),
    'Coolant_Temperature': (61.7, 195.5, 78.4),
}

input_values = {}
# Use Streamlit columns for a cleaner, side-by-side layout
col1, col2 = st.columns(2)
columns = [col1, col2]

for i, col_name in enumerate(FEATURE_COLS):
    # Determine which column to place the widget in
    current_col = columns[i % 2] 

    min_val, max_val, default_val = ranges[col_name]
    
    # Clean up names for display in the UI and set units
    label = col_name.replace('_', ' ')
    unit = ""
    if "RPM" in col_name:
        unit = " (rev/min)"
    elif "Pressure" in col_name:
        unit = " (bar/kPa)"
    elif "Temperature" in col_name:
        unit = " (°C)"


    with current_col:
        input_values[col_name] = st.slider(
            label=f"{label}{unit}",
            min_value=min_val,
            max_value=max_val,
            value=default_val,
            step=0.1,
            help=f"Current reading for {label}. Full data range: [{min_val}, {max_val}]"
        )

# --- Prediction Logic ---
if st.button("Predict Engine Condition", type="primary"):
    # 1. Prepare data for the model
    # Convert input_values dictionary to a DataFrame row
    input_df = pd.DataFrame([input_values], columns=FEATURE_COLS)

    # 2. Scale the input data using the pre-trained StandardScaler
    # The scaler must be fitted on the same data distribution as the model
    input_scaled = scaler.transform(input_df)

    # 3. Make prediction
    # The output is a numpy array, we take the first element (the prediction)
    prediction = model.predict(input_scaled)[0]
    
    # 4. Display results
    st.subheader("Prediction Result")
    
    if prediction == 1:
        st.error(
            "🔴 FAULTY (1): Immediate Maintenance Required! "
            "High probability of engine failure detected. Check for high RPM, low pressures, or extreme temperatures."
        )
    else:
        st.success(
            "🟢 NORMAL (0): Operating within expected parameters. "
            "Engine health is currently good."
        )
        
    st.caption(f"Raw Model Prediction (0=Normal, 1=Faulty): {prediction}")