import gradio as gr
import random
import transformers
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util

# Load AI models
similarity_model = SentenceTransformer('all-MiniLM-L6-v2')
generator = pipeline("text-generation", model="gpt2")

# Freud AI knowledge base
FREUD_KNOWLEDGE = [
    "Sigmund Freud (1856–1939) was an Austrian neurologist and the founder of psychoanalysis.",
    "Freud developed the structural model of the mind, which consists of the id, ego, and superego.",
    "Freud's theory of psychosexual development includes stages such as oral, anal, phallic, latent, and genital.",
    "The Oedipus complex is a key Freudian concept describing a child's unconscious desire for the opposite-sex parent.",
    "Psychoanalysis is a therapeutic method developed by Freud that explores unconscious motivations and conflicts.",
    "Freud's defense mechanisms include repression, denial, projection, displacement, and sublimation.",
    "Freud's 'Interpretation of Dreams' (1899) proposed that dreams reveal unconscious desires and repressed memories.",
    "Freud's theory suggests that childhood experiences shape adult personality and behavior.",
    "Freud introduced the free association technique to help patients uncover unconscious thoughts.",
    "The concept of libido, in Freudian theory, refers to the energy derived from sexual instincts that drives human behavior."
]

# List of random greetings and farewells
RANDOM_GREETINGS = [
    "Nice to meet you, {name}! Let's explore Freud together!",
    "Glad to have you here, {name}! What would you like to know?",
    "Welcome, {name}! Freud had some fascinating theories. What can I help with?",
    "Great to meet you, {name}! I'm ready to discuss Freud whenever you are!",
    "Hello, {name}! Freud's work is intriguing. How can I assist you today?"
]

RANDOM_FAREWELLS = [
    "Goodbye, {name}! Have a wonderful rest of your day!",
    "Take care, {name}! Hope to chat with you again soon!",
    "See you later, {name}! Stay curious about Freud!",
    "Farewell, {name}! Keep exploring psychology!",
    "Goodbye, {name}! Wishing you a fantastic day ahead!"
]

# AI processing functions
def search_relevant_answer(question):
    """Find the most relevant Freud-related response."""
    embeddings = similarity_model.encode(FREUD_KNOWLEDGE, convert_to_tensor=True)
    query_embedding = similarity_model.encode(question, convert_to_tensor=True)
    similarities = util.pytorch_cos_sim(query_embedding, embeddings)[0]
    best_match_idx = similarities.argmax().item()
    return FREUD_KNOWLEDGE[best_match_idx]

def chatbot(name, user_input):
    """Handles user interactions with the chatbot."""
    if "bye" in user_input.lower() or "exit" in user_input.lower():
        return random.choice(RANDOM_FAREWELLS).format(name=name)

    if "thank" in user_input.lower():
        return f"You're welcome, {name}! Is there anything else I can help you with?"

    if any(word in user_input.lower() for word in ["freud", "psychoanalysis", "dreams", "childhood", "ego", "id", "superego"]):
        relevant_answer = search_relevant_answer(user_input)
        ai_response = generator(relevant_answer, max_length=100, num_return_sequences=1)[0]['generated_text']
        return ai_response

    return f"I only answer questions about Sigmund Freud, {name}. Please ask me something related to his life and work."

# Gradio interface
iface = gr.Interface(
    fn=chatbot,
    inputs=["text", "text"],
    outputs="text",
    title="Freud AI Chatbot",
    description="Ask me anything about Sigmund Freud and his work!",
    examples=[["John", "Who was Freud?"], ["Alice", "Tell me about the ego"], ["Sam", "What is psychoanalysis?"]]
)

# Launch chatbot
iface.launch()