import numpy as np
import torch

from .ctc_postprocess import BaseRecLabelDecode


class CharLabelDecode(BaseRecLabelDecode):
    """Convert between text-label and text-index."""

    def __init__(self,
                 character_dict_path=None,
                 use_space_char=True,
                 **kwargs):
        super(CharLabelDecode, self).__init__(character_dict_path,
                                              use_space_char)

    def __call__(self, preds, label=None, *args, **kwargs):
        if len(preds) >= 4:
            preds_id = preds[0]
            preds_prob = preds[1]
            char_preds = preds[2]
            if isinstance(preds_id, torch.Tensor):
                preds_id = preds_id.numpy()
            if isinstance(preds_prob, torch.Tensor):
                preds_prob = preds_prob.numpy()
            if preds_id[0][0] == 2:
                preds_idx = preds_id[:, 1:]
                preds_prob = preds_prob[:, 1:]
                # char_preds = char_preds[:, 1:]
            else:
                preds_idx = preds_id
            char_preds = char_preds.numpy()
            char_preds_idx = char_preds.argmax(-1) + 4
            char_preds_prob = char_preds.max(-1)
            text, text_box = self.decode(preds_idx, preds_prob, char_preds_idx,
                                         char_preds_prob)
        else:
            preds_logit = preds[0].numpy()
            char_preds = preds[1].numpy()
            # if isinstance(preds, torch.Tensor):
            #     preds = preds.numpy()
            preds_idx = preds_logit.argmax(axis=2)
            preds_prob = preds_logit.max(axis=2)
            char_preds_idx = char_preds.argmax(-1) + 4
            char_preds_prob = char_preds.max(-1)
            text, text_box = self.decode(preds_idx, preds_prob, char_preds_idx,
                                         char_preds_prob)

        if label is None:
            return text, text_box
        label = self.decode(label[:, 1:])
        return text, text_box, label

    def add_special_char(self, dict_character):
        dict_character = ['blank', '<unk>', '<s>', '</s>'] + dict_character
        return dict_character

    def decode(
        self,
        text_index,
        text_prob=None,
        char_text_index=None,
        char_text_prob=None,
        is_remove_duplicate=False,
    ):
        """convert text-index into text-label."""
        result_list = []
        box_result_list = []
        batch_size = len(text_index)
        for batch_idx in range(batch_size):
            char_list = []
            conf_list = []
            char_box_list = []
            conf_box_list = []
            for idx in range(len(text_index[batch_idx])):
                try:
                    char_idx = self.character[int(text_index[batch_idx][idx])]
                    if char_text_index is not None:
                        char_box_idx = self.character[int(
                            char_text_index[batch_idx][idx])]
                except:
                    continue
                if char_idx == '</s>':  # end
                    break
                char_list.append(char_idx)

                if char_text_index is not None:
                    char_box_list.append(char_box_idx)

                if text_prob is not None:
                    conf_list.append(text_prob[batch_idx][idx])
                else:
                    conf_list.append(1)

                if char_text_prob is not None:
                    conf_box_list.append(char_text_prob[batch_idx][idx])
                else:
                    conf_box_list.append(1)
            text = ''.join(char_list)
            result_list.append((text, np.mean(conf_list).tolist()))

            if char_text_index is not None:
                text_box = ''.join(char_box_list)
                box_result_list.append(
                    (text_box, np.mean(conf_box_list).tolist()))
        if char_text_index is not None:
            return result_list, box_result_list
        return result_list