import streamlit as st
import pandas as pd
import joblib

# -----------------------------
# Page Config
# -----------------------------
st.set_page_config(page_title="Engine Condition App", layout="wide")

st.title("🚗 Engine Condition Classification")

# -----------------------------
# Load Model
# -----------------------------
@st.cache_resource
def load_model():
    model = joblib.load("best_model.pkl")
    return model

model = load_model()

# ✅ IMPORTANT: use feature names from model training
feature_names = model.feature_names_in_

# -----------------------------
# SINGLE PREDICTION
# -----------------------------
st.header("🔧 Manual Prediction")

input_values = []

cols = st.columns(3)

for i, feature in enumerate(feature_names):
    with cols[i % 3]:
        value = st.number_input(feature, value=0.0)
        input_values.append(value)

if st.button("Predict Engine Condition"):

    input_df = pd.DataFrame([input_values], columns=feature_names)

    prediction = model.predict(input_df)[0]
    probability = model.predict_proba(input_df)[0][1]

    if prediction == 1:
        st.error("⚠ Engine Needs Maintenance")
    else:
        st.success("✅ Engine Operating Normally")

    st.write(f"**Failure Probability:** {probability:.3f}")

# -----------------------------
# BATCH PREDICTION
# -----------------------------
st.header("📂 Batch Prediction")

uploaded_file = st.file_uploader("Upload CSV file", type=["csv"])

if uploaded_file is not None:
    try:
        df = pd.read_csv(uploaded_file)

        # Safety row limit (prevents Space crash)
        if len(df) > 20000:
            st.warning("⚠ File too large. Max 10,000 rows allowed.")
        else:
            # Check missing columns
            missing_cols = [col for col in feature_names if col not in df.columns]

            if missing_cols:
                st.error(f"Missing required columns: {missing_cols}")
            else:
                df = df[feature_names]

                probabilities = model.predict_proba(df)[:, 1]
                df["Failure_Probability"] = probabilities
                df["Prediction"] = (probabilities >= 0.5).astype(int)

                st.success("✅ Predictions completed!")

                st.dataframe(df.head())

                csv = df.to_csv(index=False).encode("utf-8")

                st.download_button(
                    label="Download Results",
                    data=csv,
                    file_name="engine_predictions.csv",
                    mime="text/csv",
                )

    except Exception as e:
        st.error(f"Error processing file: {e}")