import os
import random

import torch
import torchvision.transforms as transforms
from PIL import Image


def recalculate_box_and_verify_if_valid(x, y, w, h, image_size, original_image_size, min_box_size):
    scale = image_size / min(original_image_size)
    crop_y = (original_image_size[1] * scale - image_size) // 2
    crop_x = (original_image_size[0] * scale - image_size) // 2
    x0 = max(x * scale - crop_x, 0)
    y0 = max(y * scale - crop_y, 0)
    x1 = min((x + w) * scale - crop_x, image_size)
    y1 = min((y + h) * scale - crop_y, image_size)
    if (x1 - x0) * (y1 - y0) / (image_size * image_size) < min_box_size:
        return False, (None, None, None, None)
    return True, (x0, y0, x1, y1)


class COCODataset(torch.utils.data.Dataset):
    def __init__(
        self,
        data_path,
        image_path,
        image_size=512,
        min_box_size=0.01,
        max_boxes_per_data=8,
        tokenizer=None,
    ):
        super().__init__()
        self.min_box_size = min_box_size
        self.max_boxes_per_data = max_boxes_per_data
        self.image_size = image_size
        self.image_path = image_path
        self.tokenizer = tokenizer
        self.transforms = transforms.Compose(
            [
                transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )

        self.data_list = torch.load(data_path, map_location="cpu")

    def __getitem__(self, index):
        if self.max_boxes_per_data > 99:
            assert False, "Are you sure setting such large number of boxes per image?"

        out = {}

        data = self.data_list[index]
        image = Image.open(os.path.join(self.image_path, data["file_path"])).convert("RGB")
        original_image_size = image.size
        out["pixel_values"] = self.transforms(image)

        annos = data["annos"]

        areas, valid_annos = [], []
        for anno in annos:
            # x, y, w, h = anno['bbox']
            x0, y0, x1, y1 = anno["bbox"]
            x, y, w, h = x0, y0, x1 - x0, y1 - y0
            valid, (x0, y0, x1, y1) = recalculate_box_and_verify_if_valid(
                x, y, w, h, self.image_size, original_image_size, self.min_box_size
            )
            if valid:
                anno["bbox"] = [x0, y0, x1, y1]
                areas.append((x1 - x0) * (y1 - y0))
                valid_annos.append(anno)

        # Sort according to area and choose the largest N objects
        wanted_idxs = torch.tensor(areas).sort(descending=True)[1]
        wanted_idxs = wanted_idxs[: self.max_boxes_per_data]
        valid_annos = [valid_annos[i] for i in wanted_idxs]

        out["boxes"] = torch.zeros(self.max_boxes_per_data, 4)
        out["masks"] = torch.zeros(self.max_boxes_per_data)
        out["text_embeddings_before_projection"] = torch.zeros(self.max_boxes_per_data, 768)

        for i, anno in enumerate(valid_annos):
            out["boxes"][i] = torch.tensor(anno["bbox"]) / self.image_size
            out["masks"][i] = 1
            out["text_embeddings_before_projection"][i] = anno["text_embeddings_before_projection"]

        prob_drop_boxes = 0.1
        if random.random() < prob_drop_boxes:
            out["masks"][:] = 0

        caption = random.choice(data["captions"])

        prob_drop_captions = 0.5
        if random.random() < prob_drop_captions:
            caption = ""
        caption = self.tokenizer(
            caption,
            max_length=self.tokenizer.model_max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        out["caption"] = caption

        return out

    def __len__(self):
        return len(self.data_list)