import streamlit as st
import requests

# Set the title and introduction
st.title('Psychiatric Patient Diagnosis and Treatment App')
st.write("""
Welcome to the Psychiatric Patient Diagnosis and Treatment App. 
Please provide the following information to receive a personalized diagnosis and treatment recommendation.
Your responses will remain confidential.
""")

# Sidebar for inputs
st.sidebar.header('Patient Information')

# Input for Mental Health History
mental_health_history = st.sidebar.text_area('Mental Health History', 
                                             placeholder="Please describe your past psychiatric history...")

# Input for Current Symptoms
current_symptoms = st.sidebar.multiselect(
    'Current Symptoms', 
    ['Anxiety', 'Depression', 'Insomnia', 'Mood Swings', 'Irritability', 'Fatigue', 
     'Difficulty Concentrating', 'Panic Attacks', 'Other']
)

# Input for Medication Preferences
medication_pref = st.sidebar.selectbox(
    'Medication Preferences', 
    ['No Preference', 'Prefer Natural Remedies', 'Prefer Prescription Medication', 'Prefer to Avoid Medication']
)

# Input for Treatment Preferences
treatment_pref = st.sidebar.radio(
    'Treatment Preferences', 
    ['Cognitive Behavioral Therapy (CBT)', 'Medication Management', 'Therapy & Medication', 
     'Mindfulness-Based Therapy', 'Other']
)

# Input for Preferred Communication Method
communication_pref = st.sidebar.selectbox(
    'Preferred Communication Method', 
    ['Text Only', 'Text and Call']
)

# Show additional fields for Text and Call option
patient_name = None
patient_contact_number = None

if communication_pref == 'Text and Call':
    patient_name = st.sidebar.text_input('Patient Name')
    patient_contact_number = st.sidebar.text_input('Contact Number', help="Please enter your phone number")

# Main section for displaying results
st.header('Personalized Diagnosis and Treatment Plan')

# Button to submit the form
if st.sidebar.button('Submit'):
    if communication_pref == 'Text and Call' and (not patient_name or not patient_contact_number):
        st.error("Please provide both name and contact number to receive the call.")
    else:
        with st.spinner('Generating your personalized diagnosis...'):
            try:
                # Prepare the prompt
                prompt = (
                    f"Mental Health History: {mental_health_history}\n"
                    f"Current Symptoms: {', '.join(current_symptoms)}\n"
                    f"Medication Preferences: {medication_pref}\n"
                    f"Treatment Preferences: {treatment_pref}\n\n"
                    "Please provide a detailed psychiatric diagnosis and comprehensive treatment plan "
                    "including therapy recommendations and medication options (if applicable). "
                    "Format your response with clear sections for Diagnosis and Treatment Plan using markdown headers."
                )

                # Call AIML API
                response = requests.post(
                    "https://api.aimlapi.com/v1/chat/completions",
                    headers={
                        "Content-Type": "application/json",
                        "Authorization": f"Bearer {st.secrets['aimlapi_key']}"
                    },
                    json={
                        "model": "anthropic/claude-3.7-sonnet",
                        "messages": [
                            {
                                "role": "user",
                                "content": "You are a world-class psychiatrist who specializes in providing "
                                           "personalized mental health diagnoses and treatment plans. " + prompt
                            }
                        ],
                        "max_tokens": 1500,  # Increased slightly to allow for comprehensive responses
                        "temperature": 0.7
                    },
                    timeout=30  # Added timeout to prevent hanging
                )

                # Handle the response
                if response.status_code in (200, 201):  # Accept both success status codes
                    data = response.json()
                    if 'choices' in data and len(data['choices']) > 0:
                        diagnosis_and_treatment = data['choices'][0]['message']['content']
                        
                        # Display the formatted response
                        st.markdown(diagnosis_and_treatment)
                        
                        # Inform the patient that a call will be placed soon
                        if communication_pref == 'Text and Call':
                            st.success(f"Thank you, {patient_name}. A call will be placed to you at "
                                      f"{patient_contact_number} within the next 24 hours.")
                    else:
                        st.error("No response received from the AI model.")
                else:
                    error_msg = f"API request failed with status code {response.status_code}"
                    try:
                        error_details = response.json()
                        error_msg += f"\nDetails: {error_details.get('message', 'No additional details')}"
                    except:
                        error_msg += f"\nResponse: {response.text}"
                    st.error(error_msg)

            except requests.exceptions.Timeout:
                st.error("The request timed out. Please try again.")
            except Exception as e:
                st.error(f"An unexpected error occurred: {str(e)}")
else:
    st.info("Please complete the form in the sidebar and click 'Submit' to receive your personalized plan.")

# Additional UI/UX features for a calming experience
st.markdown("""
<style>
body {
    background-color: #f4f7f8;
    font-family: 'Arial', sans-serif;
}
.stSidebar {
    background-color: #f0f2f6;
}
.stButton > button {
    background-color: #4CAF50;
    color: white;
    border-radius: 5px;
    padding: 0.5rem 1rem;
}
.stMarkdown h1, .stMarkdown h2, .stMarkdown h3 {
    color: #2e86ab;
}
.stTextArea textarea, .stTextInput input {
    border-radius: 5px;
}
</style>
""", unsafe_allow_html=True)

# Footer with copyright and year
st.markdown("---")
st.markdown("""
<div style="text-align: center; padding: 1rem; color: #666;">
    © 2024 PsychGuide by Engr. Sumayyea Salahuddin. All rights reserved.<br>
    <small>This app provides general information only and is not a substitute for professional medical advice.</small>
</div>
""", unsafe_allow_html=True)