#!/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 os
import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

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

from .coco import make_coco_transforms
from detectron2.utils.file_io import PathManager


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
