import streamlit as st
import pandas as pd
import joblib

# 1. Page Configuration
st.set_page_config(page_title="Personality Predictor", page_icon="🧠")
st.title("Personality Insight Tool")
st.write("Enter social behavior metrics to predict personality type.")

# 2. Load Assets
# Cached for your 8GB RAM MacBook Pro
@st.cache_resource
def load_assets():
    model = joblib.load('personality_lr_model.joblib')
    num_imp = joblib.load('numeric_imputer.joblib')
    cat_imp = joblib.load('cat_imputer.joblib')
    return model, num_imp, cat_imp

try:
    best_lr, num_imputer, cat_imputer = load_assets()
except Exception as e:
    st.error(f"Error loading model files: {e}")

# 3. User Input Form
with st.form("prediction_form"):
    st.subheader("Social Metrics")
    # Organized to match your Training Data Order
    time_alone = st.number_input("Time Spent Alone (Hours)", min_value=0, max_value=24, value=5)
    stage_fear = st.selectbox("Do you have stage fear?", ["Yes", "No"])
    social_events = st.number_input("Social Events per Month", min_value=0, value=2)
    going_outside = st.number_input("Times Going Outside per Week", min_value=0, value=3)
    drained = st.selectbox("Do you feel drained after socializing?", ["Yes", "No"])
    friends_size = st.number_input("Friend Circle Size", min_value=0, value=5)
    post_freq = st.number_input("Social Media Post Frequency", min_value=0, value=1)
    
    submit = st.form_submit_button("Predict Personality")

# 4. Prediction Logic
if submit:
    # A. Create DataFrame in the EXACT order of your .info() output
    input_data = pd.DataFrame({
        'Time_spent_Alone': [time_alone],
        'Stage_fear': [stage_fear],
        'Social_event_attendance': [social_events],
        'Going_outside': [going_outside],
        'Drained_after_socializing': [drained],
        'Friends_circle_size': [friends_size],
        'Post_frequency': [post_freq]
    })

    # B. Define Column Groups for Imputation
    numeric_cols = ['Time_spent_Alone', 'Social_event_attendance', 'Going_outside', 'Friends_circle_size', 'Post_frequency']
    categorical_cols = ['Stage_fear', 'Drained_after_socializing']
    
    # C. Apply Imputers (Mirroring your 97% training accuracy logic)
    # Note: Scikit-learn transform handles the columns within input_data
    input_data[numeric_cols] = num_imputer.transform(input_data[numeric_cols])
    input_data[categorical_cols] = cat_imputer.transform(input_data[categorical_cols])
    
    # D. Map Categorical Strings to Numbers
    binary_mapping = {'Yes': 1, 'No': 0}
    input_data['Stage_fear'] = input_data['Stage_fear'].map(binary_mapping)
    input_data['Drained_after_socializing'] = input_data['Drained_after_socializing'].map(binary_mapping)
    
    # E. Predict
    # The order is already correct because of how we built input_data in step A
    prediction = best_lr.predict(input_data)
    
    # F. Map result (1: Introvert, 0: Extrovert)
    result = "Introvert" if prediction[0] == 1 else "Extrovert"
    
    # G. Display Result
    st.divider()
    if result == "Introvert":
        st.info(f"The predicted personality type is: **{result}**")
    else:
        st.success(f"The predicted personality type is: **{result}**")