from typing import List, Tuple, Union
from .abstract import Reranker
from tqdm.auto import tqdm

DecodedOutput = Union['torch.Tensor', Tuple['torch.Tensor', 'torch.Tensor']]

# copied from PyGaggle repo
# https://github.com/castorini/pygaggle/blob/master/pygaggle/rerank/transformer.py
prediction_tokens = {
    'castorini/monot5-base-msmarco':             ['▁false', '▁true'],
    'castorini/monot5-base-msmarco-10k':         ['▁false', '▁true'],
    'castorini/monot5-large-msmarco':            ['▁false', '▁true'],
    'castorini/monot5-large-msmarco-10k':        ['▁false', '▁true'],
    'castorini/monot5-base-med-msmarco':         ['▁false', '▁true'],
    'castorini/monot5-3b-med-msmarco':           ['▁false', '▁true'],
    'castorini/monot5-3b-msmarco-10k':           ['▁false', '▁true'],
    # 'unicamp-dl/mt5-13b-mmarco-100k':            ['▁false', '▁true'], # bug
    'unicamp-dl/mt5-13b-mmarco-100k':            ['▁'  , '▁true'],
    'unicamp-dl/mt5-base-en-msmarco':            ['▁no'   , '▁yes'],
    'unicamp-dl/ptt5-base-pt-msmarco-10k-v2':    ['▁não'  , '▁sim'],
    'unicamp-dl/ptt5-base-pt-msmarco-100k-v2':   ['▁não'  , '▁sim'],
    'unicamp-dl/ptt5-base-en-pt-msmarco-100k-v2':['▁não'  , '▁sim'],
    'unicamp-dl/mt5-base-en-pt-msmarco-v2':      ['▁no'   , '▁yes'],
    'unicamp-dl/mt5-base-mmarco-v2':             ['▁no'   , '▁yes'],
    'unicamp-dl/mt5-base-en-pt-msmarco-v1':      ['▁no'   , '▁yes'],
    'unicamp-dl/mt5-base-mmarco-v1':             ['▁no'   , '▁yes'],
    'unicamp-dl/ptt5-base-pt-msmarco-10k-v1':    ['▁não'  , '▁sim'],
    'unicamp-dl/ptt5-base-pt-msmarco-100k-v1':   ['▁não'  , '▁sim'],
    'unicamp-dl/ptt5-base-en-pt-msmarco-10k-v1': ['▁não'  , '▁sim'],
    'unicamp-dl/mt5-3B-mmarco-en-pt':            ['▁'  , '▁true']
}


def _get_prediction_tokens(pretrained_model_name_or_path: str,
                           tokenizer: 'PreTrainedTokenizer', token_false=None, token_true=None):
    if not (token_false and token_true):
        if pretrained_model_name_or_path in prediction_tokens:
            token_false, token_true = prediction_tokens[pretrained_model_name_or_path]
            token_false_id = tokenizer.get_vocab()[token_false]
            token_true_id  = tokenizer.get_vocab()[token_true]
            return token_false_id, token_true_id
        else:
            raise Exception(f"We don't know the indexes for the non-relevant/relevant tokens for\
                    the checkpoint {pretrained_model_name_or_path} and you did not provide any.")
    else:
        token_false_id = tokenizer.get_vocab()[token_false]
        token_true_id  = tokenizer.get_vocab()[token_true]
        return token_false_id, token_true_id


class MT5Reranker(Reranker):
    prompt = 'Query: {query} Document: {document} Relevant:'

    def __init__(self, name=None, config=None, **kwargs):
        super().__init__(name, config, **kwargs)

        from transformers import T5ForConditionalGeneration, AutoTokenizer, PreTrainedTokenizer

        assert 'model_name_or_path' in self.config
        self.q_max_length = int(self.config.get('q_max_length', 180))
        self.d_max_length = int(self.config.get('d_max_length', 512))
        self.batch_size = int(self.config.get('batch_size', 32))

        self.model = T5ForConditionalGeneration.from_pretrained(self.config['model_name_or_path'])
        self.tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(self.config['model_name_or_path'], use_fast=True)
        self.model.eval()

        if self.config.get('use_gpu', True):
            self.model = self.model.bfloat16().cuda()

        self.token_false_id, self.token_true_id  = _get_prediction_tokens(
            self.config['model_name_or_path'], self.tokenizer
        )

    def tokenize_pairs(self, batch: List[Tuple[str, str]]) -> 'BatchEncoding':
        return self.tokenizer(
            [ self.prompt.format(query=q, document=d) for q, d in batch ],
            return_tensors='pt',
            return_attention_mask=True, 
            padding='max_length',
            truncation=True,
            max_length=self.q_max_length + self.d_max_length + 9
        )

    # copied and modified from PyGaggle repo
    # https://github.com/castorini/pygaggle/blob/master/pygaggle/model/decode.py
    def forward(
            self, input_ids: 'torch.Tensor', attention_mask: 'torch.Tensor' = None, 
        ) -> DecodedOutput:
        encoder_outputs = self.model.get_encoder()(input_ids, attention_mask=attention_mask)
        
        import torch
        with torch.no_grad():
            decode_ids = torch.full(
                (input_ids.size(0), 1),
                self.model.config.decoder_start_token_id,
                dtype=torch.long, device=input_ids.device
            )

            model_inputs = self.model.prepare_inputs_for_generation(
                decode_ids,
                encoder_outputs=encoder_outputs,
                past_key_values=None,
                attention_mask=attention_mask,
                use_cache=True
            )
            outputs = self.model(**model_inputs)  # (batch_size, cur_len, vocab_size)
            true_false_logits = outputs[0][:, -1, [self.token_false_id, self.token_true_id]]  # (batch_size, vocab_size)

            return torch.nn.functional.log_softmax(true_false_logits, dim=1)[:, 1].contiguous()
        
    async def score(self, queries, passages, candidate_length = None, with_progress=False, **kwargs):
        if candidate_length is None:
            if len(queries) == len(passages): # pairwise
                candidate_length = [1]*len(queries)
            if len(passages) % len(queries) == 0: # assume same number of candidates
                candidate_length = [len(passages)//len(queries)]*len(queries)
            else:
                raise ValueError("mT5 Reranker does not support all pair scoring.")
        
        assert len(candidate_length) == len(queries)
        assert sum(candidate_length) == len(passages)

        import torch
        expanded_queries = sum([
            [query]*l for query, l in zip(queries, candidate_length)
        ], [])
        qidx = sum([
            [qid]*l for qid, l in zip(range(len(queries)), candidate_length)
        ], [])

        all_scores = torch.concat([
            self.forward(
                **self.tokenize_pairs(list(zip(
                    expanded_queries[i: i+self.batch_size], passages[i: i+self.batch_size]
                ))).to(self.model.device)
            ).cpu()
            for i in tqdm(range(0, len(passages), self.batch_size), 
                          disable=not with_progress, dynamic_ncols=True)
        ]).flatten().tolist()

        offsets = [0] + [ sum(candidate_length[:i+1]) for i in range(len(candidate_length)) ]

        return [
            all_scores[sidx: eidx] for sidx, eidx in zip(offsets[:-1], offsets[1:])
        ]

