import streamlit as st
from transformers import AutoTokenizer
from llama_cpp import Llama
#from transformers import pipeline
#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)

# Path to your GGUF model file
MODEL_PATH = "/path/to/QuantFactory/Mental-Health-FineTuned-Mistral-7B-Instruct-v0.2-GGUF"

# Load Llama model
@st.cache_resource
def load_llama_model():
    try:
        return Llama(model_path=MODEL_PATH, n_threads=8)  # Adjust `n_threads` based on your system
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

llama_model = load_llama_model()

# Load dataset for context
@st.cache_resource
def load_counseling_dataset():
    return load_dataset("Amod/mental_health_counseling_conversations")

dataset = load_counseling_dataset()

# Streamlit app
st.title("Mental Health Counseling Chat")
st.markdown("""
Welcome to the Mental Health Counseling Chat application. 
This platform is designed to provide supportive, positive, and encouraging responses based on mental health counseling expertise.
""")

# Check if the model loaded correctly
if llama_model is None:
    st.error("The text generation model could not be loaded. Please check the model path and configuration.")
else:
    # Explore dataset for additional context or resources (optional)
    if st.checkbox("Show Example Questions and Answers from Dataset"):
        sample = dataset["train"].shuffle(seed=42).select(range(3))  # Display 3 random samples
        for example in sample:
            st.markdown(f"**Question:** {example.get('context', 'N/A')}")
            st.markdown(f"**Answer:** {example.get('response', 'N/A')}")
            st.markdown("---")

    # User input for mental health concerns
    user_input = st.text_area("Your question or concern:", placeholder="Type here...")

    if st.button("Get Supportive Response"):
        if user_input.strip():
            try:
                # Generate response using Llama
                prompt = f"User: {user_input}\nCounselor:"
                response = llama_model(prompt, max_tokens=200, stop=["\n", "User:"])
                
                st.subheader("Counselor's Response:")
                st.write(response["choices"][0]["text"].strip())
            except Exception as e:
                st.error(f"An error occurred while generating the response: {e}")
        else:
            st.error("Please enter a question or concern to receive a response.")

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

    st.sidebar.info("This application is not a replacement for professional counseling. If you are in crisis, please seek professional help immediately.")