# 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"]])