# coding: utf-8

import operator
from typing import Callable, Dict, List, Optional

import numpy as np

from supervisely.annotation.annotation import Annotation
from supervisely.annotation.label import Label
from supervisely.annotation.obj_class import ObjClass
from supervisely.geometry.bitmap import Bitmap, SkeletonizeMethod
from supervisely.geometry.point_location import PointLocation
from supervisely.geometry.polygon import Polygon
from supervisely.geometry.polyline import Polyline
from supervisely.geometry.rectangle import Rectangle
from supervisely.geometry.alpha_mask import AlphaMask

if not hasattr(np, "int"):
    np.int = np.int_
if not hasattr(np, "bool"):
    np.bool = np.bool_


def skeletonize_bitmap(
    ann: Annotation, classes: List[str], method_id: SkeletonizeMethod
) -> Annotation:
    """Extracts skeletons from bitmap figures.

    :param ann: Input annotation.
    :type ann: Annotation
    :param classes: List of classes to skeletonize.
    :type classes: List[str]
    :param method_id: Algorithm of processing. See supervisely.geometry.bitmap.SkeletonizeMethod enum.
    :type method_id: SkeletonizeMethod
    :return: Annotation with skeletonized labels.
    :rtype: Annotation
    :raises RuntimeError: If input class is not a Bitmap.
    """

    def _skel(label: Label):
        if label.obj_class.name not in classes:
            return [label]

        if not isinstance(label.geometry, Bitmap):
            raise RuntimeError("Input class must be a Bitmap.")

        return [label.clone(geometry=label.geometry.skeletonize(method_id))]

    return ann.transform_labels(_skel)


def approximate_vector(ann: Annotation, classes: List[str], epsilon: float) -> Annotation:
    """Approximates vector figures: lines and polygons.

    :param ann: Input annotations.
    :type ann: Annotation
    :param classes: List of classes to apply transformation.
    :type classes: List[str]
    :param epsilon: Approximation accuracy (maximum distance between the original curve and its approximation)
    :type epsilon: float
    :return: Annotation with approximated vector figures of selected classes.
    :rtype: Annotation
    :raises RuntimeError: If input class is not a Polygon or a Line.
    """

    def _approx(label: Label):
        if label.obj_class.name not in classes:
            return [label]

        if not isinstance(label.geometry, (Polygon, Polyline)):
            raise RuntimeError("Input class must be a Polygon or a Line.")

        return [label.clone(geometry=label.geometry.approx_dp(epsilon))]

    return ann.transform_labels(_approx)


def add_background(ann: Annotation, bg_class: ObjClass) -> Annotation:
    """Adds background rectangle (size equals to image size) to annotation.

    :param ann: Input annotation.
    :type ann: Annotation
    :param bg_class: ObjClass instance for background class label.
    :type bg_class: ObjClass
    :return: Annotation with added background rectangle.
    :rtype: Annotation
    """
    img_size = ann.img_size
    rect = Rectangle(0, 0, img_size[0] - 1, img_size[1] - 1)
    new_label = Label(rect, bg_class)
    return ann.add_label(new_label)


def drop_object_by_class(ann: Annotation, classes: List[str]) -> Annotation:
    """Removes labels of specified classes from annotation.

    :param ann: Input annotation.
    :type ann: Annotation
    :param classes: List of classes to remove.
    :type classes: List[str]
    :return: Annotation with removed labels of specified classes.
    :rtype: Annotation
    """

    def _filter(label: Label):
        if label.obj_class.name in classes:
            return [label]
        return []

    return ann.transform_labels(_filter)


def filter_objects_by_area(
    ann: Annotation,
    classes: List[str],
    comparator: Optional[Callable] = operator.lt,
    thresh_percent: Optional[float] = None,
) -> Annotation:  # @ TODO: add size mode
    """Deletes labels less (or greater) than specified percentage of image area.

    :param ann: Input annotation.
    :type ann: Annotation
    :param classes: List of classes to filter.
    :type classes: List[str]
    :param comparator: Comparison function.
    :type comparator: Callable, optional
    :param thresh_percent: Threshold percent value of image area.
    :type thresh_percent: float, optional
    :return: Annotation containing filtered labels.
    :rtype: Annotation
    """
    imsize = ann.img_size
    img_area = float(imsize[0] * imsize[1])

    def _del_filter_percent(label: Label):
        if label.obj_class.name in classes:
            fig_area = label.area
            area_percent = 100.0 * fig_area / img_area
            if comparator(area_percent, thresh_percent):  # satisfied condition
                return []  # action 'delete'
        return [label]

    return ann.transform_labels(imsize, _del_filter_percent)


def bitwise_mask(
    ann: Annotation,
    class_mask: str,
    classes_to_correct: List[str],
    bitwise_op: Optional[Callable[[np.ndarray, np.ndarray], np.ndarray]] = np.logical_and,
) -> Annotation:
    """Performs bitwise operation between two masks. Uses one target mask to correct all others.

    :param ann: Input annotation.
    :type ann: Annotation
    :param class_mask: Class name of target mask.
    :type class_mask: str
    :param classes_to_correct: List of classes which will be corrected using target mask.
    :type classes_to_correct: List[str]
    ;param bitwise_op: Bitwise numpy function to process masks.For example: "np.logical_or", "np.logical_and",
    :type bitwise_op: Callable[[np.ndarray, np.ndarray], np.ndarray], optional
    :return: Annotation containing corrected Bitmaps.
    :rtype: Annotation
    :raises RuntimeError: If input class is not a Bitmap.
    """
    imsize = ann.img_size

    def find_mask_class(labels, class_mask_name):
        for label in labels:
            if label.obj_class.name == class_mask_name:
                if not isinstance(label.geometry, Bitmap):
                    raise RuntimeError("Class <{}> must be a Bitmap.".format(class_mask_name))
                return label

    mask_label = find_mask_class(ann.labels, class_mask)
    if mask_label is not None:
        target_original, target_mask = mask_label.geometry.origin, mask_label.geometry.data
        full_target_mask = np.full(imsize, False, bool)

        full_target_mask[
            target_original.row : target_original.row + target_mask.shape[0],
            target_original.col : target_original.col + target_mask.shape[1],
        ] = target_mask

        def perform_op(label):
            if label.obj_class.name not in classes_to_correct or label.obj_class.name == class_mask:
                return [label]

            if not isinstance(label.geometry, Bitmap):
                raise RuntimeError("Input class must be a Bitmap.")

            new_geom = label.geometry.bitwise_mask(full_target_mask, bitwise_op)
            return [label.clone(geometry=new_geom)] if new_geom is not None else []

        res_ann = ann.transform_labels(perform_op)
    else:
        res_ann = ann.clone()

    return res_ann


def find_contours(
    ann: Annotation, classes_mapping: Dict[str, ObjClass]
) -> Annotation:  # @TODO: approximation dropped
    """Finds contours of Bitmaps and converts them to Polygons.

    :param ann: Input annotation.
    :type ann: Annotation
    :param classes_mapping: Dict matching source class names and new ObjClasses
    :type classes_mapping: Dict[str, ObjClass]
    :return: Annotation with Bitmaps converted to contours Polygons.
    :rtype: Annotation
    :raises RuntimeError: If input class is not a Bitmap.
    """

    def to_contours(label: Label):
        new_obj_cls = classes_mapping.get(label.obj_class.name)
        if new_obj_cls is None:
            return [label]
        if not isinstance(label.geometry, Bitmap) and not isinstance(label.geometry, AlphaMask):
            raise RuntimeError("Input class must be a Bitmap or an AlphaMask.")

        return [
            Label(geometry=geom, obj_class=new_obj_cls) for geom in label.geometry.to_contours()
        ]

    return ann.transform_labels(to_contours)


def extract_labels_from_mask(
    mask: np.ndarray, color_id_to_obj_class: Dict[int, ObjClass]
) -> List[Label]:
    """
    Extract multiclass instances from grayscale mask and save it to labels list.
    :param mask: multiclass grayscale mask
    :type mask: np.ndarray
    :param color_id_to_obj_class: dict of objects classes assigned to color id (e.g. {1: ObjClass('cat), ...})
    :type color_id_to_obj_class: Dict[int, ObjClass]
    :return: list of labels with bitmap geometry
    :rtype: List[Label]
    """
    from scipy import ndimage
    from skimage import measure

    zero_offset = 1 if 0 in color_id_to_obj_class else 0
    if zero_offset > 0:
        mask = mask + zero_offset

    labeled, labels_count = measure.label(mask, connectivity=1, return_num=True)
    objects_slices = ndimage.find_objects(labeled)
    labels = []

    for object_index, slices in enumerate(objects_slices, start=1):
        crop = mask[slices]
        sub_mask = crop * (labeled[slices] == object_index).astype(np.int)

        class_index = np.max(sub_mask) - zero_offset

        if class_index in color_id_to_obj_class:
            bitmap = Bitmap(
                data=sub_mask.astype(np.bool),
                origin=PointLocation(slices[0].start, slices[1].start),
            )
            label = Label(geometry=bitmap, obj_class=color_id_to_obj_class.get(class_index))
            labels.append(label)
    return labels
