#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import json
import os
from pathlib import Path

import numpy as np
import torch
from detectron2.utils.file_io import PathManager
from detr.util.box_ops import masks_to_boxes
from panopticapi.utils import rgb2id
from PIL import Image

from .coco import make_coco_transforms


class CocoPanoptic:
    def __init__(
        self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True
    ):
        with PathManager.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):
        ann_info = (
            self.coco["annotations"][idx]
            if "annotations" in self.coco
            else self.coco["images"][idx]
        )
        img_path = os.path.join(
            self.img_folder, ann_info["file_name"].replace(".png", ".jpg")
        )
        ann_path = os.path.join(self.ann_folder, ann_info["file_name"])

        with PathManager.open(img_path, "rb") as f:
            img = Image.open(f).convert("RGB")
        w, h = img.size
        if "segments_info" in ann_info:
            with PathManager.open(ann_path, "rb") as f:
                masks = np.asarray(Image.open(f), 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
        target["labels"] = labels

        target["boxes"] = masks_to_boxes(masks)

        target["size"] = torch.as_tensor([int(h), int(w)])
        target["orig_size"] = torch.as_tensor([int(h), int(w)])
        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

    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):
    if "manifold" in args.coco_path:
        root = args.coco_path
        PATHS = {
            "train": (
                os.path.join(root, "coco_train2017"),
                "manifold://fair_vision_data/tree/detectron2/json_dataset_annotations/coco/panoptic_train2017.json",
            ),
            "val": (
                os.path.join(root, "coco_val2017"),
                "manifold://fair_vision_data/tree/detectron2/json_dataset_annotations/coco/panoptic_val2017.json",
            ),
        }
        img_folder_path, ann_file = PATHS[image_set]
        ann_folder = os.path.join(root, f"coco_panoptic_{image_set}2017")
    else:
        img_folder_root = Path(args.coco_path)
        ann_folder_root = Path(args.coco_panoptic_path)
        assert (
            img_folder_root.exists()
        ), f"provided COCO path {img_folder_root} does not exist"
        assert (
            ann_folder_root.exists()
        ), f"provided COCO path {ann_folder_root} does not exist"
        mode = "panoptic"
        PATHS = {
            "train": ("train2017", Path("annotations") / f"{mode}_train2017.json"),
            "val": ("val2017", Path("annotations") / f"{mode}_val2017.json"),
        }
        img_folder, ann_file = PATHS[image_set]
        img_folder_path = img_folder_root / img_folder

        ann_folder = ann_folder_root / f"{mode}_{img_folder}"
        ann_file = ann_folder_root / ann_file

    dataset = CocoPanoptic(
        img_folder_path,
        ann_folder,
        ann_file,
        transforms=make_coco_transforms(image_set),
        return_masks=args.masks,
    )

    return dataset
