import torch
from torch.utils.data.dataset import TensorDataset
from transformers import AutoTokenizer, AutoModel
from torch.utils.data import DataLoader
from tqdm import tqdm

from parlai.agents.programr.utils.logging.ylogger import YLogger, timed_task


class TextEmbedder:
    def __init__(self, tokenizer_filepath, model_filepath):
        self._tokenizer = AutoTokenizer.from_pretrained(tokenizer_filepath)
        self._model = AutoModel.from_pretrained(model_filepath)

    @property
    def tokenizer(self):
        return self._tokenizer

    @property
    def model(self):
        return self._model


    @timed_task
    def create_sentence_embeddings_no_batching(self, document):
        try:
            #Tokenize questions
            encoded_input = self._tokenizer(document, padding=True, truncation=True, max_length=128, return_tensors='pt')

            # if torch.cuda.is_available():
            #     encoded_input = encoded_input.to('cuda')

            with torch.no_grad():
                model_output = self._model(**encoded_input)
            
            #Perform pooling. In this case, mean pooling
            pooled_embeddings = self._mean_pooling(model_output, encoded_input['attention_mask'])

            return pooled_embeddings
        except Exception as e:
            print("Error creating sentence embeddings - {}".format(e))

    @timed_task
    def create_sentence_embeddings(self, document):
        try:
            # Tokenize questions
            encoded_input = self._tokenizer(document, padding=True, truncation=True, max_length=128,
                                            return_tensors='pt')

            if torch.cuda.is_available():
                print("Putting encoded_input onto cuda.")
                encoded_input = encoded_input.to('cuda:0')

            dataloader = DataLoader(TensorDataset(encoded_input['input_ids'],
                                                  encoded_input['token_type_ids'],
                                                  encoded_input['attention_mask']),
                                    batch_size=1,
                                    shuffle=False, num_workers=0)

            all_pooled_embeddings = self._compute_token_embedding(dataloader)

            return all_pooled_embeddings
        except Exception as e:
            print("Error creating sentence embeddings - {}".format(e))
            raise Exception

    """Private Methods"""
    def _compute_token_embedding(self, dataloader):
        all_pooled_embeddings = []
        with torch.no_grad():
            for i, data in tqdm(enumerate(dataloader)):
                embedded_batch = self._model(input_ids=data[0], token_type_ids=data[1], attention_mask=data[2])
                # all_pooled_embeddings.append(embedded_batch['pooler_output'])

                pooled_embeddings = self._mean_pooling(embedded_batch, data[2])
                all_pooled_embeddings.append(pooled_embeddings)

        all_pooled_embeddings = torch.cat(all_pooled_embeddings, dim=0)
        return all_pooled_embeddings

    # Mean Pooling - Take attention mask into account for correct averaging
    def _mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        return sum_embeddings / sum_mask