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

# Load models
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
emotion_analyzer = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

# Set up 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; }
    .css-ffhzg2 { font-size: 1.5em; font-weight: 500; color: #4C6D7D; }
    .stTextInput>div>div>input { background-color: #D8E3E2; }
    .stButton>button { background-color: #A9D0B6; color: white; border-radius: 5px; }
    .stButton>button:hover { background-color: #8FB79A; }
    .stTextInput>div>label { color: #4C6D7D; }
    </style>
""", unsafe_allow_html=True)

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

# Function to generate a response
def generate_response(user_input):
    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    chat_history_ids = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id, temperature=0.7, top_k=50, repetition_penalty=1.2)
    bot_output = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return bot_output

# Function to generate a comforting story based on a theme
def generate_story(theme):
    prompt = f"Tell me a short, comforting story about {theme}."
    input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors='pt')
    
    # Generate story
    story_ids = model.generate(input_ids, max_length=150, temperature=0.7, top_p=0.9, repetition_penalty=1.2, do_sample=True, num_return_sequences=1)
    story = tokenizer.decode(story_ids[0], skip_special_tokens=True)
    return story

# Mindfulness Meditation
st.write("### 🧘 Guided Meditation")
if st.button("Play a 5-minute Guided Meditation"):
    meditation_text = "Focus on your breath. Inhale deeply, hold for a moment, and exhale slowly. Let go of any tension."
    tts_meditation = gTTS(meditation_text, lang='en')
    tts_meditation.save("meditation.mp3")
    st.audio("meditation.mp3", format="audio/mp3")

# Interactive Storytelling
st.write("### 📖 Short Comforting Story")

# User selects a theme for the story
story_theme = st.selectbox("Choose a theme for your story:", ["hope", "courage", "healing", "friendship", "resilience"])
if st.button("Generate Story"):
    story = generate_story(story_theme)
    st.text_area("Here’s your comforting story:", story, height=250)
    
    # Convert story to speech
    tts_story = gTTS(story, lang='en')
    tts_story.save("story.mp3")
    st.audio("story.mp3", format="audio/mp3")

# User input for conversational support
user_input = st.text_input("Share what's on your mind...", placeholder="Type here...", max_chars=500)
if user_input:
    response = generate_response(user_input)
    st.text_area("Bot's Response:", response, height=250)
    
    # Text-to-speech for response
    tts = gTTS(response, lang='en')
    tts.save("response.mp3")
    st.audio("response.mp3", format="audio/mp3")

# Enhanced Activity Suggestions
st.write("### 🎨 Try a New Activity")
activities = ["journaling", "yoga", "painting", "exercise", "meditation"]
activity_choice = st.selectbox("Pick an activity:", activities)
if st.button("Find Videos"):
    try:
        search = Search(activity_choice)
        results = search.results[:3]
        for video in results:
            st.write(f"[{video.title}]({video.watch_url})")
    except Exception as e:
        st.write("Error fetching videos.")