import os
import json
import numpy as np
import pickle
from typing import Any

from pycocotools.coco import COCO
from torch.utils.data import Dataset


class DetectionMSCOCODataset(Dataset):
    def __init__(self, annotation_file: str, image_dir: str):

        self._annotation_file = annotation_file
        self._image_dir = image_dir
        self._cache_file = self._annotation_file + ".cache"

        self._coco = COCO(self._annotation_file)

        self._img_ids = self._coco.getImgIds()
        self._cat_ids = self._coco.getCatIds()
        self._ann_ids = self._coco.getAnnIds()

        self._data = "coco"
        self._classes = {
            ind: cat_id for ind, cat_id in enumerate(self._cat_ids)
        }
        self._coco_to_class_map = {
            value: key for key, value in self._classes.items()
        }

        self._load_data()
        self._db_inds = np.arange(len(self._image_names))

        self._load_coco_data()

    def _load_data(self):
        print("loading from cache file: {}".format(self._cache_file))
        if not os.path.exists(self._cache_file):
            print("No cache file found...")
            self._extract_data()
            with open(self._cache_file, "wb") as f:
                pickle.dump([self._detections, self._image_names], f)
            print("Cache file created")
        else:
            with open(self._cache_file, "rb") as f:
                self._detections, self._image_names = pickle.load(f)

    def _load_coco_data(self):
        with open(self._annotation_file, "r") as f:
            data = json.load(f)

        coco_ids = self._coco.getImgIds()
        eval_ids = {
            self._coco.loadImgs(coco_id)[0]["file_name"]: coco_id
            for coco_id in coco_ids
            }

        self._coco_categories = data["categories"]
        self._coco_eval_ids = eval_ids

    def class_name(self, cid):
        cat_id = self._classes[cid]
        cat = self._coco.loadCats([cat_id])[0]
        return cat["name"]

    def _extract_data(self):

        self._image_names = [
            self._coco.loadImgs(img_id)[0]["file_name"]
            for img_id in self._img_ids
        ]
        self._detections = {}
        for ind, (coco_image_id, image_name) in enumerate(zip(self._img_ids, self._image_names)):
            image = self._coco.loadImgs(coco_image_id)[0]
            bboxes = []
            categories = []

            for cat_id in self._cat_ids:
                annotation_ids = self._coco.getAnnIds(imgIds=image["id"], catIds=cat_id)
                annotations = self._coco.loadAnns(annotation_ids)
                category = self._coco_to_class_map[cat_id]
                for annotation in annotations:
                    bbox = np.array(annotation["bbox"])
                    bbox[[2, 3]] += bbox[[0, 1]]
                    bboxes.append(bbox)

                    categories.append(category)

            self._detections[image_name] = [{
                'bbox': bbox.astype(np.float32),
                'category_id': category,
                'category_name': self.class_name(category)
            } for bbox, category in zip(bboxes, categories)]

    def __getitem__(self, ind: int) -> Any:
        image_name = self._image_names[ind]

        return {
            'image_name': os.path.join(self._image_dir, image_name),
            'detections': self._detections[image_name]
        }

    def __len__(self) -> int:
        return len(self._img_ids)

    def get_num_classes(self) -> int:
        return len(self._cat_ids)

[((18, 21, 18, 48), 'pycocotools.coco.COCO', 'COCO', ({(18, 26, 18, 47): 'self._annotation_file'}, {}), '(self._annotation_file)', False, 'from pycocotools.coco import COCO\n'), ((39, 15, 39, 47), 'os.path.exists', 'os.path.exists', ({(39, 30, 39, 46): 'self._cache_file'}, {}), '(self._cache_file)', False, 'import os\n'), ((51, 19, 51, 31), 'json.load', 'json.load', ({(51, 29, 51, 30): 'f'}, {}), '(f)', False, 'import json\n'), ((100, 26, 100, 67), 'os.path.join', 'os.path.join', ({(100, 39, 100, 54): 'self._image_dir', (100, 56, 100, 66): 'image_name'}, {}), '(self._image_dir, image_name)', False, 'import os\n'), ((43, 16, 43, 69), 'pickle.dump', 'pickle.dump', ({(43, 28, 43, 65): '[self._detections, self._image_names]', (43, 67, 43, 68): 'f'}, {}), '([self._detections, self._image_names], f)', False, 'import pickle\n'), ((47, 54, 47, 68), 'pickle.load', 'pickle.load', ({(47, 66, 47, 67): 'f'}, {}), '(f)', False, 'import pickle\n'), ((84, 27, 84, 55), 'numpy.array', 'np.array', ({(84, 36, 84, 54): "annotation['bbox']"}, {}), "(annotation['bbox'])", True, 'import numpy as np\n')]