import os
import gradio as gr
from simpletransformers.seq2seq import Seq2SeqModel, Seq2SeqArgs


# os.environ["TOKENIZERS_PARALLELISM"] = "false"


def load_translator(model_name='Enutrof/marian-mt-en-pcm'):
    '''
    This method loads the sequence to sequence model for translation.
    :return: model
    '''
    pmodel_args = Seq2SeqArgs()
    pmodel_args.max_length = 1024
    pmodel_args.length_penalty = 1
    pmodel_args.num_beams = 20
    pmodel_args.num_return_sequences = 3

    pmodel = Seq2SeqModel(
        encoder_decoder_type="marian",
        encoder_decoder_name=model_name,
        args=pmodel_args,
        use_cuda=False
    )
    return pmodel


en_pcm_model = load_translator()


def predict(input):
    if isinstance(input, str):
        input = [input]
    predictions = en_pcm_model.predict(input)
    return [i.replace('▁', ' ') for i in predictions[0]]


# HF_TOKEN = os.getenv('english-pidgin-flagging')
# hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN,
#                                        dataset_name="English-NigerianPidgin-Result-Validation",
#                                        organization="Enutrof",
#                                        )
gr.Interface(
    fn=predict,
    inputs=gr.inputs.Textbox(lines=1, label="Input Text in English"),
    outputs=[
        gr.outputs.Textbox(label="Translated texts in 🇳🇬 Pidgin"),
        gr.outputs.Textbox(label=''),
        gr.outputs.Textbox(label=''),
    ],
    # theme="peach",
    title='English to 🇳🇬 Pidgin Automatic Translation',
    description='Type your English text in the left text box to get 🇳🇬 Pidgin translations on the right. '
                'Tell us the best translation by clicking one of the buttons below.',
    examples=[
        'Who are you?',
        'You shall not pervert justice due the stranger or the fatherless, nor take a widow’s garment as a pledge.',
        'I know every song by that artiste.',
        'They should not be permitted here.',
        'What are you looking for?',
        'I am lost please help me find my way to the market.',
    ],
    allow_flagging="manual",
    flagging_options=["translation 1 ✅", "translation 2 ✅",
                      "translation 3 ✅"],
    #flagging_callback=hf_writer,
).launch(enable_queue=True)