# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from panopticapi.utils import rgb2id
# from util.box_ops import masks_to_boxes

from .construction import make_construction_transforms

import logging


def box_xywh_to_xyxy(x):
    xs, ys, w, h = x.unbind(-1)
    b = [xs, ys, (xs + w), (ys + h)]
    return torch.stack(b, dim=-1)


def masks_to_boxes(segments):
    boxes = []
    labels = []
    iscrowd = []
    area = []

    for ann in segments:
        if len(ann["bbox"]) == 4:
            boxes.append(ann["bbox"])
            area.append(ann['area'])
        else:
            boxes.append([0, 0, 2, 2])
            area.append(4)
        labels.append(ann["category_id"])
        iscrowd.append(ann['iscrowd'])
    
    if len(boxes) == 0 and len(labels) == 0:
        boxes.append([0, 0, 2, 2])
        labels.append(1)
        area.append(4)
        iscrowd.append(0)

    boxes = torch.tensor(boxes, dtype=torch.int64)
    labels = torch.tensor(labels, dtype=torch.int64)
    iscrowd = torch.tensor(iscrowd)
    area = torch.tensor(area)
    boxes = box_xywh_to_xyxy(boxes)
    return boxes, labels, iscrowd, area

class ConstructionPanoptic:
    def __init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True):
        with open(ann_file, "r") as f:
            self.coco = json.load(f)

        # sort 'images' field so that they are aligned with 'annotations'
        # i.e., in alphabetical order
        self.coco["images"] = sorted(self.coco["images"], key=lambda x: x["id"])
        # sanity check
        if "annotations" in self.coco:
            for img, ann in zip(self.coco["images"], self.coco["annotations"]):
                assert img["file_name"][:-4] == ann["file_name"][:-4]

        self.img_folder = img_folder
        self.ann_folder = ann_folder
        self.ann_file = ann_file
        self.transforms = transforms
        self.return_masks = return_masks

    def __getitem__(self, idx):
        try:
            ann_info = (
                self.coco["annotations"][idx]
                if "annotations" in self.coco
                else self.coco["images"][idx]
            )
            img_path = Path(self.img_folder) / ann_info["file_name"].replace(".png", ".jpg")
            ann_path = Path(self.ann_folder) / ann_info["file_name"]

            img = Image.open(img_path).convert("RGB")
            w, h = img.size
            if "segments_info" in ann_info:
                masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
                masks = rgb2id(masks)

                ids = np.array([ann["id"] for ann in ann_info["segments_info"]])
                masks = masks == ids[:, None, None]

                masks = torch.as_tensor(masks, dtype=torch.uint8)
                
                # labels = torch.tensor(
                #     [ann["category_id"] for ann in ann_info["segments_info"]],
                #     dtype=torch.int64,
                # )


            target = {}
            target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
            if self.return_masks:
                target['masks'] = masks

            boxes, labels, iscrowd, area = masks_to_boxes(ann_info["segments_info"])

            target['labels'] = labels

            # Instead of finding boxes, just take the one from json info available 
            # target["boxes"] = masks_to_boxes(ann_info["segments_info"])
            target["boxes"] = boxes


            target['size'] = torch.as_tensor([int(h), int(w)])
            target['orig_size'] = torch.as_tensor([int(h), int(w)])

            target['iscrowd'] = iscrowd
            target['area'] = area
            # if "segments_info" in ann_info:
            #     for name in ['iscrowd', 'area']:
            #         target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])

            if self.transforms is not None:
                img, target = self.transforms(img, target)

            return img, target

        except Exception as e:
            logging.error(ann_info)
            raise e


    def __len__(self):
        return len(self.coco['images'])

    def get_height_and_width(self, idx):
        img_info = self.coco['images'][idx]
        height = img_info['height']
        width = img_info['width']
        return height, width


def build(image_set, args):
    root = Path(args.data_path)
    assert (
        root.exists()
    ), f"provided Panoptic path {root} does not exist"

    mode = "panoptic"

    PATHS = {
        "train": ("images", f"{mode}", f"{mode}.json"),
        "val": ("images", f"val_{mode}", f"val_{mode}.json"),
    }

    img_folder, ann_folder, ann_file = PATHS[image_set]
    img_folder_path = root / img_folder
    ann_folder_path = root / ann_folder
    ann_file = root / ann_file

    dataset = ConstructionPanoptic(
        img_folder_path,
        ann_folder_path,
        ann_file,
        transforms=make_construction_transforms(image_set),
        return_masks=args.masks,
    )

    return dataset

[((20, 11, 20, 33), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((45, 12, 45, 50), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((46, 13, 46, 52), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((47, 14, 47, 35), 'torch.tensor', 'torch.tensor', ({(47, 27, 47, 34): 'iscrowd'}, {}), '(iscrowd)', False, 'import torch\n'), ((48, 11, 48, 29), 'torch.tensor', 'torch.tensor', ({(48, 24, 48, 28): 'area'}, {}), '(area)', False, 'import torch\n'), ((142, 11, 142, 31), 'pathlib.Path', 'Path', ({(142, 16, 142, 30): 'args.data_path'}, {}), '(args.data_path)', False, 'from pathlib import Path\n'), ((55, 24, 55, 36), 'json.load', 'json.load', ({(55, 34, 55, 35): 'f'}, {}), '(f)', False, 'import json\n'), ((99, 33, 99, 115), 'torch.tensor', 'torch.tensor', ({(99, 46, 99, 114): "[ann_info['image_id'] if 'image_id' in ann_info else ann_info['id']]"}, {}), "([ann_info['image_id'] if 'image_id' in ann_info else ann_info[\n    'id']])", False, 'import torch\n'), ((78, 23, 78, 44), 'pathlib.Path', 'Path', ({(78, 28, 78, 43): 'self.img_folder'}, {}), '(self.img_folder)', False, 'from pathlib import Path\n'), ((79, 23, 79, 44), 'pathlib.Path', 'Path', ({(79, 28, 79, 43): 'self.ann_folder'}, {}), '(self.ann_folder)', False, 'from pathlib import Path\n'), ((85, 24, 85, 37), 'panopticapi.utils.rgb2id', 'rgb2id', ({(85, 31, 85, 36): 'masks'}, {}), '(masks)', False, 'from panopticapi.utils import rgb2id\n'), ((87, 22, 87, 80), 'numpy.array', 'np.array', ({(87, 31, 87, 79): "[ann['id'] for ann in ann_info['segments_info']]"}, {}), "([ann['id'] for ann in ann_info['segments_info']])", True, 'import numpy as np\n'), ((90, 24, 90, 65), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((127, 12, 127, 35), 'logging.error', 'logging.error', ({(127, 26, 127, 34): 'ann_info'}, {}), '(ann_info)', False, 'import logging\n'), ((81, 18, 81, 38), 'PIL.Image.open', 'Image.open', ({(81, 29, 81, 37): 'img_path'}, {}), '(img_path)', False, 'from PIL import Image\n'), ((84, 35, 84, 55), 'PIL.Image.open', 'Image.open', ({(84, 46, 84, 54): 'ann_path'}, {}), '(ann_path)', False, 'from PIL import Image\n')]