"""
translation program for simple text
1. detect language from langdetect
2. translate to target language given by user

Example from
https://www.thepythoncode.com/article/machine-translation-using-huggingface-transformers-in-python 

user_input:
    string: string to be translated
    target_lang: language to be translated to

Returns:
    string: translated string of text
"""

import gradio as gr
import argparse

import langid
from transformers import pipeline


def detect_lang(article, target_lang):
    """
    Language Detection using library langid

    Args:
        article (string): article that user wish to translate
        target_lang (string): language user want to translate article into

    Returns:
        string: detected language short form
    """

    result_lang = langid.classify(article)
    return result_lang[0]


def opus_trans(article, target_language):
    """
    Translation by Helsinki-NLP model

    Args:
        article (string): article that user wishes to translate
        result_lang (string): detected language in short form
        target_language (string): language that user wishes to translate article into

    Returns:
        string: translated piece of article based off target_language
    """
    result_lang = detect_lang(article, target_language)

    if target_language == "English":
        target_lang = "en"
    elif target_language == "Chinese":
        target_lang = "zh"
    # result_lang = detect_lang(article, target_language)

    if result_lang != target_lang:
        task_name = f"translation_{result_lang}_to_{target_lang}"
        model_name = f"Helsinki-NLP/opus-mt-{result_lang}-{target_lang}"
        translator = pipeline(task_name, model=model_name, tokenizer=model_name)
        translated = translator(article)[0]["translation_text"]
        # print(translated)
    else:
        translated = "Error"
    return translated


article = gr.Textbox()
lang_select = gr.Radio(["English", "Chinese"])
translate = gr.Interface(
    opus_trans,
    [
        article,
        lang_select,
    ],
    outputs=gr.Textbox(),
)

translate.launch()