from transformers import (
    MarianMTModel,
    MarianTokenizer,
)

class BackTranslator:
    """
    Perform Back-translation with pivot language. English -> German -> Spanish -> English
    Args:
        source_lang: Source language (default: 'en')
        pivot_lang: Pivot language (default: 'de')
        target_lang: Target language (default: 'es')
    Examples:
        back_translator = BackTranslator()
        back_translator.back_translate("Hello, how are you?")
    """
    def __init__(self, source_lang='en', pivot_lang='de', target_lang='es'):
        # Forward (English to German)
        pivot_forward_model_name = f'Helsinki-NLP/opus-mt-{source_lang}-{pivot_lang}'
        self.tokenizer_pivot_forward = MarianTokenizer.from_pretrained(pivot_forward_model_name)
        self.model_pivot_forward = MarianMTModel.from_pretrained(pivot_forward_model_name)
        
        # Pivot translation model (German to Spanish) 
        pivot_backward_model_name = f'Helsinki-NLP/opus-mt-{pivot_lang}-{target_lang}'
        self.tokenizer_pivot_backward = MarianTokenizer.from_pretrained(pivot_backward_model_name)
        self.model_pivot_backward = MarianMTModel.from_pretrained(pivot_backward_model_name)

        # Backward (Spanish to English)
        backward_model_name = f'Helsinki-NLP/opus-mt-{target_lang}-{source_lang}'
        self.tokenizer_backward = MarianTokenizer.from_pretrained(backward_model_name)
        self.model_backward = MarianMTModel.from_pretrained(backward_model_name)

    def back_translate(self, text):
        """
        Perform back-translation through German and Spanish to generate text variations.
        Args:
            text (str): The input text to be back-translated
            
        Returns:
            str: The back-translated text
        """
        # 1. English to German
        encoded_pivot = self.tokenizer_pivot_forward([text], padding=True, truncation=True, return_tensors='pt')
        generated_pivot = self.model_pivot_forward.generate(**encoded_pivot)
        pivot_text = self.tokenizer_pivot_forward.batch_decode(generated_pivot, skip_special_tokens=True)[0]

        # 2. German to Spanish
        encoded_back_pivot = self.tokenizer_pivot_backward([pivot_text], padding=True, truncation=True, return_tensors='pt')
        retranslated_pivot = self.model_pivot_backward.generate(**encoded_back_pivot)
        tgt_text_back = self.tokenizer_pivot_backward.batch_decode(retranslated_pivot, skip_special_tokens=True)[0]

        # 3. Spanish to English
        encoded_back = self.tokenizer_backward([tgt_text_back], padding=True, truncation=True, return_tensors='pt')
        retranslated = self.model_backward.generate(**encoded_back)
        src_text = self.tokenizer_backward.batch_decode(retranslated, skip_special_tokens=True)[0]
        return src_text