# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import random

import cv2
import numpy as np
from PIL import Image
from shapely import affinity
from shapely.geometry import Polygon
from torchvision.transforms import functional as F


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target

    def __repr__(self):
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
        return format_string


class Resize(object):
    def __init__(self, min_size, max_size, strict_resize):
        self.min_size = min_size
        self.max_size = max_size
        self.strict_resize = strict_resize

    # modified from torchvision to add support for max size
    def get_size(self, image_size):
        w, h = image_size
        if isinstance(self.min_size, tuple):
            if len(self.min_size) == 1:
                size = self.min_size[0]
            else:
                random_size_index = random.randint(0, len(self.min_size) - 1)
                size = self.min_size[random_size_index]
        else:
            size = self.min_size
        max_size = self.max_size
        if max_size is not None:
            min_original_size = float(min((w, h)))
            max_original_size = float(max((w, h)))
            if max_original_size / min_original_size * size > max_size:
                size = int(round(max_size * min_original_size / max_original_size))

        if (w <= h and w == size) or (h <= w and h == size):
            if self.strict_resize:
                h = h if h % 32 == 0 else (h // 32) * 32
                w = w if w % 32 == 0 else (w // 32) * 32
            return (h, w)

        if w < h:
            ow = size
            oh = int(size * h / w)
        else:
            oh = size
            ow = int(size * w / h)
        if self.strict_resize:
            oh = oh if oh % 32 == 0 else (oh // 32) * 32
            ow = ow if ow % 32 == 0 else (ow // 32) * 32

        return (oh, ow)

    def __call__(self, image, target):
        size = self.get_size(image.size)
        image = F.resize(image, size)
        if target is not None:
            target = target.resize(image.size)
        return image, target


class RandomCrop(object):
    def __init__(self, prob, crop_min_size=500, crop_max_size=1000, max_trys=50):
        self.min_size = crop_min_size
        self.max_size = crop_max_size
        self.max_trys = max_trys
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            im = np.array(image)
            w, h = image.size
            h_array = np.zeros((h), dtype=np.int32)
            w_array = np.zeros((w), dtype=np.int32)
            boxes = target.bbox.numpy()
            if len(boxes) == 0:
                return image, target
            for box in boxes:
                box = np.round(box, decimals=0).astype(np.int32)
                minx = box[0]
                maxx = box[2]
                w_array[minx:maxx] = 1
                miny = box[1]
                maxy = box[3]
                h_array[miny:maxy] = 1
            h_axis = np.where(h_array == 0)[0]
            w_axis = np.where(w_array == 0)[0]
            if len(h_axis) == 0 or len(w_axis) == 0:
                return image, target
            for _ in range(self.max_trys):
                xx = np.random.choice(w_axis, size=2)
                xmin = min(xx)
                xmax = max(xx)
                x_size = xmax - xmin
                if x_size > self.max_size or x_size < self.min_size:
                    continue
                yy = np.random.choice(h_axis, size=2)
                ymin = min(yy)
                ymax = max(yy)
                y_size = ymax - ymin
                if y_size > self.max_size or y_size < self.min_size:
                    continue
                box_in_area = (
                    (boxes[:, 0] >= xmin)
                    & (boxes[:, 1] >= ymin)
                    & (boxes[:, 2] <= xmax)
                    & (boxes[:, 3] <= ymax)
                )
                if len(np.where(box_in_area)[0]) == 0:
                    continue
                im = im[ymin:ymax, xmin:xmax]
                target = target.crop([xmin, ymin, xmax, ymax])
                return Image.fromarray(im), target
            return image, target
        else:
            return image, target


# class RandomCropFixSize(object):
#     def __init__(self, prob, crop_size=512, max_trys=50):
#         self.crop_size = crop_size
#         self.max_trys = max_trys
#         self.prob = prob

#     def __call__(self, image, target):
#         if random.random() < self.prob:
#             im = np.array(image)
#             w, h = image.size
#             h_array = np.zeros((h), dtype=np.int32)
#             w_array = np.zeros((w), dtype=np.int32)
#             boxes = target.bbox.numpy()
#             if len(boxes) == 0:
#                 return image, target
#             for box in boxes:
#                 box = np.round(box, decimals=0).astype(np.int32)
#                 minx = box[0]
#                 maxx = box[2]
#                 w_array[minx:maxx] = 1
#                 miny = box[1]
#                 maxy = box[3]
#                 h_array[miny:maxy] = 1
#             h_axis = np.where(h_array == 0)[0]
#             w_axis = np.where(w_array == 0)[0]
#             if len(h_axis) == 0 or len(w_axis) == 0:
#                 return image, target
#             for _ in range(self.max_trys):
#                 xx = np.random.choice(w_axis, size=2)
#                 xmin = min(xx)
#                 xmax = max(xx)
#                 x_size = xmax - xmin
#                 if x_size > self.max_size or x_size < self.min_size:
#                     continue
#                 yy = np.random.choice(h_axis, size=2)
#                 ymin = min(yy)
#                 ymax = max(yy)
#                 y_size = ymax - ymin
#                 if y_size > self.max_size or y_size < self.min_size:
#                     continue
#                 box_in_area = (
#                     (boxes[:, 0] >= xmin)
#                     & (boxes[:, 1] >= ymin)
#                     & (boxes[:, 2] <= xmax)
#                     & (boxes[:, 3] <= ymax)
#                 )
#                 if len(np.where(box_in_area)[0]) == 0:
#                     continue
#                 im = im[ymin:ymax, xmin:xmax]
#                 target = target.crop([xmin, ymin, xmax, ymax])
#                 return Image.fromarray(im), target
#             return image, target
#         else:
#             return image, target


class RandomHorizontalFlip(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            image = F.hflip(image)
            target = target.transpose(0)
        return image, target


class ToTensor(object):
    def __call__(self, image, target):
        return F.to_tensor(image), target


class Normalize(object):
    def __init__(self, mean, std, to_bgr255=True):
        self.mean = mean
        self.std = std
        self.to_bgr255 = to_bgr255

    def __call__(self, image, target):
        if self.to_bgr255:
            image = image[[2, 1, 0]] * 255
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, target


class RandomBrightness(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            brightness_factor = random.uniform(0.5, 2)
            image = F.adjust_brightness(image, brightness_factor)
        return image, target


class RandomContrast(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            contrast_factor = random.uniform(0.5, 2)
            image = F.adjust_contrast(image, contrast_factor)
        return image, target


class RandomHue(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            hue_factor = random.uniform(-0.25, 0.25)
            image = F.adjust_hue(image, hue_factor)
        return image, target


class RandomSaturation(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            saturation_factor = random.uniform(0.5, 2)
            image = F.adjust_saturation(image, saturation_factor)
        return image, target


class RandomGamma(object):
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            gamma_factor = random.uniform(0.5, 2)
            image = F.adjust_gamma(image, gamma_factor)
        return image, target


class RandomRotate(object):
    def __init__(self, prob, max_theta=30, fix_rotate=False):
        self.prob = prob
        self.max_theta = max_theta
        self.fix_rotate = fix_rotate

    def __call__(self, image, target):
        if random.random() < self.prob and target is not None:
            # try:
            if self.fix_rotate:
                delta = 30
            else:
                delta = random.uniform(-1 * self.max_theta, self.max_theta)
            width, height = image.size
            ## get the minimal rect to cover the rotated image
            img_box = [[[0, 0], [width, 0], [width, height], [0, height]]]
            rotated_img_box = _quad2minrect(
                _rotate_polygons(img_box, delta, (width / 2, height / 2))
            )
            r_height = int(
                max(rotated_img_box[0][3], rotated_img_box[0][1])
                - min(rotated_img_box[0][3], rotated_img_box[0][1])
            )
            r_width = int(
                max(rotated_img_box[0][2], rotated_img_box[0][0])
                - min(rotated_img_box[0][2], rotated_img_box[0][0])
            )
            r_height = max(r_height, height + 1)
            r_width = max(r_width, width + 1)

            ## padding im
            im_padding = np.zeros((r_height, r_width, 3))
            start_h, start_w = (
                int((r_height - height) / 2.0),
                int((r_width - width) / 2.0),
            )
            end_h, end_w = start_h + height, start_w + width
            im_padding[start_h:end_h, start_w:end_w, :] = image

            M = cv2.getRotationMatrix2D((r_width / 2, r_height / 2), delta, 1)
            im = cv2.warpAffine(im_padding, M, (r_width, r_height))
            im = Image.fromarray(im.astype(np.uint8))
            target = target.rotate(
                -delta, (r_width / 2, r_height / 2), start_h, start_w
            )
            return im, target
            # except:
            #    return image, target
        else:
            return image, target


def _quad2minrect(boxes):
    ## trans a quad(N*4) to a rectangle(N*4) which has miniual area to cover it
    return np.hstack(
        (
            boxes[:, ::2].min(axis=1).reshape((-1, 1)),
            boxes[:, 1::2].min(axis=1).reshape((-1, 1)),
            boxes[:, ::2].max(axis=1).reshape((-1, 1)),
            boxes[:, 1::2].max(axis=1).reshape((-1, 1)),
        )
    )


def _boxlist2quads(boxlist):
    res = np.zeros((len(boxlist), 8))
    for i, box in enumerate(boxlist):
        # print(box)
        res[i] = np.array(
            [
                box[0][0],
                box[0][1],
                box[1][0],
                box[1][1],
                box[2][0],
                box[2][1],
                box[3][0],
                box[3][1],
            ]
        )
    return res


def _rotate_polygons(polygons, angle, r_c):
    ## polygons: N*8
    ## r_x: rotate center x
    ## r_y: rotate center y
    ## angle: -15~15

    rotate_boxes_list = []
    for poly in polygons:
        box = Polygon(poly)
        rbox = affinity.rotate(box, angle, r_c)
        if len(list(rbox.exterior.coords)) < 5:
            print("img_box_ori:", poly)
            print("img_box_rotated:", rbox)
        # assert(len(list(rbox.exterior.coords))>=5)
        rotate_boxes_list.append(rbox.boundary.coords[:-1])
    res = _boxlist2quads(rotate_boxes_list)
    return res