import cv2
import numpy as np
import torch
from mmdet.registry import VISUALIZERS


class SegMaskHelper:
    def __init__(self):
        pass

    # Pad the masks to image size (bug in RTMDet config?)
    # @timer_func
    def align_masks_with_image(self, result, img):
        masks = list()

        img = img[..., ::-1].copy()

        for j, mask in enumerate(result.pred_instances.masks):
            numpy_mask = mask.cpu().numpy()
            mask = cv2.resize(
                numpy_mask.astype(np.uint8),
                (img.shape[1], img.shape[0]),
                interpolation=cv2.INTER_NEAREST,
            )

            # Pad the mask to match the size of the image
            padded_mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
            padded_mask[: mask.shape[0], : mask.shape[1]] = mask
            mask = padded_mask
            mask = torch.from_numpy(mask)
            masks.append(mask)

        stacked_masks = torch.stack(masks)
        result.pred_instances.masks = stacked_masks

        return result

    # Crops the images using masks and put the cropped images on a white background
    # @timer_func
    def crop_masks(self, result, img):
        cropped_imgs = list()
        polygons = list()

        for j, mask in enumerate(result.pred_instances.masks):
            np_array = mask.cpu().numpy()
            contours, _ = cv2.findContours(
                np_array.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
            )  # fix so only one contour (the largest one) is extracted
            largest_contour = max(contours, key=cv2.contourArea)

            epsilon = 0.003 * cv2.arcLength(largest_contour, True)
            approx_poly = cv2.approxPolyDP(largest_contour, epsilon, True)
            approx_poly = np.squeeze(approx_poly)
            approx_poly = approx_poly.tolist()
            polygons.append(approx_poly)

            x, y, w, h = cv2.boundingRect(largest_contour)

            # Crop masked region and put on white background
            masked_region = img[y : y + h, x : x + w]
            white_background = np.ones_like(masked_region)
            white_background.fill(255)
            masked_region_on_white = cv2.bitwise_and(
                white_background, masked_region, mask=np_array.astype(np.uint8)[y : y + h, x : x + w]
            )

            cv2.bitwise_not(white_background, white_background, mask=np_array.astype(np.uint8)[y : y + h, x : x + w])
            res = white_background + masked_region_on_white

            cropped_imgs.append(res)

        return cropped_imgs, polygons

    def visualize_result(self, result, img, model_visualizer):
        visualizer = VISUALIZERS.build(model_visualizer)
        visualizer.add_datasample("result", img, data_sample=result, draw_gt=False)

        return visualizer.get_image()

    def _translate_line_coords(self, region_mask, line_polygons):
        region_mask = region_mask.cpu().numpy()
        region_masks_binary = (region_mask > 0).astype(np.uint8)

        box = cv2.boundingRect(region_masks_binary)
        translated_line_polygons = [[[a + box[0], b + box[1]] for [a, b] in poly] for poly in line_polygons]

        return translated_line_polygons