#
# Copyright 2021 Intel Corporation.
#
# LEGAL NOTICE: Your use of this software and any required dependent software
# (the "Software Package") is subject to the terms and conditions of
# the Intel(R) OpenVINO(TM) Distribution License for the Software Package,
# which may also include notices, disclaimers, or license terms for
# third party or open source software included in or with the Software Package,
# and your use indicates your acceptance of all such terms. Please refer
# to the "third-party-programs.txt" or other similarly-named text file
# included with the Software Package for additional details.

import os
import json
import cv2
import numpy as np

from compression.api import DataLoader


class COCOLoader(DataLoader):
    def __init__(self, config):
        super().__init__(config)
        self.images_path = config.images_path
        self.annotation_path = config.annotation_path
        self.images = os.listdir(self.images_path)
        self.labels = None
        self.data, self.bbox = self.prepare_annotation()

    def prepare_annotation(self):
        with open(self.annotation_path) as f:
            file = json.load(f)
        self.labels = [i['id'] for i in file['categories']]
        data = {}
        for idx, image in enumerate(file['images']):
            data[idx] = {'file_name': image['file_name'], 'image_id': image['id'],
                         'height': image['height'], 'width': image['width']}
        bbox = {}
        for i in file['annotations']:
            if i['image_id'] not in bbox.keys():
                bbox[i['image_id']] = {'bbox': [self.prepare_bbox(*i['bbox'])],
                                       'category_id': [i['category_id']],
                                       'iscrowd': [i['iscrowd']]}
            else:
                bbox[i['image_id']]['bbox'].append(self.prepare_bbox(*i['bbox']))
                bbox[i['image_id']]['category_id'].append(i['category_id'])
                bbox[i['image_id']]['iscrowd'].append(i['iscrowd'])

        return data, bbox

    def __getitem__(self, index):
        """ Returns (img_id, img_annotation), image"""
        if index >= len(self):
            raise IndexError

        shape_image = self.data[index]['height'], self.data[index]['width']

        bbox = np.array(self.bbox.get(self.data[index]['image_id'], {}).get('bbox', []))

        if bbox.size != 0:
            # x_maxs = np.max(bbox, axis=1)
            # y_maxs = np.max(bbox, axis=1)
            # x_mins = np.min(bbox, axis=1)
            # y_mins = np.min(bbox, axis=1)
            x_maxs = bbox[:, 2]
            y_maxs = bbox[:, 3]
            x_mins = bbox[:, 0]
            y_mins = bbox[:, 1]
        else:
            x_maxs, y_maxs, x_mins, y_mins = [], [], [], []

        labels = np.array(self.bbox.get(self.data[index]['image_id'], {}).get('category_id', []))
        iscrowd = np.array(self.bbox.get(self.data[index]['image_id'], {}).get('iscrowd', []))

        annotation = {'boxes': bbox, 'labels': labels, 'iscrowd': iscrowd,
                      'x_maxs': x_maxs, 'x_mins': x_mins, 'y_maxs': y_maxs, 'y_mins': y_mins}
        annotation = (index, [annotation, shape_image])
        return annotation, self._read_and_preprocess_image(os.path.join(self.images_path, self.data[index]['file_name']))

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

    @staticmethod
    def _read_and_preprocess_image(image_path):
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        image = cv2.resize(image, (640, 640)).astype(np.float32)
        image /= 255
        return image.transpose(2, 0, 1)[::-1]

    @staticmethod
    def prepare_bbox(x, y, weight, height):
        return np.array([x, y, x + weight, y + height])