import gradio as gr
import torch
import numpy as np
import fasttext
import os
import urllib
import huggingface_hub
from transformers import MBartForConditionalGeneration, MBart50Tokenizer


MODEL_URL_MYV_MUL = 'slone/mbart-large-51-myv-mul-v1'
MODEL_URL_MUL_MYV = 'slone/mbart-large-51-mul-myv-v1'
MODEL_URL_LANGID = 'https://huggingface.co/slone/fastText-LID-323/resolve/main/lid.323.ftz'
MODEL_PATH_LANGID = 'lid.323.ftz'


HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(
    hf_token=HF_TOKEN,
    dataset_name="myv-translation-2022-demo-flags-v2",
    organization="slone",
    private=True,
)


lang_to_code = {
    'Эрзянь | Erzya': 'myv_XX',
    'Русский | Рузонь | Russian': 'ru_RU',
    'Suomi | Суоминь | Finnish': 'fi_FI',
    'Deutsch | Немецень | German': 'de_DE',
    'Español | Испанонь | Spanish': 'es_XX',
    'English | Англань ': 'en_XX',
    'हिन्दी | Хинди | Hindi': 'hi_IN',
    '漢語 | Китаень | Chinese': 'zh_CN',
    'Türkçe | Турконь | Turkish': 'tr_TR',
    'Українська | Украинань | Ukrainian': 'uk_UA',
    'Français | Французонь | French': 'fr_XX',
    'العربية | Арабонь | Arabic': 'ar_AR',
}


def fix_tokenizer(tokenizer, extra_lang='myv_XX'):
    """Add a new language id to a MBART 50 tokenizer (because it is not serialized) and shift the mask token id."""
    old_len = len(tokenizer) - int(extra_lang in tokenizer.added_tokens_encoder)
    tokenizer.lang_code_to_id[extra_lang] = old_len-1
    tokenizer.id_to_lang_code[old_len-1] = extra_lang
    tokenizer.fairseq_tokens_to_ids["<mask>"] = len(tokenizer.sp_model) + len(tokenizer.lang_code_to_id) + tokenizer.fairseq_offset

    tokenizer.fairseq_tokens_to_ids.update(tokenizer.lang_code_to_id)
    tokenizer.fairseq_ids_to_tokens = {v: k for k, v in tokenizer.fairseq_tokens_to_ids.items()}
    if extra_lang not in tokenizer._additional_special_tokens:
        tokenizer._additional_special_tokens.append(extra_lang)
    tokenizer.added_tokens_encoder = {}


def translate(
        text, model, tokenizer,
        src='ru_RU',
        trg='myv_XX',
        max_length='auto',
        num_beams=3,
        repetition_penalty=5.0,
        train_mode=False, n_out=None,
        **kwargs
):
    tokenizer.src_lang = src
    encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
    if max_length == 'auto':
        max_length = int(32 + 1.5 * encoded.input_ids.shape[1])
    if train_mode:
        model.train()
    else:
        model.eval()
    generated_tokens = model.generate(
        **encoded.to(model.device),
        forced_bos_token_id=tokenizer.lang_code_to_id[trg],
        max_length=max_length,
        num_beams=num_beams,
        repetition_penalty=repetition_penalty,
        # early_stopping=True,
        num_return_sequences=n_out or 1,
        **kwargs
    )
    out = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    if isinstance(text, str) and n_out is None:
        return out[0]
    return out


def translate_rerank(
    text, model, tokenizer,
    src='ru_RU', trg='myv_XX', max_length='auto', num_beams=3, repetition_penalty=5.0, train_mode=False,
    n=5, diversity_penalty=3.0, lang='myv', max_score=0.3, order_penalty=0.01,
    verbose=False,
    **kwargs
):
    texts = translate(
        text, model, tokenizer, src, trg,
        max_length=max_length, train_mode=train_mode, repetition_penalty=repetition_penalty,
        num_beams=n,
        num_beam_groups=n,
        diversity_penalty=diversity_penalty,
        n_out=n,
        **kwargs
    )
    scores = [get_mean_lang_score(t, lang=lang, max_score=max_score) for t in texts]
    pen_scores = scores - order_penalty * np.arange(n)
    if verbose:
        print(texts)
        print(scores)
        print(pen_scores)
    return texts[np.argmax(pen_scores)]


def get_mean_lang_score(text, lang='myv', k=300, max_score=0.3):
    words = text.split() + [text]
    res = []
    for langs, scores in zip(*langid_model.predict(words, k=k)):
        d = dict(zip([l[9:] for l in langs], scores))
        score = min(d.get(lang, 0), max_score) / max_score
        res.append(score)
    # print(res)
    return np.mean(res)


def translate_wrapper(text, src, trg, correct=None):
    src = lang_to_code.get(src)
    trg = lang_to_code.get(trg)
    if src == trg:
        return 'Please choose two different languages'
    if src == 'myv_XX':
        model = model_myv_mul
    elif trg == 'myv_XX':
        model = model_mul_myv
    else:
        return 'Please translate to or from Erzya'
    print(text, src, trg)
    fn = translate_rerank if trg == 'myv_XX' else translate
    result = fn(
        text=text,
        model=model,
        tokenizer=tokenizer,
        src=src,
        trg=trg,
    )
    return result


article = """
Те эрзянь кельсэ автоматической васенце ютавтыця. Тонавкстнэ улить – [сёрмадовкссо](https://arxiv.org/abs/2209.09368).

Это первый автоматический переводчик для эрзянского языка. Подробности – в [статье](https://arxiv.org/abs/2209.09368). 
Пожалуйста, оставляйте своё мнение о качестве переводов с помощью кнопок с эмодзи!

This is the first automatic translator for the Erzya language. The details are in the [paper](https://arxiv.org/abs/2209.09368).
Please leave your feedback about the quality of translations using the buttons with emojis.

The code and models for translation can be found in the repository: https://github.com/slone-nlp/myv-nmt. 
"""

fix_instruction = 'Если перевод модели неправильный, впишите сюда правильный текст, снова нажмите "Исполнить", и затем "bad 🙁". ' \
                  'Тогда к нам в базу попадёт пометка, что перевод был неверным, и его исправление.'


interface = gr.Interface(
    translate_wrapper,
    [
        gr.Textbox(label="Text / текст", lines=2, placeholder='text to translate / текст ютавтозь'),
        gr.Dropdown(list(lang_to_code.keys()), type="value", label='source language / васенце кель', value=list(lang_to_code.keys())[0]),
        gr.Dropdown(list(lang_to_code.keys()), type="value", label='target language / эрявикс кель', value=list(lang_to_code.keys())[1]),
        gr.Textbox(label="Correct translation", lines=2, placeholder=fix_instruction),
    ],
    "text",
    allow_flagging="manual",
    flagging_options=["good 🙂", "50/50 😐", "bad 🙁"],
    flagging_callback=hf_writer,
    title='Эрзянь ютавтыця | Эрзянский переводчик | Erzya translator',
    article=article,
)


if __name__ == '__main__':
    model_mul_myv = MBartForConditionalGeneration.from_pretrained(MODEL_URL_MUL_MYV)
    model_myv_mul = MBartForConditionalGeneration.from_pretrained(MODEL_URL_MYV_MUL)
    if torch.cuda.is_available():
        model_mul_myv.cuda()
        model_myv_mul.cuda()
    tokenizer = MBart50Tokenizer.from_pretrained(MODEL_URL_MYV_MUL)
    fix_tokenizer(tokenizer)

    if not os.path.exists(MODEL_PATH_LANGID):
        print('downloading LID model...')
        urllib.request.urlretrieve(MODEL_URL_LANGID, MODEL_PATH_LANGID)
    langid_model = fasttext.load_model(MODEL_PATH_LANGID)

    interface.launch()