"""
BackTranslation class
-----------------------------------

"""

import random

from transformers import MarianMTModel, MarianTokenizer

from textattack.shared import AttackedText

from .sentence_transformation import SentenceTransformation


class BackTranslation(SentenceTransformation):
    """A type of sentence level transformation that takes in a text input,
    translates it into target language and translates it back to source
    language.

    letters_to_insert (string): letters allowed for insertion into words
    (used by some char-based transformations)

    src_lang (string): source language
    target_lang (string): target language, for the list of supported language check bottom of this page
    src_model: translation model from huggingface that translates from source language to target language
    target_model: translation model from huggingface that translates from target language to source language
    chained_back_translation: run back translation in a chain for more perturbation (for example, en-es-en-fr-en)

    Example::

        >>> from textattack.transformations.sentence_transformations import BackTranslation
        >>> from textattack.constraints.pre_transformation import RepeatModification, StopwordModification
        >>> from textattack.augmentation import Augmenter

        >>> transformation = BackTranslation()
        >>> constraints = [RepeatModification(), StopwordModification()]
        >>> augmenter = Augmenter(transformation = transformation, constraints = constraints)
        >>> s = 'What on earth are you doing here.'

        >>> augmenter.augment(s)
    """

    def __init__(
        self,
        src_lang="en",
        target_lang="es",
        src_model="Helsinki-NLP/opus-mt-ROMANCE-en",
        target_model="Helsinki-NLP/opus-mt-en-ROMANCE",
        chained_back_translation=0,
    ):
        self.src_lang = src_lang
        self.target_lang = target_lang
        self.target_model = MarianMTModel.from_pretrained(target_model)
        self.target_tokenizer = MarianTokenizer.from_pretrained(target_model)
        self.src_model = MarianMTModel.from_pretrained(src_model)
        self.src_tokenizer = MarianTokenizer.from_pretrained(src_model)
        self.chained_back_translation = chained_back_translation

    def translate(self, input, model, tokenizer, lang="es"):
        # change the text to model's format
        src_texts = []
        if lang == "en":
            src_texts.append(input[0])
        else:
            if ">>" and "<<" not in lang:
                lang = ">>" + lang + "<< "
            src_texts.append(lang + input[0])

        # tokenize the input
        encoded_input = tokenizer.prepare_seq2seq_batch(src_texts, return_tensors="pt")

        # translate the input
        translated = model.generate(**encoded_input)
        translated_input = tokenizer.batch_decode(translated, skip_special_tokens=True)
        return translated_input

    def _get_transformations(self, current_text, indices_to_modify):
        transformed_texts = []
        current_text = current_text.text

        # to perform chained back translation, a random list of target languages are selected from the provided model
        if self.chained_back_translation:
            list_of_target_lang = random.sample(
                self.target_tokenizer.supported_language_codes,
                self.chained_back_translation,
            )
            for target_lang in list_of_target_lang:
                target_language_text = self.translate(
                    [current_text],
                    self.target_model,
                    self.target_tokenizer,
                    target_lang,
                )
                src_language_text = self.translate(
                    target_language_text,
                    self.src_model,
                    self.src_tokenizer,
                    self.src_lang,
                )
                current_text = src_language_text[0]
            return [AttackedText(current_text)]

        # translates source to target language and back to source language (single back translation)
        target_language_text = self.translate(
            [current_text], self.target_model, self.target_tokenizer, self.target_lang
        )
        src_language_text = self.translate(
            target_language_text, self.src_model, self.src_tokenizer, self.src_lang
        )
        transformed_texts.append(AttackedText(src_language_text[0]))
        return transformed_texts


"""
List of supported languages
['fr',
 'es',
 'it',
 'pt',
 'pt_br',
 'ro',
 'ca',
 'gl',
 'pt_BR<<',
 'la<<',
 'wa<<',
 'fur<<',
 'oc<<',
 'fr_CA<<',
 'sc<<',
 'es_ES',
 'es_MX',
 'es_AR',
 'es_PR',
 'es_UY',
 'es_CL',
 'es_CO',
 'es_CR',
 'es_GT',
 'es_HN',
 'es_NI',
 'es_PA',
 'es_PE',
 'es_VE',
 'es_DO',
 'es_EC',
 'es_SV',
 'an',
 'pt_PT',
 'frp',
 'lad',
 'vec',
 'fr_FR',
 'co',
 'it_IT',
 'lld',
 'lij',
 'lmo',
 'nap',
 'rm',
 'scn',
 'mwl']
"""
