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

# Download and load the model
model_path = hf_hub_download(repo_id="JohnsonSAimlarge/engine-failure-predict", filename="engine_failure_model.joblib")
model = joblib.load(model_path)

# ------------------------------
# Streamlit UI
# ------------------------------
st.title("🔧 Engine Failure Prediction System")
st.write("""
This application predicts the likelihood of engine failure based on sensor readings and operational parameters.
Please enter **Engine Sensor Data** below to get a prediction.
""")

# ------------------------------
# User Inputs
# ------------------------------
st.subheader("Engine Operational Parameters")

col1, col2 = st.columns(2)

with col1:
    engine_rpm = st.number_input(
        "Engine RPM (Revolutions Per Minute)", 
        min_value=0, 
        max_value=10000, 
        value=3000,
        help="Normal range: 500-8000 RPM"
    )
    
    lub_oil_pressure = st.number_input(
        "Lubricating Oil Pressure (bar)", 
        min_value=0.0, 
        max_value=10.0, 
        value=4.5,
        step=0.1,
        help="Normal range: 2.0-6.0 bar"
    )
    
    fuel_pressure = st.number_input(
        "Fuel Pressure (bar)", 
        min_value=0.0, 
        max_value=10.0, 
        value=4.0,
        step=0.1,
        help="Normal range: 2.0-6.0 bar"
    )

with col2:
    coolant_pressure = st.number_input(
        "Coolant Pressure (bar)", 
        min_value=0.0, 
        max_value=5.0, 
        value=2.5,
        step=0.1,
        help="Normal range: 1.5-3.5 bar"
    )
    
    lub_oil_temp = st.number_input(
        "Lubricating Oil Temperature (°C)", 
        min_value=0, 
        max_value=200, 
        value=75,
        help="Normal range: 50-120°C"
    )
    
    coolant_temp = st.number_input(
        "Coolant Temperature (°C)", 
        min_value=0, 
        max_value=150, 
        value=80,
        help="Normal range: 60-100°C"
    )

# ------------------------------
# Prepare Input for Prediction
# ------------------------------
input_data = {
    "Engine rpm": engine_rpm,
    "Lub oil pressure": lub_oil_pressure,
    "Fuel pressure": fuel_pressure,
    "Coolant pressure": coolant_pressure,
    "lub oil temp": lub_oil_temp,
    "Coolant temp": coolant_temp
}

input_df = pd.DataFrame([input_data])

# Display input summary
st.subheader("Input Summary")
st.dataframe(input_df, use_container_width=True)

# ------------------------------
# Prediction
# ------------------------------
if st.button("🔍 Predict Engine Condition", type="primary"):
    try:
        prediction = model.predict(input_df)[0]
        probability = model.predict_proba(input_df)[0][1]
        
        # Use custom threshold for imbalanced dataset
        # Adjust based on your model's optimal threshold
        classification_threshold = 0.5
        prediction = (probability >= classification_threshold).astype(int)
        
        st.markdown("---")
        st.subheader("Prediction Results")
        
        if prediction == 1:
            st.error(f"⚠️ **ENGINE FAILURE PREDICTED** - Immediate maintenance required!")
            st.error(f"**Failure Probability: {probability:.2%}**")
            st.warning("""
            **Recommended Actions:**
            - Stop engine operation immediately
            - Conduct thorough inspection
            - Check all sensor readings
            - Consult maintenance team
            """)
        else:
            st.success(f"✅ **ENGINE CONDITION NORMAL** - No immediate action required")
            st.success(f"**Failure Probability: {probability:.2%}**")
            st.info("""
            **Maintenance Recommendations:**
            - Continue regular monitoring
            - Schedule routine maintenance as planned
            - Keep monitoring sensor readings
            """)
        
        # Display confidence meter
        st.subheader("Confidence Level")
        confidence = max(probability, 1 - probability)
        st.progress(confidence)
        st.write(f"Model Confidence: {confidence:.2%}")
        
    except Exception as e:
        st.error(f"Error during prediction: {str(e)}")
        st.info("Please check your input values and try again.")

# ------------------------------
# Additional Information
# ------------------------------
with st.expander("ℹ️ About This Model"):
    st.write("""
    **Model Information:**
    - Algorithm: XGBoost with SMOTE for class balancing
    - Test Accuracy: 64.42%
    - Precision: 76.42%
    - Recall: 63.01%
    - Dataset: 19,535 engine records
    
    **Most Important Features:**
    1. Engine RPM (38.3%)
    2. Fuel Pressure (16.2%)
    3. Oil Temperature (13.7%)
    
    **Model Repository:** [JohnsonSAimlarge/engine-failure-predictor](https://huggingface.co/JohnsonSAimlarge/engine-failure-predictor)
    """)

with st.expander("📊 Feature Ranges & Guidelines"):
    st.write("""
    | Parameter | Normal Range | Critical Threshold |
    |-----------|--------------|-------------------|
    | Engine RPM | 500-8000 | >8000 or <500 |
    | Lub Oil Pressure | 2.0-6.0 bar | <2.0 or >6.0 |
    | Fuel Pressure | 2.0-6.0 bar | <2.0 or >6.0 |
    | Coolant Pressure | 1.5-3.5 bar | <1.5 or >3.5 |
    | Lub Oil Temp | 50-120°C | >120°C |
    | Coolant Temp | 60-100°C | >100°C |
    """)

# Footer
st.markdown("---")
st.caption("Engine Failure Prediction System | Powered by XGBoost & Hugging Face")