from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
import gradio as gr

# Load a conversational model (replace with your Bangla LLM)
chatbot_model_name = "csebuetnlp/banglat5_base"  # Replace with your fine-tuned model
chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
chatbot_model = AutoModelForSeq2SeqLM.from_pretrained(chatbot_model_name)

# Function to generate empathetic responses
def generate_empathetic_response(user_input):
    try:
        # Prepare the input for the chatbot
        input_text = f"User: {user_input}\nChatbot:"
        inputs = chatbot_tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
        
        # Generate response
        outputs = chatbot_model.generate(
            **inputs,
            max_length=512,
            num_beams=5,
            early_stopping=True,
            no_repeat_ngram_size=2
        )
        
        # Decode and clean the response
        response = chatbot_tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Ensure the response is culturally appropriate and empathetic
        response = add_empathy_and_cultural_tone(response)
        return response
    except Exception as e:
        return f"Error generating response: {str(e)}"

# Function to add empathy and cultural tone to the response
def add_empathy_and_cultural_tone(response):
    # Add empathetic phrases
    empathetic_phrases = [
        "বুঝতে পারছি", "কষ্ট পেয়েছো নিশ্চয়", "আমি তোমার পাশে আছি", 
        "এটা খুব কঠিন সময় তোমার জন্য", "তোমার অনুভূতি আমি বুঝতে পারছি"
    ]
    
    # Add regional dialects and informal pauses
    regional_phrases = [
        "ভাইয়া", "হুজুর", "থাইকা", "পাল্লায় পড়া", "টাকা ফুরাইয়া গেছে", 
        "পকেট শূন্য", "মনটা ভাঙা কাঁচের মতো"
    ]
    
    # Randomly add empathetic and regional phrases to the response
    import random
    if random.random() > 0.5:
        response = random.choice(empathetic_phrases) + "। " + response
    if random.random() > 0.5:
        response = response + " " + random.choice(regional_phrases)
    
    return response

# Function to handle the chatbot conversation
def chat_with_bot(user_input, history=[]):
    # Generate a response
    bot_response = generate_empathetic_response(user_input)
    
    # Update conversation history
    history.append((user_input, bot_response))
    
    # Format the conversation history for display
    formatted_history = "\n".join([f"User: {user}\nChatbot: {bot}" for user, bot in history])
    
    return formatted_history, history

# Define the Gradio interface
iface = gr.Interface(
    fn=chat_with_bot,
    inputs=[
        gr.Textbox(label="Your Message", placeholder="আপনার সমস্যা বা অনুভূতি শেয়ার করুন..."),
        gr.State()  # To store conversation history
    ],
    outputs=[
        gr.Textbox(label="Conversation History", lines=10),
        gr.State()  # To pass updated history back
    ],
    live=False,
    title="Bangla Emotional Support Chatbot",
    description="আপনার সমস্যা বা অনুভূতি শেয়ার করুন। আমি আপনার পাশে আছি।",
    theme='Ajaxon6255/Emerald_Isle',
    examples=[
        ["আমার টাকা ফুরিয়ে গেছে এবং আমি ঋণে ডুবে আছি।"],
        ["আমার ভাইয়ের সাথে আমার সম্পর্ক খারাপ হয়ে গেছে।"],
        ["আমি খুব একা বোধ করছি এবং কাউকে পাশে পাচ্ছি না।"]
    ]
)

# Launch the Gradio app
iface.launch(inline=False)