import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from gtts import gTTS
from pytube import Search
import random
import os

# Load pretrained models
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
# Load GPT-2 model and tokenizer for story generation
gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
gpt2_model = AutoModelForCausalLM.from_pretrained("gpt2-medium")

emotion_classifier = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", return_all_scores=True)

def generate_story(theme):
    # A detailed prompt for generating a comforting story about the selected theme
    story_prompt = f"Write a comforting, detailed, and heartwarming story about {theme}. The story should include a character who faces a tough challenge, finds hope, and ultimately overcomes the situation with a positive resolution."
    
    # Generate story using GPT-2 with adjusted parameters
    input_ids = gpt2_tokenizer.encode(story_prompt, return_tensors='pt')
    
    story_ids = gpt2_model.generate(
        input_ids,
        max_length=450,   # Generate slightly shorter but focused stories
        temperature=0.7,  # Balanced creativity without too much randomness
        top_p=0.9,        # Encourage diversity in output
        top_k=50,         # Limit to more probable words
        repetition_penalty=1.2,  # Prevent repetitive patterns
        num_return_sequences=1
    )
    
    # Decode the generated text
    story = gpt2_tokenizer.decode(story_ids[0], skip_special_tokens=True)
    
    # Clean up the generated story by removing the initial prompt
    cleaned_response = story.replace(story_prompt, "").strip()
    
    return cleaned_response


def generate_response(user_input):
    # Limit user input length to prevent overflow issues
    truncated_input = user_input[:200]
    
    # Construct a simpler prompt for generating empathetic responses
    prompt = f"The user is feeling: '{truncated_input}'. Respond with empathy, compassion, and encouragement."
    
    # Encode the prompt
    input_ids = gpt2_tokenizer.encode(prompt, return_tensors='pt')
    
    # Generate the response
    response_ids = gpt2_model.generate(
        input_ids,
        max_length=120,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        repetition_penalty=1.2,
        num_return_sequences=1
    )
    
    # Decode and clean up the generated response
    response = gpt2_tokenizer.decode(response_ids[0], skip_special_tokens=True)
    cleaned_response = response.replace(prompt, "").strip()
    
    return cleaned_response




# Analyze user input for emotional tone
def get_emotion(user_input):
    emotions = emotion_classifier(user_input)
    emotions_sorted = sorted(emotions[0], key=lambda x: x['score'], reverse=True)
    return emotions_sorted[0]['label']

# Function to fetch YouTube videos
def fetch_youtube_videos(activity):
    search = Search(f"{activity} for mental health relaxation")
    search_results = search.results[:3]
    videos = []
    for video in search_results:
        video_url = f"https://www.youtube.com/watch?v={video.video_id}"
        videos.append((video.title, video_url))
    return videos

# Streamlit page configuration
st.set_page_config(page_title="Grief and Loss Support Bot 🌿", page_icon="🌿", layout="centered")
st.markdown("<style>.css-1d391kg { background-color: #F3F7F6; }</style>", unsafe_allow_html=True)

st.title("Grief and Loss Support Bot 🌿")
st.subheader("Your compassionate companion in tough times 💚")

# Sidebar for Meditation and Story Generation
with st.sidebar:
    st.header("🧘 Guided Meditation")
    if st.button("Play Meditation"):
        meditation_audio = "meditation.mp3"
        if not os.path.exists(meditation_audio):
            tts = gTTS("Take a deep breath. Relax and let go of any tension...", lang='en')
            tts.save(meditation_audio)
        st.audio(meditation_audio, format="audio/mp3")

# Generating a comforting story
st.sidebar.header("📖 Short Comforting Story")
story_theme = st.selectbox("Choose a theme for your story:", ["courage", "healing", "hope"])
if st.sidebar.button("Generate Story"):
    with st.spinner("Generating your story..."):
        story = generate_story(story_theme)
    st.text_area("Here's your story:", story, height=300)



# User input section
user_input = st.text_input("Share what's on your mind. I am here to listen...", placeholder="Type here...", max_chars=500, key="user_input_1")


# Initialize session state
if 'previous_responses' not in st.session_state:
    st.session_state.previous_responses = []
if 'badges' not in st.session_state:
    st.session_state.badges = []


# Initialize session state
if 'badges' not in st.session_state:
    st.session_state.badges = []

if user_input:
    with st.spinner("Analyzing your input..."):
        # Get the emotion of the user input
        emotion = get_emotion(user_input)
        
        # Generate an empathetic response
        response = generate_response(user_input)
        
        # Display the bot's response
        st.text_area("Bot's Response:", response, height=250)
    
    # Assign badges based on the detected emotion
    if emotion in ["joy", "optimism"]:
        badge = "🌟 Positivity Badge"
        if badge not in st.session_state.badges:
            st.session_state.badges.append(badge)
            st.success(f"Congratulations! You've earned a {badge}!")

    # Suggest activities based on emotion
    st.info("🎨 Try a New Activity")
    activities = ["exercise", "yoga", "journaling", "painting", "meditation", "swimming"]
    selected_activity = st.selectbox("Pick an activity:", activities)

    if st.button("Find Videos"):
        videos = fetch_youtube_videos(selected_activity)
        if videos:
            for title, url in videos:
                st.write(f"[{title}]({url})")
        else:
            st.write(f"No results found for '{selected_activity}'.")

# Crisis resources
if user_input and any(word in user_input.lower() for word in ["suicide", "help", "depressed"]):
    st.warning("Please reach out to a crisis hotline for immediate support.")
    st.write("[Find emergency resources here](https://www.helpguide.org/find-help.htm)")


# Generate audio response
if user_input:
    tts = gTTS(response, lang='en')
    audio_file = "response.mp3"
    tts.save(audio_file)
    st.audio(audio_file, format="audio/mp3")

# Display badgesz
if st.session_state.badges:
    st.sidebar.header("🏅 Achievements")
    for badge in st.session_state.badges:
        st.sidebar.write(badge)