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

# Configuration
HF_TOKEN = os.getenv("HF_TOKEN")  # optional if model is public
HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "dhani10/engine-condition-model")
MODEL_FILE = os.getenv("MODEL_FILE", "best_engine_model.joblib")

# Use a writable cache (esp. on Spaces)
os.environ.setdefault("HF_HOME", "/tmp/huggingface")
os.environ.setdefault("HF_HUB_CACHE", "/tmp/huggingface/hub")
os.makedirs(os.environ["HF_HUB_CACHE"], exist_ok=True)

if HF_TOKEN:
    try:
        login(HF_TOKEN)
    except Exception:
        pass

@st.cache_resource
def load_model():
    p = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=MODEL_FILE,
        repo_type="model",
        token=HF_TOKEN,
        cache_dir=os.environ["HF_HUB_CACHE"],
    )
    return joblib.load(p)

model = load_model()

# Engine sensor features
ENGINE_FEATURES = [
    "Engine rpm",
    "Lub oil pressure", 
    "Fuel pressure",
    "Coolant pressure",
    "lub oil temp",
    "Coolant temp"
]

st.set_page_config(page_title="Engine Condition Monitor", layout="centered")
st.title("🏭 Engine Condition Monitoring System")
st.caption("Enter sensor readings to predict engine condition (Normal/Faulty)")

# Add information about the model
with st.expander("ℹ️ About this Model"):
    st.write('''
    This model predicts engine condition based on real-time sensor readings:
    - **Normal (0)**: Engine operating within normal parameters
    - **Faulty (1)**: Engine showing signs of potential failure
    
    **Typical Operating Ranges:**
    - Engine RPM: 600-2500
    - Lub Oil Pressure: 2.0-4.0 bar
    - Fuel Pressure: 8.0-15.0 bar  
    - Coolant Pressure: 1.5-4.0 bar
    - Lub Oil Temp: 75-110°C
    - Coolant Temp: 70-100°C
    ''')

# Sensor input form
with st.form("engine_predict_form"):
    st.subheader("🔧 Sensor Readings")
    
    col1, col2 = st.columns(2)
    
    with col1:
        engine_rpm = st.slider(
            "Engine RPM", 
            min_value=0, 
            max_value=3000, 
            value=1800,
            help="Engine rotations per minute"
        )
        
        lub_oil_pressure = st.slider(
            "Lub Oil Pressure (bar)", 
            min_value=0.0, 
            max_value=6.0, 
            value=3.1,
            step=0.1,
            help="Lubricating oil pressure in bar"
        )
        
        fuel_pressure = st.slider(
            "Fuel Pressure (bar)", 
            min_value=0.0, 
            max_value=20.0, 
            value=12.0,
            step=0.1,
            help="Fuel system pressure in bar"
        )
    
    with col2:
        coolant_pressure = st.slider(
            "Coolant Pressure (bar)", 
            min_value=0.0, 
            max_value=5.0, 
            value=2.9,
            step=0.1,
            help="Cooling system pressure in bar"
        )
        
        lub_oil_temp = st.slider(
            "Lub Oil Temp (°C)", 
            min_value=0, 
            max_value=150, 
            value=92,
            help="Lubricating oil temperature in °C"
        )
        
        coolant_temp = st.slider(
            "Coolant Temp (°C)", 
            min_value=0, 
            max_value=150, 
            value=89,
            help="Coolant temperature in °C"
        )
    
    submitted = st.form_submit_button("🔍 Predict Engine Condition")

if submitted:
    # Build input data
    ui_row = {
        "Engine rpm": float(engine_rpm),
        "Lub oil pressure": float(lub_oil_pressure),
        "Fuel pressure": float(fuel_pressure),
        "Coolant pressure": float(coolant_pressure),
        "lub oil temp": float(lub_oil_temp),
        "Coolant temp": float(coolant_temp)
    }

    # Create DataFrame with expected columns
    row = pd.DataFrame([ui_row])

    try:
        # Make prediction
        prediction = model.predict(row)[0]
        probability = None
        if hasattr(model, "predict_proba"):
            probability = model.predict_proba(row)[0]

        # Display results
        st.subheader("🎯 Prediction Result")
        
        if prediction == 1:
            st.error("🚨 **FAULTY CONDITION DETECTED**")
            st.warning("⚠️ Engine shows signs of potential failure. Immediate maintenance recommended.")
            
            if probability is not None:
                st.metric(
                    "Confidence Score", 
                    f"{probability[1]:.1%}",
                    delta=f"Faulty probability",
                    delta_color="inverse"
                )
        else:
            st.success("✅ **NORMAL OPERATION**")
            st.info("🌡️ Engine operating within normal parameters.")
            
            if probability is not None:
                st.metric(
                    "Confidence Score", 
                    f"{probability[0]:.1%}",
                    delta=f"Normal probability",
                    delta_color="normal"
                )

        # Display detailed probabilities
        if probability is not None:
            col1, col2 = st.columns(2)
            with col1:
                st.progress(probability[0], text=f"Normal: {probability[0]:.1%}")
            with col2:
                st.progress(probability[1], text=f"Faulty: {probability[1]:.1%}")

        # Show input values
        with st.expander("📊 Sensor Readings Used"):
            st.dataframe(row.T.rename(columns={0: "Value"}))

        # Add maintenance recommendations for faulty conditions
        if prediction == 1:
            st.subheader("🔧 Recommended Actions")
            issues = []
            
            if lub_oil_pressure < 2.5:
                issues.append("Low lubricating oil pressure")
            if fuel_pressure > 13.0:
                issues.append("High fuel pressure")
            if lub_oil_temp > 105:
                issues.append("High lubricating oil temperature")
            if coolant_temp > 95:
                issues.append("High coolant temperature")
            
            if issues:
                st.write("Potential issues detected:")
                for issue in issues:
                    st.write(f"• {issue}")
            
            st.write('''
            **Immediate Steps:**
            1. Check oil levels and quality
            2. Inspect cooling system
            3. Verify fuel system components
            4. Consult maintenance manual
            ''')

    except Exception as e:
        st.error(f"❌ Prediction failed: {e}")
        st.write("Expected features:", ENGINE_FEATURES)

# Add footer
st.markdown("---")
st.caption('''
**Engine Condition Prediction System** | 
Predictive Maintenance Model | 
[View Model on Hugging Face](https://huggingface.co/dhani10/engine-condition-model)
''')