import math

import cv2
import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore


class TritonPythonModel:
    def initialize(self, args):
        self.logger = pb_utils.Logger

        self.text_threshold = 0.7
        self.link_threshold = 0.35
        self.low_text = 0.4

    def _getDetBoxes_core(self, textmap, linkmap, text_threshold, link_threshold, low_text):
        # prepare data
        linkmap = linkmap.copy()
        textmap = textmap.copy()
        img_h, img_w = textmap.shape

        """ labeling method """
        _, text_score = cv2.threshold(textmap, low_text, 1, 0)
        _, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)

        text_score_comb = np.clip(text_score + link_score, 0, 1)
        nLabels, labels, stats, _ = cv2.connectedComponentsWithStats(
            text_score_comb.astype(np.uint8),
            connectivity=4
        )

        detected_text_box_list = []
        for k in range(1, nLabels):
            # size filtering
            size = stats[k, cv2.CC_STAT_AREA]
            if size < 10: continue

            # thresholding
            if np.max(textmap[labels == k]) < text_threshold: continue

            # make segmentation map
            segmap = np.zeros(textmap.shape, dtype=np.uint8)
            segmap[labels == k] = 255
            segmap[np.logical_and(link_score == 1, text_score == 0)] = 0  # remove link area
            x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
            w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
            niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
            sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
            # boundary check
            if sx < 0: sx = 0
            if sy < 0: sy = 0
            if ex >= img_w: ex = img_w
            if ey >= img_h: ey = img_h
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
            segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)

            # make box
            np_contours = np.roll(np.array(np.where(segmap != 0)), 1, axis=0).transpose().reshape(-1, 2)
            rectangle = cv2.minAreaRect(np_contours)
            box = cv2.boxPoints(rectangle)

            # align diamond-shape
            w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
            box_ratio = max(w, h) / (min(w, h) + 1e-5)
            if abs(1 - box_ratio) <= 0.1:
                l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
                t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
                box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)

            # make clock-wise order
            startidx = box.sum(axis=1).argmin()
            box = np.roll(box, 4 - startidx, 0)
            box = np.array(box)

            margin = int(min(w, h)* 0.02)
            box[0][1] = max(0, box[0][1] - margin)
            box[1][1] = max(0, box[1][1] - margin)

            detected_text_box_list.append(box)

        return detected_text_box_list

    def _postprocess(self, td_output, resize_ratio):
        out = td_output[0]
        score_text = np.float32(out[:, :, 0])
        score_link = np.float32(out[:, :, 1])

        detected_text_box_list = self._getDetBoxes_core(
            textmap=score_text,
            linkmap=score_link,
            text_threshold=self.text_threshold,
            link_threshold=self.link_threshold,
            low_text=self.low_text
        )

        for i in range(len(detected_text_box_list)):
            detected_text_box_list[i] = detected_text_box_list[i] * 1/resize_ratio * 2

        if len(detected_text_box_list) == 0:
            output_np = np.zeros((0, 4, 2), dtype=np.float32)
        else:
            output_np = np.array(detected_text_box_list, dtype=np.float32)

        out_tensor_0 = pb_utils.Tensor("detected_text_box_list", output_np)

        return out_tensor_0

    def execute(self, requests):
        responses = []
        for request in requests:
            try:
                td_output = pb_utils.get_input_tensor_by_name(request, "td_output").as_numpy()
                resize_ratio = pb_utils.get_input_tensor_by_name(request, "resize_ratio").as_numpy()
                resize_ratio = resize_ratio[0]

                out_tensor_0 = self._postprocess(td_output, resize_ratio)

            except pb_utils.TritonModelException as e:
                responses.append(pb_utils.InferenceResponse(error=pb_utils.TritonError(str(e), pb_utils.TritonError.BAD_REQUEST)))
                continue

            responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0]))

        return responses