import PIL  # version 1.2.0
import torch
import os
import torchvision.transforms.functional as F
import numpy as np
import random


def intersect(boxes1, boxes2):
    """
    Find intersection of every box combination between two sets of box
    boxes1: bounding boxes 1, a tensor of dimensions (n1, 4)
    boxes2: bounding boxes 2, a tensor of dimensions (n2, 4)

    Out: Intersection each of boxes1 with respect to each of boxes2,
         a tensor of dimensions (n1, n2)
    """
    n1 = boxes1.size(0)
    n2 = boxes2.size(0)
    max_xy = torch.min(
        boxes1[:, 2:].unsqueeze(1).expand(n1, n2, 2),
        boxes2[:, 2:].unsqueeze(0).expand(n1, n2, 2),
    )

    min_xy = torch.max(
        boxes1[:, :2].unsqueeze(1).expand(n1, n2, 2),
        boxes2[:, :2].unsqueeze(0).expand(n1, n2, 2),
    )
    inter = torch.clamp(max_xy - min_xy, min=0)  # (n1, n2, 2)
    return inter[:, :, 0] * inter[:, :, 1]  # (n1, n2)


def find_IoU(boxes1, boxes2):
    """
    Find IoU between every boxes set of boxes
    boxes1: a tensor of dimensions (n1, 4) (left, top, right , bottom)
    boxes2: a tensor of dimensions (n2, 4)

    Out: IoU each of boxes1 with respect to each of boxes2, a tensor of
         dimensions (n1, n2)

    Formula:
    (box1 ∩ box2) / (box1 u box2) = (box1 ∩ box2) / (area(box1) + area(box2) - (box1 ∩ box2 ))
    """
    inter = intersect(boxes1, boxes2)
    area_boxes1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    area_boxes2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])

    area_boxes1 = area_boxes1.unsqueeze(1).expand_as(inter)  # (n1, n2)
    area_boxes2 = area_boxes2.unsqueeze(0).expand_as(inter)  # (n1, n2)
    union = area_boxes1 + area_boxes2 - inter
    return inter / union


def random_crop(image, boxes, labels, difficulties=None):
    """
    image: A PIL image
    boxes: Bounding boxes, a tensor of dimensions (#objects, 4)
    labels: labels of object, a tensor of dimensions (#objects)
    difficulties: difficulties of detect object, a tensor of dimensions (#objects)

    Out: cropped image , new boxes, new labels, new difficulties
    """
    if type(image) == PIL.Image.Image:
        image = F.to_tensor(image)
    original_h = image.size(1)
    original_w = image.size(2)

    while True:
        mode = random.choice([0.1, 0.3, 0.5, 0.9, None])

        if mode is None:
            return F.to_pil_image(image), boxes, labels, difficulties

        new_image = image
        new_boxes = boxes
        new_difficulties = difficulties
        new_labels = labels
        for _ in range(50):
            # Crop dimensions: [0.3, 1] of original dimensions
            new_h = random.uniform(0.3 * original_h, original_h)
            new_w = random.uniform(0.3 * original_w, original_w)

            # Aspect ratio constraint b/t .5 & 2
            if new_h / new_w < 0.5 or new_h / new_w > 2:
                continue

            # Crop coordinate
            left = random.uniform(0, original_w - new_w)
            right = left + new_w
            top = random.uniform(0, original_h - new_h)
            bottom = top + new_h
            crop = torch.FloatTensor(
                [int(left), int(top), int(right), int(bottom)]
            )

            # Calculate IoU  between the crop and the bounding boxes
            overlap = find_IoU(crop.unsqueeze(0), boxes)  # (1, #objects)
            overlap = overlap.squeeze(0)

            # If not a single bounding box has a IoU of greater than the minimum, try again
            if overlap.shape[0] == 0:
                continue
            if overlap.max().item() < mode:
                continue

            # Crop
            new_image = image[
                :, int(top) : int(bottom), int(left) : int(right)
            ]  # (3, new_h, new_w)

            # Center of bounding boxes
            center_bb = (boxes[:, :2] + boxes[:, 2:]) / 2.0

            # Find bounding box has been had center in crop
            center_in_crop = (
                (center_bb[:, 0] > left)
                * (center_bb[:, 0] < right)
                * (center_bb[:, 1] > top)
                * (center_bb[:, 1] < bottom)
            )  # ( #objects)

            if not center_in_crop.any():
                continue

            # take matching bounding box
            new_boxes = boxes[center_in_crop, :]

            # take matching labels
            new_labels = labels[center_in_crop]

            # take matching difficulities
            if difficulties is not None:
                new_difficulties = difficulties[center_in_crop]
            else:
                new_difficulties = None

            # Use the box left and top corner or the crop's
            new_boxes[:, :2] = torch.max(new_boxes[:, :2], crop[:2])

            # adjust to crop
            new_boxes[:, :2] -= crop[:2]

            new_boxes[:, 2:] = torch.min(new_boxes[:, 2:], crop[2:])

            # adjust to crop
            new_boxes[:, 2:] -= crop[:2]

            return (
                F.to_pil_image(new_image),
                new_boxes,
                new_labels,
                new_difficulties,
            )
