import torch
from transformers import T5Tokenizer


class T5TextConditionProcessor:

    def __init__(self, tokens_length, processor_path):
        self.tokens_length = tokens_length
        self.processor = T5Tokenizer.from_pretrained(processor_path)

    def encode(self, text=None, negative_text=None):
        encoded = self.processor(text, max_length=self.tokens_length, truncation=True)
        pad_length = self.tokens_length - len(encoded['input_ids'])
        input_ids = encoded['input_ids'] + [self.processor.pad_token_id] * pad_length
        attention_mask = encoded['attention_mask'] + [0] * pad_length
        condition_model_input = {
            'input_ids': torch.tensor(input_ids, dtype=torch.long),
            'attention_mask': torch.tensor(attention_mask, dtype=torch.long)
        }

        if negative_text is not None:
            negative_encoded = self.processor(negative_text, max_length=self.tokens_length, truncation=True)
            negative_input_ids = negative_encoded['input_ids'][:len(encoded['input_ids'])]
            negative_input_ids[-1] = self.processor.eos_token_id
            negative_pad_length = self.tokens_length - len(negative_input_ids)
            negative_input_ids = negative_input_ids + [self.processor.pad_token_id] * negative_pad_length
            negative_attention_mask = encoded['attention_mask'] + [0] * pad_length
            negative_condition_model_input = {
                'input_ids': torch.tensor(negative_input_ids, dtype=torch.long),
                'attention_mask': torch.tensor(negative_attention_mask, dtype=torch.long)
            }
        else:
            negative_condition_model_input = None
        return condition_model_input, negative_condition_model_input