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

# 1. SETUP & CONFIGURATION
st.set_page_config(
    page_title="Mental Health Prediction",
    page_icon="🧠",
    layout="centered")

# Load the trained model and tools
@st.cache_resource
def load_model_pipeline():
    pipeline = joblib.load('src/depression_model.pkl')
    return pipeline

try:
    data = load_model_pipeline()
    model = data["model"]
    encoders = data["encoders"]
    sleep_mapping = data["sleep_mapping"]
except FileNotFoundError:
    st.error("Model file not found! Please upload 'depression_model_pipeline.pkl'.")
    st.stop()

# 2. UI DESIGN (Sidebar & Main Inputs)
st.title("🧠 Mental Health Prediction")
st.markdown("""
    This application analyze lifestyle and demographic factors to predict the likelihood of depression.
""")

with st.form("prediction_form"):
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Personal Info")
        gender = st.selectbox("Gender", ["Male", "Female"])
        age = st.number_input("Age", min_value=10, max_value=100, value=25)
        city = st.text_input("City", "Visakhapatnam") # Default to a known city or handle 'Unknown'
        
        st.subheader("Work & Study")
        occupation = st.selectbox("Occupation Status", ["Student", "Working Professional", "Retired", "Unemployed"])
        profession = st.text_input("Profession (e.g. Engineer, Student)", "Student")
        work_study_hours = st.slider("Work/Study Hours (per day)", 0.0, 16.0, 8.0)
        
    with col2:
        st.subheader("Well-being & Habits")
        sleep_options = list(sleep_mapping.keys())
        display_sleep = sorted([k for k in sleep_options if isinstance(k, str) and len(k) < 20])
        sleep_input = st.selectbox("Sleep Duration", display_sleep, index=display_sleep.index("7-8 hours") if "7-8 hours" in display_sleep else 0)
        
        dietary = st.selectbox("Dietary Habits", ["Healthy", "Moderate", "Unhealthy"])
        degree = st.selectbox("Degree Level", ["Undergraduate", "Postgraduate", "PhD", "High School", "Class 12"])
        
    st.subheader("Self Assessment")
    
    c1, c2, c3 = st.columns(3)
    with c1:
        academic_pressure = st.slider("Academic Pressure (0-5)", 0.0, 5.0, 0.0)
        work_pressure = st.slider("Work Pressure (0-5)", 0.0, 5.0, 0.0)
        cgpa = st.number_input("CGPA (if student)", 0.0, 10.0, 0.0)
    
    with c2:
        study_satisfaction = st.slider("Study Satisfaction (0-5)", 0.0, 5.0, 0.0)
        job_satisfaction = st.slider("Job Satisfaction (0-5)", 0.0, 5.0, 0.0)
        financial_stress = st.slider("Financial Stress (0-5)", 0.0, 5.0, 0.0)
        
    with c3:
        suicidal_thoughts = st.selectbox("History of Suicidal Thoughts", ["Yes", "No"])
        family_history = st.selectbox("Family History of Mental Illness", ["Yes", "No"])

    submitted = st.form_submit_button("Analyze Result")

# 3. PREDICTION LOGIC
if submitted:
    
    try:
        def safe_encode(encoder, value):
            try:
                return encoder.transform([str(value)])[0]
            except ValueError:
                # If unseen category, fallback to mode or first class (Basic handling)
                return 0 

        sleep_val = sleep_mapping.get(sleep_input, 7.0) # Default to 7 if error

        input_data = pd.DataFrame({
            'gender': [safe_encode(encoders['gender'], gender)],
            'age': [float(age)],
            'city': [safe_encode(encoders['city'], city)],
            'occupation_status': [safe_encode(encoders['occupation_status'], occupation)],
            'profession': [safe_encode(encoders['profession'], profession)],
            'academic_pressure': [float(academic_pressure)],
            'work_pressure': [float(work_pressure)],
            'cgpa': [float(cgpa)],
            'study_satisfaction': [float(study_satisfaction)],
            'job_satisfaction': [float(job_satisfaction)],
            'sleep_duration': [float(sleep_val)],
            'dietary_habits': [safe_encode(encoders['dietary_habits'], dietary)],
            'degree': [safe_encode(encoders['degree'], degree)],
            'suicidal_thoughts': [safe_encode(encoders['suicidal_thoughts'], suicidal_thoughts)],
            'work_study_hours': [float(work_study_hours)],
            'financial_stress': [float(financial_stress)],
            'family_history_mental_illness': [safe_encode(encoders['family_history_mental_illness'], family_history)]})

        prediction_prob = model.predict_proba(input_data)[0][1] # Probability of Class 1
        prediction_class = (prediction_prob > 0.5).astype(int)

        # 4. DISPLAY RESULTS
        if prediction_class == 1:
            st.error(f"⚠️ **Result: High Risk of Depression Detected**")
            st.write(f"**Confidence Score:** {prediction_prob*100:.2f}%")
        else:
            st.success(f"✅ **Result: Low Risk of Depression**")
            st.write(f"**Confidence Score:** {(1-prediction_prob)*100:.2f}% (Safe)")
    except Exception as e:
        st.error(f"An error occurred during prediction: {e}")