import gradio as gr
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import requests
from PIL import Image
from io import BytesIO

# Load model
model = SentenceTransformer('all-mpnet-base-v2')

# Load emoji dictionaries
def kitchen_txt_to_dict(filepath):
    emoji_dict = {}
    with open(filepath, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split(' ', 1)
            if len(parts) == 2:
                emoji, desc = parts
                emoji_dict[emoji] = desc
    return emoji_dict

file_path_emotion = 'google-emoji-kitchen-emotion.txt'
file_path_item = 'google-emoji-kitchen-item.txt'

emotion_dict = kitchen_txt_to_dict(file_path_emotion)
event_dict = kitchen_txt_to_dict('google-emoji-kitchen-item.txt')

# Precompute embeddings
emotion_embeddings = {emoji: model.encode(desc) for emoji, desc in emotion_dict.items()}
event_embeddings = {emoji: model.encode(desc) for emoji, desc in event_dict.items()}

# Helper functions
def find_top_emojis(embedding, emoji_embeddings, top_n=1):
    similarities = [
        (emoji, cosine_similarity([embedding], [e_embed])[0][0])
        for emoji, e_embed in emoji_embeddings.items()
    ]
    similarities.sort(key=lambda x: x[1], reverse=True)
    return [emoji for emoji, _ in similarities[:top_n]]

def get_emoji_kitchen_url(emoji1, emoji2, size=256):
    return f"https://emojik.vercel.app/s/{emoji1}_{emoji2}?size={size}"

def fetch_image(url):
    try:
        response = requests.get(url)
        if response.status_code == 200 and "image" in response.headers.get("Content-Type", ""):
            return Image.open(BytesIO(response.content))
        else:
            return None
    except:
        return None

# Main function for Gradio
def sentence_to_emojis(sentence):
    embedding = model.encode(sentence)
    
    top_emotion = find_top_emojis(embedding, emotion_embeddings, top_n=1)[0]
    top_event = find_top_emojis(embedding, event_embeddings, top_n=1)[0]
    
    mashup_url = get_emoji_kitchen_url(top_emotion, top_event)
    mashup_image = fetch_image(mashup_url)
    
    return top_emotion, top_event, mashup_image

# Gradio interface
demo = gr.Interface(
    fn=sentence_to_emojis,
    inputs=gr.Textbox(lines=2, placeholder="Type a sentence..."),
    outputs=[
        gr.Text(label="Top Emotion Emoji"),
        gr.Text(label="Top Event Emoji"),
        gr.Image(label=" Kitchen Emoji")
    ],
    title="Sentence → Emoji Mashup",
    description="Get the top emotion and event emoji from your sentence, and view the mashup!"
)

demo.launch(share=True)