import streamlit as st
import pandas as pd
import numpy as np
import altair as alt
import datetime

def main():
    st.set_page_config(
        page_title="AfyaMind Mental Health Platform",
        page_icon="🧠",
        layout="wide"
    )

    # Sample Data Generation (Updated for Mental Health)
    np.random.seed(42)
    date_rng = pd.date_range(start='2025-01-01', end='2025-12-31', freq='D')
    data = pd.DataFrame(date_rng, columns=['date'])
    data['user_id'] = np.random.randint(1000, 9999, len(data))
    data['mood_score'] = np.random.randint(1, 11, len(data))  # 1-10 scale
    data['phq9_score'] = np.random.randint(0, 28, len(data))  # PHQ-9: 0-27
    data['gad7_score'] = np.random.randint(0, 22, len(data))  # GAD-7: 0-21
    data['crisis_flag'] = np.random.choice([True, False], p=[0.05, 0.95], size=len(data))
    data['appointment_booked'] = np.random.choice([True, False], size=len(data))
    data['engagement_points'] = np.random.randint(0, 100, len(data))  # Gamification
    data['therapist_id'] = np.random.randint(5000, 5999, len(data))
    data['message_sample'] = np.random.choice(
        ["Feeling okay", "Hopeless today", "Doing great", "Need help"],
        size=len(data)
    )

    # Sidebar Filters
    st.sidebar.header("Filters")
    selected_user = st.sidebar.selectbox("User ID", ['All'] + sorted(data['user_id'].unique()))
    selected_date = st.sidebar.slider(
        "Date Range",
        min_value=datetime.date(2025, 1, 1),
        max_value=datetime.date(2025, 12, 31),
        value=(datetime.date(2025, 1, 1), datetime.date(2025, 12, 31))
    )

    # Apply Filters
    filtered_data = data.copy()
    if selected_user != 'All':
        filtered_data = filtered_data[filtered_data['user_id'] == selected_user]
    filtered_data = filtered_data[
        (filtered_data['date'].dt.date >= selected_date[0]) &
        (filtered_data['date'].dt.date <= selected_date[1])
    ]

    # Main Content Tabs (Updated for AfyaMind Features)
    tab1, tab2, tab3, tab4, tab5 = st.tabs([
        "User Mood Analytics",
        "Crisis Detection",
        "Therapy Appointments",
        "Clinician Dashboard",
        "Engagement & Analytics"
    ])

    # Tab 1: User Mood Analytics
    with tab1:
        st.subheader("User Mood Trends")
        col1, col2, col3 = st.columns(3)
        col1.metric("Avg Mood Score", f"{filtered_data['mood_score'].mean().round(1)}/10")
        col2.metric("Avg PHQ-9 Score", f"{filtered_data['phq9_score'].mean().round(1)}/27")
        col3.metric("Avg GAD-7 Score", f"{filtered_data['gad7_score'].mean().round(1)}/21")

        mood_chart = alt.Chart(filtered_data).mark_line().encode(
            x='date:T',
            y='mood_score:Q',
            tooltip=['date:T', 'mood_score:Q', 'phq9_score:Q']
        ).properties(title="Mood Score Over Time")
        st.altair_chart(mood_chart, use_container_width=True)

    # Tab 2: Crisis Detection
    with tab2:
        st.subheader("Crisis Detection & Response")
        crisis_data = filtered_data[filtered_data['crisis_flag'] == True]
        st.write(f"Users flagged for crisis: {len(crisis_data)}")
        st.dataframe(crisis_data[['date', 'user_id', 'message_sample', 'phq9_score']])

        # Simulate NLP-based distress signal detection
        distress_messages = filtered_data[filtered_data['message_sample'].str.contains("hopeless|need help", case=False)]
        if not distress_messages.empty:
            st.warning("Distress Signals Detected - Auto-escalated to Therapists")
            st.dataframe(distress_messages[['date', 'user_id', 'message_sample']])

    # Tab 3: Therapy Appointments
    with tab3:
        st.subheader("Appointment Scheduling")
        appointments = filtered_data[filtered_data['appointment_booked'] == True]
        st.dataframe(appointments[['date', 'user_id', 'therapist_id']])
        st.write(f"Total Appointments Booked: {len(appointments)}")
        if st.button("Simulate SMS Reminder"):
            st.success("SMS Reminder Sent: 'Your appointment is tomorrow at 10 AM.'")

    # Tab 4: Clinician Dashboard
    with tab4:
        st.subheader("Clinician Support Dashboard")
        therapist_load = filtered_data.groupby('therapist_id').agg({
            'user_id': 'nunique',
            'appointment_booked': 'sum'
        }).reset_index()
        therapist_load.columns = ['Therapist ID', 'Unique Clients', 'Appointments']
        st.dataframe(therapist_load)

        # Burnout Alert
        overloaded = therapist_load[therapist_load['Appointments'] > 20]
        if not overloaded.empty:
            st.error(f"Therapist Overload Detected: {len(overloaded)} therapists exceed 20 sessions.")

        # AI-Generated User Summary
        if selected_user != 'All':
            user_summary = filtered_data[filtered_data['user_id'] == selected_user].iloc[-1]
            st.write("AI-Generated Session Prep Summary:")
            st.write(f"User {selected_user} - Mood: {user_summary['mood_score']}/10, PHQ-9: {user_summary['phq9_score']}, Last Message: '{user_summary['message_sample']}'")

    # Tab 5: Engagement & Analytics
    with tab5:
        st.subheader("Engagement & Population Insights")
        col1, col2 = st.columns(2)
        col1.metric("Avg Engagement Points", f"{filtered_data['engagement_points'].mean().round(1)}")
        col2.metric("Mood Tracking Adoption", f"{(len(filtered_data['mood_score'].notna()) / len(filtered_data)) * 100:.1f}%")

        # Gamified Engagement Chart
        points_chart = alt.Chart(filtered_data).mark_bar().encode(
            x=alt.X('engagement_points:Q', bin=True),
            y='count()',
            tooltip=['engagement_points:Q', 'count()']
        ).properties(title="Engagement Points Distribution")
        st.altair_chart(points_chart, use_container_width=True)

        # Download Anonymized Data
        st.download_button(
            label="Download Anonymized Data (CSV)",
            data=filtered_data.drop(columns=['user_id', 'therapist_id']).to_csv(index=False),
            file_name="afyamind_analytics.csv",
            mime="text/csv"
        )

if __name__ == "__main__":
    main()