from collections import defaultdict
import json
import numpy as np
import os
import PIL.Image
import PIL.ImageDraw

from ..meta.getter_dataset import GetterDataset
import cv2

try:
    from pycocotools import mask as coco_mask
except ImportError:
    pass


class COCOInstancesBaseDataset(GetterDataset):
    def __init__(
        self,
        data_dir="auto",
        split="train",
        year="2017",
        use_crowded=False,
        read_img_with="pil",
    ):
        if year == "2017" and split in ["minival", "valminusminival"]:
            raise ValueError(
                "coco2017 dataset does not support given split: {}".format(split)
            )

        super(COCOInstancesBaseDataset, self).__init__()
        self.use_crowded = use_crowded
        self.read_img_with = read_img_with

        if split in ["val", "minival", "valminusminival"]:
            img_split = "val"
        else:
            img_split = "train"

        self.img_root = os.path.join(data_dir, "{}{}".format(img_split, year))
        tmp = self.img_root
        if not os.path.exists(self.img_root):
            self.img_root = os.path.join(data_dir, "images")
        assert os.path.exists(
            self.img_root
        ), "we have tried img_root: {} and {}, either found, did u make sure it exists?".format(
            tmp, self.img_root
        )
        anno_path = os.path.join(
            data_dir, "annotations", "instances_{}{}.json".format(split, year)
        )

        self.data_dir = data_dir
        annos = json.load(open(anno_path, "r"))

        self.id_to_prop = {}
        for prop in annos["images"]:
            self.id_to_prop[prop["id"]] = prop
        self.ids = sorted(list(self.id_to_prop.keys()))

        self.cat_ids = [cat["id"] for cat in annos["categories"]]

        self.id_to_anno = defaultdict(list)
        for anno in annos["annotations"]:
            self.id_to_anno[anno["image_id"]].append(anno)

        self.add_getter("img", self._get_image)
        self.add_getter("mask", self._get_mask)
        self.add_getter(["bbox", "label", "area", "crowded"], self._get_annotations)

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

    def _get_image(self, i):
        img_path = os.path.join(
            self.img_root, self.id_to_prop[self.ids[i]]["file_name"]
        )
        if self.read_img_with == "cv2":
            img = cv2.imread(img_path)
        else:
            img = PIL.Image.open(img_path)
        return img

    def _get_mask(self, i):
        # List[{'segmentation', 'area', 'iscrowd',
        #       'image_id', 'bbox', 'category_id', 'id'}]
        annotation = self.id_to_anno[self.ids[i]]
        H = self.id_to_prop[self.ids[i]]["height"]
        W = self.id_to_prop[self.ids[i]]["width"]

        mask = []
        crowded = []
        for anno in annotation:
            msk = self._segm_to_mask(anno["segmentation"], (H, W))
            # FIXME: some of minival annotations are malformed.
            if msk.shape != (H, W):
                continue
            mask.append(msk)
            crowded.append(anno["iscrowd"])
        mask = np.array(mask, dtype=np.bool)
        crowded = np.array(crowded, dtype=np.bool)
        if len(mask) == 0:
            mask = np.zeros((0, H, W), dtype=np.bool)

        if not self.use_crowded:
            not_crowded = np.logical_not(crowded)
            mask = mask[not_crowded]
        return mask

    def _get_annotations(self, i):
        # List[{'segmentation', 'area', 'iscrowd',
        #       'image_id', 'bbox', 'category_id', 'id'}]
        annotation = self.id_to_anno[self.ids[i]]
        bbox = np.array([ann["bbox"] for ann in annotation], dtype=np.float32)
        if len(bbox) == 0:
            bbox = np.zeros((0, 4), dtype=np.float32)
        # (x, y, width, height)  -> (x_min, y_min, x_max, y_max)
        bbox[:, 2] = bbox[:, 0] + bbox[:, 2]
        bbox[:, 3] = bbox[:, 1] + bbox[:, 3]
        # (x_min, y_min, x_max, y_max) -> (y_min, x_min, y_max, x_max)
        bbox = bbox[:, [1, 0, 3, 2]]

        label = np.array(
            [self.cat_ids.index(ann["category_id"]) for ann in annotation],
            dtype=np.int32,
        )

        area = np.array([ann["area"] for ann in annotation], dtype=np.float32)

        crowded = np.array([ann["iscrowd"] for ann in annotation], dtype=np.bool)

        # Remove invalid boxes
        bbox_area = np.prod(bbox[:, 2:] - bbox[:, :2], axis=1)
        keep_mask = np.logical_and(bbox[:, 0] <= bbox[:, 2], bbox[:, 1] <= bbox[:, 3])
        keep_mask = np.logical_and(keep_mask, bbox_area > 0)

        if not self.use_crowded:
            keep_mask = np.logical_and(keep_mask, np.logical_not(crowded))

        bbox = bbox[keep_mask]
        label = label[keep_mask]
        area = area[keep_mask]
        crowded = crowded[keep_mask]
        return bbox, label, area, crowded

    def _segm_to_mask(self, segm, size):
        # Copied from pycocotools.coco.COCO.annToMask
        H, W = size
        if isinstance(segm, list):
            # polygon -- a single object might consist of multiple parts
            # we merge all parts into one mask rle code
            mask = np.zeros((H, W), dtype=np.uint8)
            mask = PIL.Image.fromarray(mask)
            for sgm in segm:
                xy = np.array(sgm).reshape((-1, 2))
                xy = [tuple(xy_i) for xy_i in xy]
                PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
            mask = np.asarray(mask)
        elif isinstance(segm["counts"], list):
            rle = coco_mask.frPyObjects(segm, H, W)
            mask = coco_mask.decode(rle)
        else:
            mask = coco_mask.decode(segm)
        return mask.astype(np.bool)


class COCOBboxDataset(COCOInstancesBaseDataset):
    def __init__(
        self,
        data_dir="auto",
        split="train",
        year="2017",
        use_crowded=False,
        read_img_with="pil",
        return_area=False,
        return_crowded=False,
    ):
        super(COCOBboxDataset, self).__init__(
            data_dir, split, year, use_crowded, read_img_with
        )
        keys = ("img", "bbox", "label")
        if return_area:
            keys += ("area",)
        if return_crowded:
            keys += ("crowded",)
        self.keys = keys
