import streamlit as st
from transformers import AutoTokenizer
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM
from datasets import load_dataset

# Replace with the direct image URL
flower_image_url = "https://i.postimg.cc/hG2FG85D/2.png"

# Inject custom CSS for the background with a centered and blurred image
st.markdown(
    f"""
    <style>
    /* Container for background */
    html, body {{
        margin: 0;
        padding: 0;
        overflow: hidden;
    }}
    [data-testid="stAppViewContainer"] {{
        position: relative;
        z-index: 1; /* Ensure UI elements are above the background */
    }}
    /* Blurred background image */
    .blurred-background {{
        position: fixed;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
        z-index: -1; /* Send background image behind all UI elements */
        background-image: url("{flower_image_url}");
        background-size: cover;
        background-position: center;
        filter: blur(10px); /* Adjust blur ratio here */
        opacity: 0.8; /* Optional: Add slight transparency for a subtle effect */
    }}
    </style>
    """,
    unsafe_allow_html=True
)

# Add the blurred background div
st.markdown('<div class="blurred-background"></div>', unsafe_allow_html=True)

# Load the fine-tuned model and tokenizer
@st.cache_resource
def load_model_and_tokenizer():
    config = PeftConfig.from_pretrained("zementalist/llama-3-8B-chat-psychotherapist")
    base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
    model = PeftModel.from_pretrained(base_model, "zementalist/llama-3-8B-chat-psychotherapist")
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
    return model, tokenizer

model, tokenizer = load_model_and_tokenizer()

# Load dataset for reference (optional)
@st.cache_resource
def load_dataset_reference():
    return load_dataset("Amod/mental_health_counseling_conversations")

dataset = load_dataset_reference()

# Streamlit App Configuration
st.title("Mental Well-Being Support Application")
st.markdown("""
Welcome to the Mental Well-Being Support Application. This platform is designed to provide positive, supportive, and encouraging responses to your mental health concerns. Our responses are powered by a fine-tuned AI model based on expert psychologists' answers.
""")

# User Input Section
st.header("Your Mental Health Journey")
user_query = st.text_area("Please share your thoughts or questions:", placeholder="Write here...")

# Generate AI Response
if st.button("Get Supportive Response"):
    if user_query.strip():
        # Generate response
        inputs = tokenizer(f"User: {user_query}\nAI:", return_tensors="pt")
        outputs = model.generate(inputs.input_ids, max_length=200, temperature=0.7, num_return_sequences=1)
        ai_response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("AI:")[-1].strip()

        # Display the response
        st.subheader("Your Supportive Response:")
        st.write(ai_response)
    else:
        st.error("Please enter a question or concern to get a response.")

# Additional Resources Section
st.sidebar.header("Resources")
st.sidebar.markdown("""
- [Mental Health Foundation](https://www.mentalhealth.org)
- [Mind](https://www.mind.org.uk)
- [National Suicide Prevention Lifeline](https://suicidepreventionlifeline.org)
""")

# Footer
st.sidebar.info("This application is not a replacement for professional help. If you're in crisis, please contact a mental health professional.")