import streamlit as st
import pandas as pd
import joblib
import numpy as np

# Page Configuration
st.set_page_config(
    page_title="Mushroom Edibility Predictor",
    page_icon="🍄",
    layout="centered")

# Title and Description
st.title("🍄 Mushroom Edibility Predictor")
st.markdown("""
This app predicts whether a mushroom is **Edible** or **Poisonous** based on its physical characteristics.
""")

st.divider()

# Load Model 
@st.cache_resource
def load_artifacts():
    try:
        return joblib.load('src/mushroom_pipeline.pkl')
    except FileNotFoundError:
        st.error("Model file 'mushroom_pipeline.pkl' not found. Please upload it to the repository.")
        return None

artifacts = load_artifacts()

if artifacts:
    model = artifacts['model']
    encoders = artifacts['encoders']
    valid_labels = artifacts['valid_labels']
    cat_cols = artifacts['cat_cols']
    num_cols = artifacts['num_cols']
    imputer = artifacts['num_imputer']

    # User Input 
    st.sidebar.header("🍄 Mushroom Features")
    st.sidebar.write("Please select the characteristics below:")
    
    input_data = {}

    for col in cat_cols:
        options = sorted(valid_labels[col])
        label = col.replace('-', ' ').title()
        input_data[col] = st.sidebar.selectbox(f"{label}", options)

    for col in num_cols:
        label = col.replace('-', ' ').title()
        input_data[col] = st.sidebar.number_input(f"{label}", value=0.0)

    st.subheader("Results")
    
    if st.button("Predict Edibility 🔍", type="primary"):
        
        df = pd.DataFrame([input_data])
        
        df[num_cols] = imputer.transform(df[num_cols])
        
        for col in cat_cols:
            val = df.loc[0, col]
            
            if val not in valid_labels[col]:
                val = 'Other'
                
            le = encoders[col]
            try:
                transformed_val = le.transform([val])[0]
                df.loc[0, col] = transformed_val
            except:
                fallback_val = le.classes_[0] 
                df.loc[0, col] = le.transform([fallback_val])[0]
                
        for col in cat_cols:
            df[col] = df[col].astype(int)
        
        for col in num_cols:
            df[col] = df[col].astype(float)

        df = df[model.feature_names_in_]
        
        prediction = model.predict(df)[0]
        probability = model.predict_proba(df)[0][1] # Probability of being Poisonous

        if prediction == 1: # 1: Poisonous
            st.error(f"### ⚠️ POISONOUS 🍄")
            st.write(f"**Confidence:** {probability*100:.2f}% chance of being poisonous.")
            st.warning("Do not eat this mushroom!")
        else: # 0: Edible
            st.success(f"### ✅ EDIBLE 🍄")
            st.write(f"**Confidence:** {(1-probability)*100:.2f}% chance of being edible.")
            st.balloons()

else:
    st.stop()