import translation_direction_detection as tdd
from nmtscore import NMTScorer
import gradio as gr

detector = tdd.TranslationDirectionDetector(NMTScorer("m2m100_418M"))

def translate_direction(text1, lang1, text2, lang2):
    lang_to_code = {"English": 'en',
        "German": 'de',
        "French": 'fr',
        "Czech": 'cs',
        "Ukrainian": 'uk',
        "Chinese": 'zh',
        "Russian": 'ru',
        "Bengali": 'bn',
        "Hindi": 'hi',
        "Xhosa": 'xh',
        "Zulu": 'zu',
    }
    if "\n" in text1 or "\n" in text2:
        sentence1 = text1.split("\n")
        sentence2 = text2.split("\n")
    else:
        sentence1 = text1
        sentence2 = text2
    result = detector.detect(sentence1, sentence2, lang_to_code[lang1], lang_to_code[lang2])
    return result

iface = gr.Interface(
    fn=translate_direction,
    inputs=[
        gr.Textbox(placeholder="Enter a single sentence or multiple sentences separated by a line break (Shift+Enter) in language 1 here.", label="Text 1"),
        gr.Dropdown(choices=["English", "German", "French", "Czech", "Ukranian", "Chinese", "Russian", "Bengali", "Hindi", "Xhosa", "Zulu"], label="Language of Text 1"),
        gr.Textbox(placeholder="Enter a single sentence or multiple sentences separated by a line break (Shift+Enter) in language 2 here.", label="Text 2"),
        gr.Dropdown(choices=["English", "German", "French", "Czech", "Ukranian", "Chinese", "Russian", "Bengali", "Hindi", "Xhosa", "Zulu"], label="Language of Text 2")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Translation Direction Detector",
    description="Detects the translation direction between two sentences using the M2M100 418M translation model.",
    theme="dark"
)

iface.launch()