import gradio as gr
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from gtts import gTTS
import pdfplumber
import os

# Load models
def load_model():
    try:
        # Sentiment Analysis Pipeline
        sentiment_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
        sentiment_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
        sentiment_analyzer = pipeline("sentiment-analysis", model=sentiment_model, tokenizer=sentiment_tokenizer)
        
        # Emotion Detection
        emotion_model = AutoModelForSequenceClassification.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
        emotion_tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
        emotion_analyzer = pipeline("sentiment-analysis", model=emotion_model, tokenizer=emotion_tokenizer)
        
        # Sentence Transformer for PDF Insights
        book_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Text Generator
        generator = pipeline("text-generation", model="gpt2")
        
        return sentiment_analyzer, emotion_analyzer, book_model, generator
    except Exception as e:
        print(f"Error loading models: {str(e)}")
        return None, None, None, None

# Text-to-Speech Function
def text_to_speech(text):
    try:
        tts = gTTS(text, lang='en')
        audio_file = "response.mp3"
        tts.save(audio_file)
        return audio_file
    except Exception as e:
        print(f"Error in text-to-speech conversion: {str(e)}")
        return None

# Load and extract text from PDFs in repository
def extract_text_from_pdf(pdf_path):
    text = ""
    try:
        with pdfplumber.open(pdf_path) as pdf:
            for page in pdf.pages:
                text += page.extract_text() or ""  # Handle empty pages gracefully
    except Exception as e:
        return f"Error extracting text: {e}"
    return text

# Load book text from repository
pdf_files = ['Diagnostic_and_statistical_manual_of_mental_disorders_DSM-5.pdf', 'Theories_of_Personality.pdf']
all_text = ""
for pdf_file in pdf_files:
    pdf_path = os.path.join("repository", pdf_file)  # Adjust path to repository
    all_text += extract_text_from_pdf(pdf_path) + "\n"

# Chunk the combined text into smaller segments
def chunk_text(text, chunk_size=300):
    sentences = text.split('. ')
    chunks = []
    current_chunk = ""
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= chunk_size:
            current_chunk += sentence + ". "
        else:
            chunks.append(current_chunk.strip())
            current_chunk = sentence + ". "
    if current_chunk:
        chunks.append(current_chunk.strip())
    return chunks

chunks = chunk_text(all_text)

# Embed and index text chunks
book_model = SentenceTransformer('all-MiniLM-L6-v2')
embeddings = book_model.encode(chunks, convert_to_tensor=True).detach().cpu().numpy()
embedding_dim = embeddings.shape[1]
index = faiss.IndexFlatL2(embedding_dim)
index.add(embeddings)

# Main Assistant Function
def virtual_psychologist_assistant(user_input, sentiment_analyzer, emotion_analyzer, generator):
    try:
        # Sentiment and Emotion Analysis
        sentiment = sentiment_analyzer(user_input)[0]['label']
        emotion = emotion_analyzer(user_input)[0]['label']
        
        # Embed the query and search FAISS for similar chunks
        query_embedding = book_model.encode([user_input], convert_to_tensor=True).detach().cpu().numpy()
        k = 3  # Number of similar chunks to retrieve
        _, retrieved_indices = index.search(query_embedding, k)
        retrieved_chunks = [chunks[idx] for idx in retrieved_indices[0] if idx < len(chunks)]
        
        # Combine retrieved chunks and generate a response
        context = " ".join(retrieved_chunks) if retrieved_chunks else "No relevant book information found."
        prompt = f"User query: {user_input}\n\nRelevant information from books: {context}\n\nSupportive response:"
        response = generator(prompt, max_new_tokens=150, num_return_sequences=1)[0]["generated_text"]
        
        # Final response with emotion and sentiment analysis
        final_response = f"Emotion Detected: {emotion}\nSentiment: {sentiment}\n\nGenerated Response:\n{response}"
        
        # Convert response to audio
        audio_file = text_to_speech(final_response)
        
        return final_response, audio_file
    except Exception as e:
        print(f"Error during response generation: {str(e)}")
        return f"An error occurred: {str(e)}", None

# Gradio Interface
def gradio_interface(user_input):
    sentiment_analyzer, emotion_analyzer, book_model, generator = load_model()
    
    if None in (sentiment_analyzer, emotion_analyzer, book_model, generator):
        return "Error loading models. Please try again later.", None
    
    response, audio_file = virtual_psychologist_assistant(user_input, sentiment_analyzer, emotion_analyzer, generator)
    return response, audio_file

# Create Gradio interface
iface = gr.Interface(fn=gradio_interface,
                     inputs="text",
                     outputs=["text", "audio"],
                     title="Enhanced Virtual Psychologist",
                     description="A virtual psychologist providing responses with emotional and sentiment insights, supported by book-based information.",
                     live=False)

iface.launch()