"""MDL Retriever."""

from typing import List, Optional

import numpy as np
import torch
import tqdm
from transformers import AutoModelForCausalLM

from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever.icl_topk_retriever import TopkRetriever
from opencompass.openicl.utils.logging import get_logger
from opencompass.registry import ICL_PROMPT_TEMPLATES, ICL_RETRIEVERS

logger = get_logger(__name__)


@ICL_RETRIEVERS.register_module()
class MDLRetriever(TopkRetriever):
    """MDL Retriever, subclass of `TopkRetriever`. MDL is a abbreviation of
    Minimum Description Length, specially designed for ppl evaluation. You may
    refer to the paper for more details: https://arxiv.org/pdf/2212.10375.pdf.

    Args:
        dataset (`BaseDataset`): Any BaseDataset instances.
            Attributes of ``reader``, ``train`` and ``test`` will be used.
        ice_separator (`Optional[str]`): The separator between each in-context
            example template when origin `PromptTemplate` is provided. Defaults
            to '\n'.
        ice_eos_token (`Optional[str]`): The end of sentence token for
            in-context example template when origin `PromptTemplate` is
            provided. Defaults to '\n'.
        ice_num (`Optional[int]`): The number of in-context example template
            when origin `PromptTemplate` is provided. Defaults to 1.
        sentence_transformers_model_name (`Optional[str]`): The name of the
            sentence transformers model. Defaults to 'all-mpnet-base-v2'.
        tokenizer_name (`Optional[str]`): The name of the tokenizer. Defaults
            to 'gpt2-xl'.
        batch_size (`Optional[int]`): The batch size for the dataloader.
            Defaults to 1.
        candidate_num (`Optional[int]`): The number of candidates to retrieve
            for each example. Defaults to 1.
        ce_model_name (`Optional[str]`): The name of the model for calculating
            MDL. Defaults to 'gpt2-xl'.
        select_time (`Optional[int]`): The number of times to select MDL.
            Defaults to 5.
        ice_template (`Optional[PromptTemplate]`): The template for in-context
            example. Defaults to None.
        prompt_template (`Optional[PromptTemplate]`): The template for prompt.
            Defaults to None.
        labels (`Optional[List]`): The labels for calculating MDL. Defaults to
            None.
        seed (`Optional[int]`): The seed for random. Defaults to 1.
    """
    metric_model = None

    def __init__(self,
                 dataset,
                 ice_separator: Optional[str] = '\n',
                 ice_eos_token: Optional[str] = '\n',
                 ice_num: Optional[int] = 1,
                 sentence_transformers_model_name: Optional[
                     str] = 'all-mpnet-base-v2',
                 tokenizer_name: Optional[str] = 'gpt2-xl',
                 batch_size: Optional[int] = 1,
                 candidate_num: Optional[int] = 1,
                 ce_model_name: Optional[str] = 'gpt2-xl',
                 select_time: Optional[int] = 5,
                 ice_template: Optional[PromptTemplate] = None,
                 prompt_template: Optional[PromptTemplate] = None,
                 labels: Optional[List] = None,
                 seed: Optional[int] = 1) -> None:
        super().__init__(dataset, ice_separator, ice_eos_token, ice_num,
                         sentence_transformers_model_name, tokenizer_name,
                         batch_size)
        self.ce_model_name = ce_model_name
        self.candidate_num = candidate_num
        self.select_time = select_time
        self.ice_template = ICL_PROMPT_TEMPLATES.build(ice_template)
        if prompt_template is not None:
            self.prompt_template = ICL_PROMPT_TEMPLATES.build(prompt_template)
        else:
            self.prompt_template = None
        self.labels = labels
        self.seed = seed

    def topk_search(self):
        np.random.seed(self.seed)
        res_list = self.forward(self.dataloader)
        rtr_idx_list = [[] for _ in range(len(res_list))]

        logger.info('Retrieving data for test set...')
        for entry in tqdm.tqdm(res_list, disable=not self.is_main_process):
            idx = entry['metadata']['id']
            embed = np.expand_dims(entry['embed'], axis=0)
            near_ids = self.index.search(
                embed, min(self.candidate_num,
                           len(self.index_ds)))[1][0].tolist()
            candidates = []
            mdl_scores = []
            for j in range(self.select_time):
                if j == 0:
                    rand_idx_list = near_ids[:self.ice_num]
                else:
                    rand_idx_list = np.random.choice(near_ids,
                                                     self.ice_num,
                                                     replace=False)
                    rand_idx_list = [int(i) for i in rand_idx_list]
                candidates.append(rand_idx_list)

                ice = self.generate_ice(rand_idx_list,
                                        ice_template=self.ice_template)
                ice = str(ice)
                mask_length = len(
                    self.tokenizer(ice + self.ice_eos_token,
                                   verbose=False)['input_ids'])
                if self.labels is None:
                    labels = self.get_labels(self.ice_template,
                                             self.prompt_template)
                else:
                    labels = self.labels
                prompt_list = []
                for label in labels:
                    prompt = self.generate_label_prompt(
                        idx, ice, label, self.ice_template,
                        self.prompt_template)
                    prompt = str(prompt)
                    prompt_list.append(prompt)
                loss_list = self.cal_ce(prompt_list, mask_length=mask_length)
                probs = np.exp(-np.array(loss_list))
                normalized_probs = probs / probs.sum(0, keepdims=True)
                neg_entropy = -entropy(normalized_probs, label_dim=0)
                mdl_scores.append(neg_entropy)

            rtr_idx_list[idx] = candidates[mdl_scores.index(max(mdl_scores))]
            rtr_idx_list[idx] = [int(i) for i in rtr_idx_list[idx]]

        return rtr_idx_list

    def retrieve(self):
        """Retrieve the in-context example index for each test example."""
        return self.topk_search()

    def cal_ce(self, input_texts: List[str], mask_length=None):
        if self.metric_model is None:
            logger.info(
                f'Load model {self.ce_model_name} for calculating MDL...')
            self.metric_model = AutoModelForCausalLM.from_pretrained(
                self.ce_model_name)
            self.metric_model.to(self.device)
        inputs = self.tokenizer(input_texts,
                                padding=True,
                                return_tensors='pt',
                                truncation=True)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        outputs = self.metric_model(**inputs)

        shift_logits = outputs.logits[..., :-1, :].contiguous()
        shift_labels = inputs['input_ids'][..., 1:].contiguous()

        loss_fct = torch.nn.CrossEntropyLoss(
            reduction='none', ignore_index=self.tokenizer.pad_token_id)
        shift_logits = shift_logits.view(-1, shift_logits.size(-1))
        loss = loss_fct(shift_logits,
                        shift_labels.view(-1)).view(shift_labels.size())
        if mask_length is not None:
            mask = torch.cat([
                torch.zeros([loss.shape[0], mask_length], dtype=torch.float),
                torch.ones([loss.shape[0], loss.shape[-1] - mask_length],
                           dtype=torch.float)
            ], -1)
            mask = mask.to(self.device)
            loss = torch.mul(mask, loss)

        lens = (inputs['input_ids'] !=
                self.tokenizer.pad_token_id).sum(-1).cpu().numpy()
        if mask_length is not None:
            lens -= mask_length
        ce_loss = loss.sum(-1).cpu().detach().numpy() / lens
        return ce_loss


def entropy(probs: np.array, label_dim: int = 0, mask=None):
    if mask is None:
        return -(probs * np.log(probs)).sum(label_dim)
    return -(mask * probs * np.log(probs)).sum(label_dim)
