import numpy as np
from sentence_transformers import SentenceTransformer
from gtts import gTTS

import gradio as gr
from transformers import pipeline

# Load Sentiment Analysis and Emotion Detection models
sentiment_analyzer = pipeline("sentiment-analysis")
emotion_detector = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

# Load a smaller pre-trained text generation model for faster output
text_generator = pipeline("text-generation", model="distilgpt2")  # Smaller model for faster inference

# Analyze sentiment and emotion
def analyze_sentiment_and_emotion(user_input):
    sentiment = sentiment_analyzer(user_input)[0]['label']
    emotion = emotion_detector(user_input)[0]['label']
    return sentiment, emotion

# Generate response based on sentiment and emotion with a supportive approach
def generate_response(user_input, sentiment, emotion):
    # Construct response prompts based on detected sentiment and emotion
    if sentiment == "POSITIVE":
        prompt = (
            f"I’m glad to hear that you’re feeling {emotion}. Maintaining positive energy is so important, "
            f"and it's wonderful to hear about your good state. Here are a few ways you can continue to nurture "
            f"this happiness: focus on the things that bring you joy, spend time with people who uplift you, "
            f"and take moments to appreciate your achievements, big and small. Remember, happiness is often found "
            f"in appreciating the little things in life. Keep this positivity going, and know that your good "
            f"energy can also inspire others around you. Is there anything specific that brings you joy?"
        )
    else:  # Assume sentiment is NEGATIVE
        prompt = (
            f"I’m here to listen and support you. It’s completely normal to feel {emotion} sometimes, and it’s "
            f"okay to acknowledge these feelings. Often, difficult emotions are a way for our minds to tell us "
            f"that something needs attention. Take your time, and consider ways to care for yourself. Whether it’s "
            f"reaching out to loved ones, taking a break, or reflecting on what brings you peace, there are steps "
            f"you can take. Remember, emotions are part of life’s journey, and finding meaning through challenging "
            f"times can lead to growth. You’re not alone in this; I’m here to help guide you. "
        )

    # Generate the response from the model
    response = text_generator(prompt, max_length=250, num_return_sequences=1, do_sample=True)[0]["generated_text"]
    
    return response

# Convert text response to audio
def text_to_audio(response_text):
    tts = gTTS(response_text, lang='en')
    tts.save("response.mp3")
    return "response.mp3"  # Return the file path for Gradio to play audio

# Process the input and generate output
def gradio_interface(user_input):
    if not user_input:
        return "Please provide complete input and submit.", None, None, None
    
    # Perform sentiment analysis and emotion detection
    sentiment, emotion = analyze_sentiment_and_emotion(user_input)
    
    # Generate a response from the text generation model
    response_text = generate_response(user_input, sentiment, emotion)
    
    # Convert to audio and return both text and audio output
    audio_file = text_to_audio(response_text)
    
    return response_text, audio_file, sentiment, emotion

# Create the Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs="text",
    outputs=["text", "audio", "text", "text"],  # Outputs: response, audio, sentiment, emotion
    live=False,  # Disable live updates to require submitting
    title="Virtual Psychologist App",
    description="Enter your thoughts or feelings, and the app will provide a thoughtful response based on your input. It also provides sentiment and emotion analysis.",
    allow_flagging="never"  # Optional: Disable flagging
)

# Launch the app
iface.launch()