from __future__ import annotations
from abc import abstractmethod
from collections import OrderedDict
from enum import Enum
from typing import Any, Dict, List, Union
from dataclasses_json import config

import paddle
from paddle.metric import Metric
from paddlenlp.transformers.tokenizer_utils import PretrainedTokenizer
from paddle_prompt.config import Tensor, Config
from paddle_prompt.schema import InputFeature
from paddle_prompt.utils import get_metric, to_list


class TokenHandler(Enum):
    first: str = 'first'
    max: str = 'max'
    mean: str = 'mean'

class Verbalizer:
    def __init__(
        self,
        tokenizer: PretrainedTokenizer,
        label_map: Dict[str, Union[str, List[str]]],
        config: Config,
        multi_token_handler: TokenHandler = TokenHandler.mean,
    ) -> None:
        self.tokenizer = tokenizer
        self.label_map = label_map

        self.label_words_ids_dict: Dict[str, List[List[int]]] = OrderedDict()
        
        label_words_ids_tensor = []
        for label, words in label_map.items():
            self.label_words_ids_dict[label] = self._map_label_words_to_label_ids(words)
            label_words_ids_tensor.append(self.label_words_ids_dict[label])
        
        self.multi_token_handler = multi_token_handler
        self.config = config
    
    def _map_label_words_to_label_ids(self, words: Union[str, List[str]]) -> List[List[int]]:
        if isinstance(words, str):
            words = [words]

        label_ids = []
        for word in words:
            encoded_features = self.tokenizer.encode(
                word,
                return_token_type_ids=False
            )
            # in paddlenlp, there always special token in the sentence, so we should remove it
            label_ids.append(encoded_features['input_ids'][1: -1])
        return label_ids

    @abstractmethod    
    def project(self, mask_label_logits: Tensor) -> Tensor:
        """project mask label logits to label distribution

        Args:
            mask_label_logits (Tensor): the output mask label logit from PLM 
                shape: [batch_size, max_token_num, vocab_size]

        Returns:
            Tensor: the mask label distribution 
                shape: [batch_size, label_num]
        """
        raise NotImplementedError

    def process_outputs(self,
                       outputs: Tensor,
                       batch,
                       **kwargs):
        r"""By default, the verbalizer will process the logits of the PLM's 
        output. 

        Args:
            logits (:obj:`torch.Tensor`): The current logits generated by pre-trained language models.
            batch (:obj:`Union[Dict, InputFeatures]`): The input features of the data.
        """

        return self.process_logits(outputs, batch=batch, **kwargs)

    @abstractmethod
    def process_logits(self, outputs: Tensor, **kwargs):
        """process the logit from the pretrained mlm mode with batch data

        Args:
            outputs (Tensor): the output from the 
            batch_data (BatchData): # TODO: define the batch data strucuture
        """
        raise NotImplementedError

    def replace_trigger_tokens(self, model_inputs, trigger_ids, trigger_mask):
        """Replaces the trigger tokens in input_ids."""
        out = model_inputs.copy()
        input_ids = model_inputs['input_ids']
        trigger_ids = trigger_ids.repeat(trigger_mask.size(0), 1)
        try:
            filled = input_ids.masked_scatter(trigger_mask, trigger_ids)
        except RuntimeError:
            filled = input_ids
        out['input_ids'] = filled
        return out

    def handle_multi_token(self, label_words_logits: Tensor, mask: Tensor):
        r"""
        Support multiple methods to handle the multi tokens produced by the tokenizer.
        We suggest using 'first' or 'max' if the some parts of the tokenization is not meaningful.
        Can broadcast to 3-d tensor.
    
        Args:
            label_words_logits (:obj:`torch.Tensor`):
        
        Returns:
            :obj:`torch.Tensor`
        """
        # TODO: to be tested
        if self.multi_token_handler == TokenHandler.first:
            label_words_logits = label_words_logits.index_select(axis=-1, index=0)
        elif self.multi_token_handler == TokenHandler.max:
            label_words_logits = label_words_logits - 1000*(1-mask.unsqueeze(0))
            label_words_logits = label_words_logits.max(dim=-1).values
        elif self.multi_token_handler == TokenHandler.mean:
            label_words_logits = (label_words_logits*mask.unsqueeze(0)).sum(dim=-1)/(mask.unsqueeze(0).sum(dim=-1)+1e-15)
        else:
            raise ValueError("multi_token_handler {} not configured".format(self.multi_token_handler))
        return label_words_logits
    
    def aggregate(self, label_words_logits: Tensor) -> Tensor:
        r"""Use weight to aggregate the logits of label words.

        Args:
            label_words_logits(:obj:`torch.Tensor`): The logits of the label words.
        
        Returns:
            :obj:`torch.Tensor`: The aggregated logits from the label words. 
        """
        label_words_logits = (label_words_logits * self.label_words_mask).sum(-1)/self.label_words_mask.sum(-1)
        return label_words_logits

    def calibrate(self, label_words_probs: Tensor, **kwargs) -> Tensor:
        r"""
        
        Args:
            label_words_probs (:obj:`torch.Tensor`): The probability distribution of the label words with the shape of [``batch_size``, ``num_classes``, ``num_label_words_per_class``]
        
        Returns:
            :obj:`torch.Tensor`: The calibrated probability of label words.
        """
        shape = label_words_probs.shape
        assert self._calibrate_logits.dim() ==  1, "self._calibrate_logits are not 1-d tensor"
        calibrate_label_words_probs = self.normalize(self.project(self._calibrate_logits.unsqueeze(0), **kwargs))
        assert calibrate_label_words_probs.shape[1:] == label_words_probs.shape[1:] \
             and calibrate_label_words_probs.shape[0]==1, "shape not match"
        label_words_probs /= (calibrate_label_words_probs+1e-15)

        # normalize # TODO Test the performance
        norm = label_words_probs.reshape(shape[0], -1).sum(dim=-1,keepdim=True) # TODO Test the performance of detaching()
        label_words_probs = label_words_probs.reshape(shape[0], -1) / norm
        label_words_probs = label_words_probs.reshape(*shape)
        return label_words_probs
    
    def compute_metric(self, mask_label_logits: Tensor, label_ids: Tensor, metric_name: str = 'acc') -> float:
        """compute the acc based on the logits and label_ids

        Args:
            mask_label_logits (Tensor): the logits from the PLM
                [batch_size, max_token_num, vocab_size]
            label_ids (Tensor): the golden truth of label ids
                [batch_size, max_token_num]

        Returns:
            float: the final label_ids
        """
        assert len(mask_label_logits) == 3, 'the shape mask_label_logits should be 3-dim'

        label_logits = self.project(mask_label_logits)
        metric: Metric = get_metric(self.config.metric_name)
        metric.update(label_logits, label_ids)
        
        return metric.accumulate()
        

        