import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.metrics.pairwise import cosine_similarity

# Load swissBERT for sentence embeddings model
model_name = "jgrosjean-mathesis/sentence-swissbert"
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_sentence_embedding(sentence, language):
    if "de" in language:
        model.set_default_language("de_CH")
    if "fr" in language:
        model.set_default_language("fr_CH")
    if "it" in language:
        model.set_default_language("it_CH")
    if "rm" in language:
        model.set_default_language("rm_CH")
    inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt", max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    token_embeddings = outputs.last_hidden_state
    attention_mask = inputs['attention_mask'].unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * attention_mask, 1)
    sum_mask = torch.clamp(attention_mask.sum(1), min=1e-9)
    embedding = sum_embeddings / sum_mask
    return embedding

def calculate_cosine_similarities(source_sentence, source_language, target_sentence_1, target_language_1, target_sentence_2, target_language_2, target_sentence_3, target_language_3):
    source_embedding = generate_sentence_embedding(source_sentence, source_language)
    target_embedding_1 = generate_sentence_embedding(target_sentence_1, target_language_1)
    target_embedding_2 = generate_sentence_embedding(target_sentence_2, target_language_2)
    target_embedding_3 = generate_sentence_embedding(target_sentence_3, target_language_3)
    cosine_score_1 = cosine_similarity(source_embedding, target_embedding_1)
    cosine_score_2 = cosine_similarity(source_embedding, target_embedding_2)
    cosine_score_3 = cosine_similarity(source_embedding, target_embedding_3)
    cosine_scores = {
        target_sentence_1: cosine_score_1[0][0],
        target_sentence_2: cosine_score_2[0][0],
        target_sentence_3: cosine_score_3[0][0]
    }
    cosine_scores_dict = dict(sorted(cosine_scores.items(), key=lambda item: item[1], reverse=True))
    cosine_scores_output = ""
    for key, value in cosine_scores_dict.items():
        cosine_scores_output += key + ": " + str(value) + "\n"
    cosine_scores_output = "**" + cosine_scores_output.replace("\n", "**\n", 1)
    return cosine_scores_output

def main():   
    demo = gr.Interface(
        fn=calculate_cosine_similarities,
        inputs=[
            gr.Textbox(lines=1, placeholder="Enter source sentence", label="Source Sentence"),
            gr.Dropdown(["de", "fr", "it", "rm"], label="Source Language"),
            gr.Textbox(lines=1, placeholder="Enter target sentence 1", label="Target Sentence 1"),
            gr.Dropdown(["de", "fr", "it", "rm"], label="Target Language 1"),
            gr.Textbox(lines=1, placeholder="Enter target sentence 2", label="Target Sentence 2"),
            gr.Dropdown(["de", "fr", "it", "rm"], label="Target Language 2"),
            gr.Textbox(lines=1, placeholder="Enter target sentence 3", label="Target Sentence 3"),
            gr.Dropdown(["de", "fr", "it", "rm"], label="Target Language 3")
            ],
        outputs= gr.Textbox(label="Cosine Similarity Scores", type="text", lines=3),
        title="Sentence Similarity Calculator",
        description="Enter a source sentence and up to three target sentences to calculate their cosine similarity.",
        examples=[
            ["Der Zug fährt um 9 Uhr in Zürich ab.", "de", "Le train arrive à Lausanne à 11 heures.", "fr", "Alla stazione di Lugano ci sono diversi binari.", "it", "A Cuera van biars trens ellas muntognas.", "rm"]
        ]
    )
    demo.launch(share=True)

if __name__ == "__main__":
    main()