import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer, util

# Function to load the selected model
def load_model(model_name):
    return SentenceTransformer(model_name)

# Function to compute similarity and classify relationship
def predict(original_sentence_input, *sentences_to_compare):
    model_name = "sartifyllc/African-Cross-Lingua-Embeddings-Model"
    model = load_model(model_name)
    result = {
        "Model Name": model_name,
        "Original Sentence": original_sentence_input,
        "Sentences to Compare": sentences_to_compare,
        "Similarity Scores": {}
    }
    
    if original_sentence_input and sentences_to_compare:
        sentences = [original_sentence_input] + list(sentences_to_compare)
        embeddings = model.encode(sentences)
        original_embedding = embeddings[0]

        for i, sentence in enumerate(sentences_to_compare, start=1):
            similarity_score = util.pytorch_cos_sim(original_embedding, embeddings[i]).item()
            result["Similarity Scores"][f"Sentence {i}"] = similarity_score
    
    return result

# Define inputs and outputs for Gradio interface
model_name_display = gr.Markdown(value="**Model Name**: sartifyllc/African-Cross-Lingua-Embeddings-Model")
original_sentence_input = gr.Textbox(lines=2, placeholder="Enter the original sentence here...", label="Original Sentence")

# Initial sentence comparison inputs
sentence_to_compare_inputs = [
    gr.Textbox(lines=2, placeholder="Enter the sentence you want to compare...", label="Sentence to Compare 1"),
    gr.Textbox(lines=2, placeholder="Enter the sentence you want to compare...", label="Sentence to Compare 2"),
    gr.Textbox(lines=2, placeholder="Enter the sentence you want to compare...", label="Sentence to Compare 3")
]

num_sentences_input = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of Sentences to Compare")

# Function to dynamically update the interface
def update_interface(num_sentences):
    return [gr.Textbox(lines=2, placeholder="Enter the sentence you want to compare...", label=f"Sentence to Compare {i}") for i in range(1, num_sentences + 1)]

def create_interface():
    inputs_dynamic = [model_name_display, original_sentence_input, num_sentences_input]
    inputs_dynamic += update_interface(3)
    return gr.Interface(
        fn=predict,
        title="African Cross-Lingua Embeddings Model's Demo",
        description="Compute the semantic similarity across various sentences among any African Languages using African-Cross-Lingua-Embeddings-Model.",
        inputs=inputs_dynamic,
        outputs=outputs,
        live=False,
        allow_flagging="never"
    )

# Create and launch the initial interface
interface = create_interface()
interface.launch(debug=True, share=True)