from PIL import Image
import cv2
import numpy as np
import math
import time
import torch
from tqdm import tqdm

from ...pytorchocr.base_ocr_v20 import BaseOCRV20
from . import pytorchocr_utility as utility
from ...pytorchocr.postprocess import build_post_process
from ...pytorchocr.modeling.backbones.rec_hgnet import ConvBNAct


class TextRecognizer(BaseOCRV20):
    def __init__(self, args, **kwargs):
        self.device = args.device
        self.rec_image_shape = [int(v) for v in args.rec_image_shape.split(",")]
        self.character_type = args.rec_char_type
        self.rec_batch_num = args.rec_batch_num
        self.rec_algorithm = args.rec_algorithm
        self.max_text_length = args.max_text_length
        postprocess_params = {
            'name': 'CTCLabelDecode',
            "character_type": args.rec_char_type,
            "character_dict_path": args.rec_char_dict_path,
            "use_space_char": args.use_space_char
        }
        if self.rec_algorithm == "SRN":
            postprocess_params = {
                'name': 'SRNLabelDecode',
                "character_type": args.rec_char_type,
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == "RARE":
            postprocess_params = {
                'name': 'AttnLabelDecode',
                "character_type": args.rec_char_type,
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == 'NRTR':
            postprocess_params = {
                'name': 'NRTRLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == "SAR":
            postprocess_params = {
                'name': 'SARLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == 'ViTSTR':
            postprocess_params = {
                'name': 'ViTSTRLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == "CAN":
            self.inverse = args.rec_image_inverse
            postprocess_params = {
                'name': 'CANLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == 'RFL':
            postprocess_params = {
                'name': 'RFLLabelDecode',
                "character_dict_path": None,
                "use_space_char": args.use_space_char
            }
        self.postprocess_op = build_post_process(postprocess_params)

        self.limited_max_width = args.limited_max_width
        self.limited_min_width = args.limited_min_width

        self.weights_path = args.rec_model_path
        self.yaml_path = args.rec_yaml_path

        network_config = utility.get_arch_config(self.weights_path)
        weights = self.read_pytorch_weights(self.weights_path)

        self.out_channels = self.get_out_channels(weights)
        if self.rec_algorithm == 'NRTR':
            self.out_channels = list(weights.values())[-1].numpy().shape[0]
        elif self.rec_algorithm == 'SAR':
            self.out_channels = list(weights.values())[-3].numpy().shape[0]

        kwargs['out_channels'] = self.out_channels
        super(TextRecognizer, self).__init__(network_config, **kwargs)

        self.load_state_dict(weights)
        self.net.eval()
        self.net.to(self.device)
        for module in self.net.modules():
            if isinstance(module, ConvBNAct):
                if module.use_act:
                    torch.quantization.fuse_modules(module, ['conv', 'bn', 'act'], inplace=True)
                else:
                    torch.quantization.fuse_modules(module, ['conv', 'bn'], inplace=True)

    def resize_norm_img(self, img, max_wh_ratio):
        imgC, imgH, imgW = self.rec_image_shape
        if self.rec_algorithm == 'NRTR' or self.rec_algorithm == 'ViTSTR':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # return padding_im
            image_pil = Image.fromarray(np.uint8(img))
            if self.rec_algorithm == 'ViTSTR':
                img = image_pil.resize([imgW, imgH], Image.BICUBIC)
            else:
                img = image_pil.resize([imgW, imgH], Image.ANTIALIAS)
            img = np.array(img)
            norm_img = np.expand_dims(img, -1)
            norm_img = norm_img.transpose((2, 0, 1))
            if self.rec_algorithm == 'ViTSTR':
                norm_img = norm_img.astype(np.float32) / 255.
            else:
                norm_img = norm_img.astype(np.float32) / 128. - 1.
            return norm_img
        elif self.rec_algorithm == 'RFL':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            resized_image = cv2.resize(
                img, (imgW, imgH), interpolation=cv2.INTER_CUBIC)
            resized_image = resized_image.astype('float32')
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
            resized_image -= 0.5
            resized_image /= 0.5
            return resized_image

        assert imgC == img.shape[2]
        max_wh_ratio = max(max_wh_ratio, imgW / imgH)
        imgW = int(imgH * max_wh_ratio)
        imgW = max(min(imgW, self.limited_max_width), self.limited_min_width)
        h, w = img.shape[:2]
        ratio = w / float(h)
        ratio_imgH = max(math.ceil(imgH * ratio), self.limited_min_width)
        resized_w = min(imgW, int(ratio_imgH))
        resized_image = cv2.resize(img, (resized_w, imgH)) /127.5 - 1
        padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
        padding_im[:, :, 0:resized_w] = resized_image.transpose((2, 0, 1))
        return padding_im

    def resize_norm_img_svtr(self, img, image_shape):

        imgC, imgH, imgW = image_shape
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        return resized_image


    def resize_norm_img_srn(self, img, image_shape):
        imgC, imgH, imgW = image_shape

        img_black = np.zeros((imgH, imgW))
        im_hei = img.shape[0]
        im_wid = img.shape[1]

        if im_wid <= im_hei * 1:
            img_new = cv2.resize(img, (imgH * 1, imgH))
        elif im_wid <= im_hei * 2:
            img_new = cv2.resize(img, (imgH * 2, imgH))
        elif im_wid <= im_hei * 3:
            img_new = cv2.resize(img, (imgH * 3, imgH))
        else:
            img_new = cv2.resize(img, (imgW, imgH))

        img_np = np.asarray(img_new)
        img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
        img_black[:, 0:img_np.shape[1]] = img_np
        img_black = img_black[:, :, np.newaxis]

        row, col, c = img_black.shape
        c = 1

        return np.reshape(img_black, (c, row, col)).astype(np.float32)

    def srn_other_inputs(self, image_shape, num_heads, max_text_length):

        imgC, imgH, imgW = image_shape
        feature_dim = int((imgH / 8) * (imgW / 8))

        encoder_word_pos = np.array(range(0, feature_dim)).reshape(
            (feature_dim, 1)).astype('int64')
        gsrm_word_pos = np.array(range(0, max_text_length)).reshape(
            (max_text_length, 1)).astype('int64')

        gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
        gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias1 = np.tile(
            gsrm_slf_attn_bias1,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias2 = np.tile(
            gsrm_slf_attn_bias2,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        encoder_word_pos = encoder_word_pos[np.newaxis, :]
        gsrm_word_pos = gsrm_word_pos[np.newaxis, :]

        return [
            encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
            gsrm_slf_attn_bias2
        ]

    def process_image_srn(self, img, image_shape, num_heads, max_text_length):
        norm_img = self.resize_norm_img_srn(img, image_shape)
        norm_img = norm_img[np.newaxis, :]

        [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = \
            self.srn_other_inputs(image_shape, num_heads, max_text_length)

        gsrm_slf_attn_bias1 = gsrm_slf_attn_bias1.astype(np.float32)
        gsrm_slf_attn_bias2 = gsrm_slf_attn_bias2.astype(np.float32)
        encoder_word_pos = encoder_word_pos.astype(np.int64)
        gsrm_word_pos = gsrm_word_pos.astype(np.int64)

        return (norm_img, encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
                gsrm_slf_attn_bias2)

    def resize_norm_img_sar(self, img, image_shape,
                            width_downsample_ratio=0.25):
        imgC, imgH, imgW_min, imgW_max = image_shape
        h = img.shape[0]
        w = img.shape[1]
        valid_ratio = 1.0
        # make sure new_width is an integral multiple of width_divisor.
        width_divisor = int(1 / width_downsample_ratio)
        # resize
        ratio = w / float(h)
        resize_w = math.ceil(imgH * ratio)
        if resize_w % width_divisor != 0:
            resize_w = round(resize_w / width_divisor) * width_divisor
        if imgW_min is not None:
            resize_w = max(imgW_min, resize_w)
        if imgW_max is not None:
            valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
            resize_w = min(imgW_max, resize_w)
        resized_image = cv2.resize(img, (resize_w, imgH))
        resized_image = resized_image.astype('float32')
        # norm
        if image_shape[0] == 1:
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
        else:
            resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        resize_shape = resized_image.shape
        padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
        padding_im[:, :, 0:resize_w] = resized_image
        pad_shape = padding_im.shape

        return padding_im, resize_shape, pad_shape, valid_ratio


    def norm_img_can(self, img, image_shape):

        img = cv2.cvtColor(
            img, cv2.COLOR_BGR2GRAY)  # CAN only predict gray scale image

        if self.inverse:
            img = 255 - img

        if self.rec_image_shape[0] == 1:
            h, w = img.shape
            _, imgH, imgW = self.rec_image_shape
            if h < imgH or w < imgW:
                padding_h = max(imgH - h, 0)
                padding_w = max(imgW - w, 0)
                img_padded = np.pad(img, ((0, padding_h), (0, padding_w)),
                                    'constant',
                                    constant_values=(255))
                img = img_padded

        img = np.expand_dims(img, 0) / 255.0  # h,w,c -> c,h,w
        img = img.astype('float32')

        return img

    def __call__(self, img_list, tqdm_enable=False, tqdm_desc="OCR-rec Predict"):
        img_num = len(img_list)
        # Calculate the aspect ratio of all text bars
        width_list = []
        for img in img_list:
            width_list.append(img.shape[1] / float(img.shape[0]))
        # Sorting can speed up the recognition process
        indices = np.argsort(np.array(width_list))

        # rec_res = []
        rec_res = [['', 0.0]] * img_num
        batch_num = self.rec_batch_num
        elapse = 0
        # for beg_img_no in range(0, img_num, batch_num):
        with tqdm(total=img_num, desc=tqdm_desc, disable=not tqdm_enable) as pbar:
            index = 0
            for beg_img_no in range(0, img_num, batch_num):
                end_img_no = min(img_num, beg_img_no + batch_num)
                norm_img_batch = []
                max_wh_ratio = width_list[indices[end_img_no - 1]]
                for ino in range(beg_img_no, end_img_no):
                    if self.rec_algorithm == "SAR":
                        norm_img, _, _, valid_ratio = self.resize_norm_img_sar(
                            img_list[indices[ino]], self.rec_image_shape)
                        norm_img = norm_img[np.newaxis, :]
                        valid_ratio = np.expand_dims(valid_ratio, axis=0)
                        valid_ratios = []
                        valid_ratios.append(valid_ratio)
                        norm_img_batch.append(norm_img)

                    elif self.rec_algorithm == "SVTR":
                        norm_img = self.resize_norm_img_svtr(img_list[indices[ino]],
                                                             self.rec_image_shape)
                        norm_img = norm_img[np.newaxis, :]
                        norm_img_batch.append(norm_img)
                    elif self.rec_algorithm == "SRN":
                        norm_img = self.process_image_srn(img_list[indices[ino]],
                                                          self.rec_image_shape, 8,
                                                          self.max_text_length)
                        encoder_word_pos_list = []
                        gsrm_word_pos_list = []
                        gsrm_slf_attn_bias1_list = []
                        gsrm_slf_attn_bias2_list = []
                        encoder_word_pos_list.append(norm_img[1])
                        gsrm_word_pos_list.append(norm_img[2])
                        gsrm_slf_attn_bias1_list.append(norm_img[3])
                        gsrm_slf_attn_bias2_list.append(norm_img[4])
                        norm_img_batch.append(norm_img[0])
                    elif self.rec_algorithm == "CAN":
                        norm_img = self.norm_img_can(img_list[indices[ino]],
                                                     max_wh_ratio)
                        norm_img = norm_img[np.newaxis, :]
                        norm_img_batch.append(norm_img)
                        norm_image_mask = np.ones(norm_img.shape, dtype='float32')
                        word_label = np.ones([1, 36], dtype='int64')
                        norm_img_mask_batch = []
                        word_label_list = []
                        norm_img_mask_batch.append(norm_image_mask)
                        word_label_list.append(word_label)
                    else:
                        norm_img = self.resize_norm_img(img_list[indices[ino]],
                                                        max_wh_ratio)
                        norm_img = norm_img[np.newaxis, :]
                        norm_img_batch.append(norm_img)
                norm_img_batch = np.concatenate(norm_img_batch)
                norm_img_batch = norm_img_batch.copy()

                if self.rec_algorithm == "SRN":
                    starttime = time.time()
                    encoder_word_pos_list = np.concatenate(encoder_word_pos_list)
                    gsrm_word_pos_list = np.concatenate(gsrm_word_pos_list)
                    gsrm_slf_attn_bias1_list = np.concatenate(
                        gsrm_slf_attn_bias1_list)
                    gsrm_slf_attn_bias2_list = np.concatenate(
                        gsrm_slf_attn_bias2_list)

                    with torch.no_grad():
                        inp = torch.from_numpy(norm_img_batch)
                        encoder_word_pos_inp = torch.from_numpy(encoder_word_pos_list)
                        gsrm_word_pos_inp = torch.from_numpy(gsrm_word_pos_list)
                        gsrm_slf_attn_bias1_inp = torch.from_numpy(gsrm_slf_attn_bias1_list)
                        gsrm_slf_attn_bias2_inp = torch.from_numpy(gsrm_slf_attn_bias2_list)

                        inp = inp.to(self.device)
                        encoder_word_pos_inp = encoder_word_pos_inp.to(self.device)
                        gsrm_word_pos_inp = gsrm_word_pos_inp.to(self.device)
                        gsrm_slf_attn_bias1_inp = gsrm_slf_attn_bias1_inp.to(self.device)
                        gsrm_slf_attn_bias2_inp = gsrm_slf_attn_bias2_inp.to(self.device)

                        backbone_out = self.net.backbone(inp) # backbone_feat
                        prob_out = self.net.head(backbone_out, [encoder_word_pos_inp, gsrm_word_pos_inp, gsrm_slf_attn_bias1_inp, gsrm_slf_attn_bias2_inp])
                    # preds = {"predict": prob_out[2]}
                    preds = {"predict": prob_out["predict"]}

                elif self.rec_algorithm == "SAR":
                    starttime = time.time()
                    # valid_ratios = np.concatenate(valid_ratios)
                    # inputs = [
                    #     norm_img_batch,
                    #     valid_ratios,
                    # ]

                    with torch.no_grad():
                        inp = torch.from_numpy(norm_img_batch)
                        inp = inp.to(self.device)
                        preds = self.net(inp)

                elif self.rec_algorithm == "CAN":
                    starttime = time.time()
                    norm_img_mask_batch = np.concatenate(norm_img_mask_batch)
                    word_label_list = np.concatenate(word_label_list)
                    inputs = [norm_img_batch, norm_img_mask_batch, word_label_list]

                    inp = [torch.from_numpy(e_i) for e_i in inputs]
                    inp = [e_i.to(self.device) for e_i in inp]
                    with torch.no_grad():
                        outputs = self.net(inp)
                        outputs = [v.cpu().numpy() for k, v in enumerate(outputs)]

                    preds = outputs

                else:
                    starttime = time.time()

                    with torch.no_grad():
                        inp = torch.from_numpy(norm_img_batch)
                        inp = inp.to(self.device)
                        preds = self.net(inp)

                with torch.no_grad():
                    rec_result = self.postprocess_op(preds)

                for rno in range(len(rec_result)):
                    rec_res[indices[beg_img_no + rno]] = rec_result[rno]
                elapse += time.time() - starttime

                # 更新进度条，每次增加batch_size，但要注意最后一个batch可能不足batch_size
                current_batch_size = min(batch_num, img_num - index * batch_num)
                index += 1
                pbar.update(current_batch_size)

        # Fix NaN values in recognition results
        for i in range(len(rec_res)):
            text, score = rec_res[i]
            if isinstance(score, float) and math.isnan(score):
                rec_res[i] = (text, 0.0)

        return rec_res, elapse