# Step 2: Import Libraries and Set Up Models
import gradio as gr
from transformers import pipeline
from gtts import gTTS
import os
import IPython.display as ipd

# Load sentiment analysis and text generation models
sentiment_model = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
text_gen_model = pipeline("text-generation", model="microsoft/DialoGPT-medium")

# Step 3: Define the Response Function
def virtual_psychologist(input_text):
    # Check for empty or invalid input
    if not input_text.strip():
        return "Please provide some input about how you're feeling.", None, None

    # Step 1: Sentiment Analysis
    sentiment = sentiment_model(input_text)[0]
    label = sentiment['label']
    confidence = sentiment['score']

    # Step 2: Display Sentiment Information
    sentiment_feedback = f"Your input sentiment is detected as **{label}** with confidence {confidence:.2f}.\n\n"

    # Step 3: Generate a Response Based on Sentiment
    if confidence > 0.7:  # Threshold for confident sentiment analysis
        if label == 'POSITIVE':
            response = "I'm glad you're feeling positive! Tell me more about what’s bringing you joy, and let’s keep this energy up together."
        elif label == 'NEGATIVE':
            # Check for sensitive keywords like "suicide" or "worthless"
            if "suicide" in input_text.lower() or "worthless" in input_text.lower():
                response = ("I'm really sorry you're feeling this way, but please know you're not alone. "
                            "It's really important to talk to someone who can provide support. Would you like to share more "
                            "about what's been making you feel this way? You matter, and it's okay to reach out for help.")
            else:
                response = "It sounds like you're going through a tough time. Want to share more about what’s on your mind? I'm here to help you navigate through it."
        else:
            response = "You seem to be feeling neutral. Do you have anything specific on your mind that you'd like to talk about?"
    else:
        response = "I'm not quite sure I understand. Could you elaborate a bit more? I'm here to listen."

    # Step 4: Generate a Longer Response for the User
    generated_text = text_gen_model(response, max_length=100, num_return_sequences=1)[0]['generated_text']

    # Combine sentiment feedback and generated text into the final output
    full_response = sentiment_feedback + generated_text

    # Convert response to speech using gTTS
    tts = gTTS(text=full_response, lang='en')
    tts.save("response.mp3")

    # Return the response text and audio file path
    response_type = "Supportive Response" if "suicide" in input_text.lower() or "worthless" in input_text.lower() else "General Response"

    # Return response as text, response type, and audio file path for Gradio
    return full_response, response_type, "response.mp3"

# Step 5: Create a Gradio interface for the app
interface = gr.Interface(
    fn=virtual_psychologist,
    inputs=gr.Textbox(lines=5, placeholder="How are you feeling today?"),
    outputs=[gr.Textbox(), gr.Textbox(), gr.Audio()],
    title="Virtual Psychologist Assistant",
    description="Share your feelings, and this assistant will analyze your sentiment and respond as a supportive psychologist."
)

# Step 6: Launch the Gradio interface
interface.launch()