import cv2
import som

import numpy as np
import supervision as sv


class Visualizer:

    def __init__(
        self,
        line_thickness: int = 2,
        mask_opacity: float = 0.1,
        text_scale: float = 0.6
    ) -> None:
        self.box_annotator = sv.BoundingBoxAnnotator(
            color_lookup=sv.ColorLookup.INDEX,
            thickness=line_thickness)
        self.mask_annotator = sv.MaskAnnotator(
            color_lookup=sv.ColorLookup.INDEX,
            opacity=mask_opacity)
        self.polygon_annotator = sv.PolygonAnnotator(
            color_lookup=sv.ColorLookup.INDEX,
            thickness=line_thickness)
        self.label_annotator = sv.LabelAnnotator(
            color=sv.Color.black(),
            text_color=sv.Color.white(),
            color_lookup=sv.ColorLookup.INDEX,
            text_position=sv.Position.CENTER_OF_MASS,
            text_scale=text_scale)

    def visualize(
        self,
        image: np.ndarray,
        detections: sv.Detections,
        with_box: bool,
        with_mask: bool,
        with_polygon: bool,
        with_label: bool
    ) -> np.ndarray:
        annotated_image = image.copy()
        if with_box:
            annotated_image = self.box_annotator.annotate(
                scene=annotated_image, detections=detections)
        if with_mask:
            annotated_image = self.mask_annotator.annotate(
                scene=annotated_image, detections=detections)
        if with_polygon:
            annotated_image = self.polygon_annotator.annotate(
                scene=annotated_image, detections=detections)
        if with_label:
            labels = list(map(str, range(len(detections))))
            annotated_image = self.label_annotator.annotate(
                scene=annotated_image, detections=detections, labels=labels)
        return annotated_image


def refine_mask(
    mask: np.ndarray,
    area_threshold: float,
    mode: str = 'islands'
) -> np.ndarray:
    """
    Refines a mask by removing small islands or filling small holes based on area
    threshold.

    Parameters:
        mask (np.ndarray): Input binary mask.
        area_threshold (float): Threshold for relative area to remove or fill features.
        mode (str): Operation mode ('islands' for removing islands, 'holes' for filling
                    holes).

    Returns:
        np.ndarray: Refined binary mask.
    """
    mask = np.uint8(mask * 255)
    operation = cv2.RETR_EXTERNAL if mode == 'islands' else cv2.RETR_CCOMP
    contours, _ = cv2.findContours(
        mask, operation, cv2.CHAIN_APPROX_SIMPLE
    )
    total_area = cv2.countNonZero(mask) if mode == 'islands' else mask.size

    for contour in contours:
        area = cv2.contourArea(contour)
        relative_area = area / total_area
        if relative_area < area_threshold:
            cv2.drawContours(
                image=mask,
                contours=[contour],
                contourIdx=-1,
                color=(0 if mode == 'islands' else 255),
                thickness=-1
            )

    return np.where(mask > 0, 1, 0).astype(bool)


def filter_masks_by_relative_area(
    masks: np.ndarray,
    min_relative_area: float = 0.02,
    max_relative_area: float = 1.0
) -> np.ndarray:
    """
    Filters out masks based on their relative area.

    Parameters:
        masks (np.ndarray): A 3D numpy array where each slice along the third dimension
            represents a mask.
        min_relative_area (float): Minimum relative area threshold for keeping a mask.
        max_relative_area (float): Maximum relative area threshold for keeping a mask.

    Returns:
        np.ndarray: A 3D numpy array of filtered masks.
    """
    mask_areas = masks.sum(axis=(1, 2))
    total_area = masks.shape[1] * masks.shape[2]
    relative_areas = mask_areas / total_area
    min_area_filter = relative_areas >= min_relative_area
    max_area_filter = relative_areas <= max_relative_area
    return masks[min_area_filter & max_area_filter]


def postprocess_masks(
    detections: sv.Detections,
    area_threshold: float = 0.01,
    min_relative_area: float = 0.01,
    max_relative_area: float = 1.0,
    iou_threshold: float = 0.9
) -> sv.Detections:
    """
    Post-processes the masks of detection objects by removing small islands and filling
    small holes.

    Parameters:
        detections (sv.Detections): Detection objects to be filtered.
        area_threshold (float): Threshold for relative area to remove or fill features.
        min_relative_area (float): Minimum relative area threshold for detections.
        max_relative_area (float): Maximum relative area threshold for detections.
        iou_threshold (float): The IoU threshold above which masks will be considered as
            overlapping.

    Returns:
        np.ndarray: Post-processed masks.
    """
    masks = detections.mask.copy()
    for i in range(len(masks)):
        masks[i] = refine_mask(
            mask=masks[i],
            area_threshold=area_threshold,
            mode='islands'
        )
        masks[i] = refine_mask(
            mask=masks[i],
            area_threshold=area_threshold,
            mode='holes'
        )
    masks = filter_masks_by_relative_area(
        masks=masks,
        min_relative_area=min_relative_area,
        max_relative_area=max_relative_area)
    masks = som.mask_non_max_suppression(
        masks=masks,
        iou_threshold=iou_threshold)

    return sv.Detections(
        xyxy=sv.mask_to_xyxy(masks),
        mask=masks
    )