# Source code for draugr.opencv_utilities.transformation.cv2_transforms

    import types
    from typing import Any, List, Tuple

    import cv2
    import numpy
    from PIL.ImageTransform import Transform
    from numpy import random

    __all__ = [
        "CV2Compose",
        "Lambda",
        "ConvertFromInts",
        "SubtractMeans",
        "CV2ToAbsoluteCoords",
        "CV2ToPercentCoords",
        "CV2Resize",
        "CV2RandomSaturation",
        "CV2RandomHue",
        "CV2RandomLightingNoise",
        "CV2ConvertColor",
        "CV2RandomContrast",
        "CV2RandomBrightness",
        "CV2RandomSampleCrop",
        "CV2Expand",
        "CV2RandomMirror",
        "CV2SwapChannels",
        "CV2PhotometricDistort",
    ]

    from draugr.opencv_utilities.bounding_boxes.evaluation import (
        jaccard_overlap_numpy,
        remove_null_boxes,
    )
    from warg import TripleNumber

    class CV2Compose(object):
        """Composes several augmentations together.
        Args:
        transforms (List[Transform]): list of transforms to compose.
        Example:
        >>> augmentations.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])"""

        def __init__(self, transforms: List[Transform]):
            self._transforms = transforms

        def __call__(self, img, boxes=None, labels=None) -> Tuple:
            for t in self._transforms:
                img, boxes, labels = t(img, boxes, labels)
                if boxes is not None:
                    boxes, labels = remove_null_boxes(boxes, labels)
            return img, boxes, labels

    class Lambda(object):
        """Applies a lambda as a transform."""

        def __init__(self, lambd: callable):
            assert isinstance(lambd, types.LambdaType)
            self.lambd = lambd

        def __call__(self, img: Any, boxes: Any = None, labels: Any = None):
            return self.lambd(img, boxes, labels)

    class ConvertFromInts(object):
        def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
            return image.astype(numpy.float32), boxes, labels

    class SubtractMeans(object):
        """description"""

        def __init__(self, mean: Tuple):
            self.mean = numpy.array(mean, dtype=numpy.float32)

        def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
            image = image.astype(numpy.float32)
            image -= self.mean
            return image.astype(numpy.float32), boxes, labels

    class CV2ToAbsoluteCoords(object):
        def __call__(self, image: Any, boxes: Any = None, labels: Any = None) -> Tuple:
            height, width, channels = image.shape
            boxes[:, 0] *= width
            boxes[:, 2] *= width
            boxes[:, 1] *= height
            boxes[:, 3] *= height

            return image, boxes, labels

    class CV2ToPercentCoords(object):
        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            height, width, channels = image.shape
            boxes[:, 0] /= width
            boxes[:, 2] /= width
            boxes[:, 1] /= height
            boxes[:, 3] /= height

            return image, boxes, labels

    class CV2Resize(object):
        """description"""

        def __init__(self, size: int = 300):
            self._size = size

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            image = cv2.resize(image, (self._size, self._size))
            return image, boxes, labels

    class CV2RandomSaturation(object):
        """description"""

        def __init__(self, lower: float = 0.5, upper: float = 1.5):
            self.lower = lower
            self.upper = upper
            assert self.upper >= self.lower, "contrast upper must be >= lower."
            assert self.lower >= 0, "contrast lower must be non-negative."

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if random.randint(2):
                image[..., 1] *= random.uniform(self.lower, self.upper)

            return image, boxes, labels

    class CV2RandomHue(object):
        """description"""

        def __init__(self, delta: float = 18.0):
            assert 0.0 <= delta <= 360.0
            self.delta = delta

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if random.randint(2):
                image[..., 0] += random.uniform(-self.delta, self.delta)
                image[..., 0][image[..., 0] > 360.0] -= 360.0
                image[..., 0][image[..., 0] < 0.0] += 360.0
            return image, boxes, labels

    class CV2RandomLightingNoise(object):
        """description"""

        def __init__(self):
            self.perms = ((0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0))

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if random.randint(2):
                swap = self.perms[random.randint(len(self.perms))]
                shuffle = CV2SwapChannels(swap)  # shuffle channels
                image = shuffle(image)
            return image, boxes, labels

    class CV2ConvertColor(object):
        """description"""

        def __init__(self, current, transform):
            self.transform = transform
            self.current = current

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if self.current == "BGR" and self.transform == "HSV":
                image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
            elif self.current == "RGB" and self.transform == "HSV":
                image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
            elif self.current == "BGR" and self.transform == "RGB":
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            elif self.current == "HSV" and self.transform == "BGR":
                image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
            elif self.current == "HSV" and self.transform == "RGB":
                image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
            else:
                raise NotImplementedError
            return image, boxes, labels

    class CV2RandomContrast(object):
        """description"""

        def __init__(self, lower: float = 0.5, upper: float = 1.5):
            self.lower = lower
            self.upper = upper
            assert self.upper >= self.lower, "contrast upper must be >= lower."
            assert self.lower >= 0, "contrast lower must be non-negative."

        # expects float image
        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if random.randint(2):
                alpha = random.uniform(self.lower, self.upper)
                image *= alpha
            return image, boxes, labels

    class CV2RandomBrightness(object):
        """description"""

        def __init__(self, delta: float = 32):
            assert delta >= 0.0
            assert delta <= 255.0
            self.delta = delta

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            if random.randint(2):
                delta = random.uniform(-self.delta, self.delta)
                image += delta
            return image, boxes, labels

    class CV2RandomSampleCrop(object):
        """Crop
        Arguments:
        img (Image): the image being input during training
        boxes (Tensor): the original bounding boxes in pt form
        labels (Tensor): the class labels for each bbox
        mode (float tuple): the min and max jaccard overlaps
        Return:
        (img, boxes, classes)
        img (Image): the cropped image
        boxes (Tensor): the adjusted bounding boxes in pt form
        labels (Tensor): the class labels for each bbox"""

        def __init__(self):
            self.sample_options = (
                # using entire original input image
                None,
                # sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9
                (0.1, None),
                (0.3, None),
                (0.7, None),
                (0.9, None),
                # randomly sample a patch
                (None, None),
            )

        def __call__(
            self,
            image: numpy.ndarray,
            boxes: numpy.ndarray = None,
            labels: numpy.ndarray = None,
        ) -> Tuple:
            # guard against no boxes
            if boxes is not None and boxes.shape[0] == 0:
                return image, boxes, labels
            height, width, _ = image.shape
            while True:
                # randomly choose a mode
                mode = random.choice(self.sample_options)
                if mode is None:
                    return image, boxes, labels

                min_iou, max_iou = mode
                if min_iou is None:
                    min_iou = float("-inf")
                if max_iou is None:
                    max_iou = float("inf")

                # max trails (50)
                for _ in range(50):
                    current_image = image

                    w = random.uniform(0.3 * width, width)
                    h = random.uniform(0.3 * height, height)

                    # aspect ratio constraint b/t .5 & 2
                    if h / w < 0.5 or h / w > 2:
                        continue

                    left = random.uniform(width - w)
                    top = random.uniform(height - h)

                    # convert to integer rect x1,y1,x2,y2
                    rect = numpy.array([int(left), int(top), int(left + w), int(top + h)])

                    # calculate IoU (jaccard overlap) b/t the cropped and gt boxes
                    overlap = jaccard_overlap_numpy(boxes, rect)

                    # is min and max overlap constraint satisfied? if not try again
                    if overlap.max() < min_iou or overlap.min() > max_iou:
                        continue

                    # cut the crop from the image
                    current_image = current_image[rect[1] : rect[3], rect[0] : rect[2], :]

                    # keep overlap with gt box IF center in sampled patch
                    centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0

                    # mask in all gt boxes that above and to the left of centers
                    m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1])

                    # mask in all gt boxes that under and to the right of centers
                    m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1])

                    # mask in that both m1 and m2 are true
                    mask = m1 * m2

                    # have any valid boxes? try again if not
                    if not mask.any():
                        continue

                    # take only matching gt boxes
                    current_boxes = boxes[mask, :].copy()

                    # take only matching gt labels
                    current_labels = labels[mask]

                    # should we use the box left and top corner or the crop's
                    current_boxes[:, :2] = numpy.maximum(current_boxes[:, :2], rect[:2])
                    # adjust to crop (by substracting crop's left,top)
                    current_boxes[:, :2] -= rect[:2]

                    current_boxes[:, 2:] = numpy.minimum(current_boxes[:, 2:], rect[2:])
                    # adjust to crop (by substracting crop's left,top)
                    current_boxes[:, 2:] -= rect[:2]

                    return current_image, current_boxes, current_labels

    class CV2Expand(object):
        """description"""

        def __init__(self, mean: float):
            self.mean = mean

        def __call__(
            self, image: numpy.ndarray, boxes: numpy.ndarray, labels: numpy.ndarray
        ) -> Tuple:
            if random.randint(2):
                return image, boxes, labels

            height, width, depth = image.shape
            ratio = random.uniform(1, 4)
            left = random.uniform(0, width * ratio - width)
            top = random.uniform(0, height * ratio - height)

            expand_image = numpy.zeros(
                (int(height * ratio), int(width * ratio), depth), dtype=image.dtype
            )
            expand_image[..., :] = self.mean
            expand_image[
                int(top) : int(top + height), int(left) : int(left + width)
            ] = image
            image = expand_image

            boxes = boxes.copy()
            boxes[:, :2] += (int(left), int(top))
            boxes[:, 2:] += (int(left), int(top))

            return image, boxes, labels

    class CV2RandomMirror(object):
        def __call__(
            self, image: numpy.ndarray, boxes: numpy.ndarray, classes: numpy.ndarray
        ) -> Tuple:
            _, width, _ = image.shape
            if random.randint(2):
                image = image[:, ::-1]
                boxes = boxes.copy()
                boxes[:, 0::2] = width - boxes[:, 2::-2]
            return image, boxes, classes

    class CV2SwapChannels(object):
        """Transforms a tensorized image by swapping the channels in the order
        specified in the swap tuple.
        Args:
        swaps (int triple): final order of channels
        eg: (2, 1, 0)"""

        def __init__(self, swaps: TripleNumber):
            self.swaps = swaps

        def __call__(self, image: numpy.ndarray) -> numpy.ndarray:
            """
            Args:
            image (Tensor): image tensor to be transformed
            Return:
            a tensor with channels swapped according to swap"""
            # if torch.is_tensor(image):
            #     image = image.data.cpu().numpy()
            # else:
            #     image = numpy.array(image)
            image = image[..., self.swaps]
            return image

    class CV2PhotometricDistort(object):
        """description"""

        def __init__(self):
            self.pd = [
                CV2RandomContrast(),  # RGB
                CV2ConvertColor(current="RGB", transform="HSV"),  # HSV
                CV2RandomSaturation(),  # HSV
                CV2RandomHue(),  # HSV
                CV2ConvertColor(current="HSV", transform="RGB"),  # RGB
                CV2RandomContrast(),  # RGB
            ]
            self.rand_brightness = CV2RandomBrightness()
            self.rand_light_noise = CV2RandomLightingNoise()

        def __call__(
            self, image: numpy.ndarray, boxes: numpy.ndarray, labels: numpy.ndarray
        ) -> Tuple:
            im = image.copy()
            im, boxes, labels = self.rand_brightness(im, boxes, labels)
            if random.randint(2):
                distort = CV2Compose(self.pd[:-1])
            else:
                distort = CV2Compose(self.pd[1:])
            im, boxes, labels = distort(im, boxes, labels)
            return self.rand_light_noise(im, boxes, labels)