import streamlit as st
import numpy as np
import pandas as pd
import joblib
from sklearn.preprocessing import LabelEncoder 

# Load the saved model and scaler
try:
    model = joblib.load(r'D:\Data_Science\Projects\health\mental_health\version_1\best_depression_model.pkl')
    scaler = joblib.load(r'D:\Data_Science\Projects\health\mental_health\version_1\best_scaler.pkl')
except FileNotFoundError:
    st.error("Model or Scaler file not found. Please check the file paths.")
    st.stop()

# Define categorical options
gender_options = ['Female', 'Male']
status_options = ['Student', 'Working Professional']
sleep_duration_options = ['5-6 hours', '7-8 hours', 'Less than 5 hours', 'More than 8 hours']
diet_options = ['Healthy', 'Moderate', 'Unhealthy']
degree_options = ['Bachelor', 'Master', 'PhD', 'Diploma', 'High School']
suicidal_thoughts_options = ['No', 'Yes']
family_history_options = ['No', 'Yes']

# Sidebar for User Input
st.sidebar.title("User Input")

gender = st.sidebar.selectbox("Gender", gender_options)
age = st.sidebar.number_input("Age", min_value=10, max_value=100, value=25)
status = st.sidebar.selectbox("Are you a Student or Working Professional?", status_options)
work_pressure = st.sidebar.slider("Work Pressure Level (0 = Low, 5 = High)", 0, 5, 2)
job_satisfaction = st.sidebar.slider("Job Satisfaction Level (0 = Low, 5 = High)", 0, 5, 3)
sleep_duration = st.sidebar.selectbox("Sleep Duration", sleep_duration_options)
diet = st.sidebar.selectbox("Dietary Habits", diet_options)
degree = st.sidebar.selectbox("Highest Degree", degree_options)
suicidal_thoughts = st.sidebar.selectbox("Have you ever had suicidal thoughts?", suicidal_thoughts_options)
work_study_hours = st.sidebar.slider("Work/Study Hours per Day", 0, 12, 6)
financial_stress = st.sidebar.slider("Financial Stress Level (0 = None, 5 = High)", 0, 5, 2)
family_history = st.sidebar.selectbox("Family History of Mental Illness", family_history_options)

# Prepare Input Data
input_data = {
    'Gender': gender,
    'Age': age,
    'Working Professional or Student': status,
    'Work Pressure': work_pressure,
    'Job Satisfaction': job_satisfaction,
    'Sleep Duration': sleep_duration,
    'Dietary Habits': diet,
    'Degree': degree,
    'Have you ever had suicidal thoughts ?': suicidal_thoughts,
    'Work/Study Hours': work_study_hours,
    'Financial Stress': financial_stress,
    'Family History of Mental Illness': family_history
}

input_df = pd.DataFrame([input_data])

# Encode categorical inputs
label_encoders = {
    'Gender': LabelEncoder().fit(gender_options),
    'Working Professional or Student': LabelEncoder().fit(status_options),
    'Sleep Duration': LabelEncoder().fit(sleep_duration_options),
    'Dietary Habits': LabelEncoder().fit(diet_options),
    'Degree': LabelEncoder().fit(degree_options),
    'Have you ever had suicidal thoughts ?': LabelEncoder().fit(suicidal_thoughts_options),
    'Family History of Mental Illness': LabelEncoder().fit(family_history_options)
}

for col, le in label_encoders.items():
    input_df[col] = le.transform(input_df[col])

# Standardize numerical inputs using the saved scaler
input_scaled = scaler.transform(input_df)

# Prediction
def predict_depression(input_scaled):
    probabilities = model.predict_proba(input_scaled)[0]
    prone_to_depression = probabilities[1] * 100  # Probability of class 1 (Depression)
    return prone_to_depression

# UI Design
st.title("Depression Proneness Predictor")
st.markdown("""
    <style>
        .main { background-color: #f0f2f6; }
        h1 { color: #4f8bfc; text-align: center; }
        .stButton>button { background-color: #4f8bfc; color: white; }
    </style>
    """, unsafe_allow_html=True)

# Display Image
try:
    st.image("depression.jpg", use_column_width=True)
except FileNotFoundError:
    st.warning("Image not found. Please check the image path.")

st.markdown("<h2 style='text-align: center;'>Predict Your Mental Health Status</h2>", unsafe_allow_html=True)

if st.button("Predict"):
    probability = predict_depression(input_scaled)
    if probability > 50:
        st.error(f"High Probability of Depression: {probability:.2f}%")
    else:
        st.success(f"Low Probability of Depression: {probability:.2f}%")

st.markdown("""
    <div style='text-align: center; color: gray;'>
        This tool is for educational purposes and not a medical diagnosis. 
        Consult a healthcare professional for proper evaluation.
    </div>
    """, unsafe_allow_html=True)
# Footer
st.markdown(
    """
    <div class="footer">
        Developed by <b>SKB</b> | © 2025. All Rights Reserved
    </div>
    """,
    unsafe_allow_html=True)