"""
Cityscape Database
"""
from __future__ import print_function

import cv2
import os
import numpy as np
import multiprocessing

from six.moves import cPickle
from multiprocessing.dummy import Pool
from imdb import IMDB
from ..processing.bbox_transform import bbox_overlaps
from ..core.helper import get_scale_factor


class Cityscape(IMDB):
    def __init__(self, image_set, root_path, dataset_path, load_memory=False, use_mask=False, panoptic=False):
        """
        fill basic information to initialize imdb
        :param image_set: train or val or trainval or test
        :param root_path: 'cache' and 'rpn_data'
        :param dataset_path: data and results
        :return: imdb object
        """
        super(Cityscape, self).__init__('cityscape', image_set, root_path, dataset_path)
        self.image_set = image_set
        self.root_path = root_path
        self.data_path = dataset_path
        self.load_memory = load_memory
        self.use_mask = use_mask
        self.panoptic = panoptic

        self.classes = ['__background__', 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle']
        self.class_id = [0, 24, 25, 26, 27, 28, 31, 32, 33]
        self.num_classes = len(self.classes)
        self.image_set_index = self.load_image_set_index()
        self.num_images = len(self.image_set_index)
        print('num_images %d' % self.num_images)

    def load_image_set_index(self):
        """
        find out which indexes correspond to given image set (train or val)
        :return:
        """
        image_set_index_file = os.path.join(self.data_path, 'imglists', self.image_set + '.lst')
        assert os.path.exists(image_set_index_file), 'Path does not exist: {}'.format(image_set_index_file)
        image_set_index = []
        with open(image_set_index_file, 'r') as f:
            for line in f:
                if len(line) > 1:
                    label = line.strip().split('\t')
                    image_set_index.append(label[1])
        return image_set_index

    def image_path_from_index(self, index):
        """
        given image index, find out full path
        :param index: index of a specific image
        :return: full path of this image
        """
        image_file = os.path.join(self.data_path, index)
        assert os.path.exists(image_file), 'Path does not exist: {}'.format(image_file)
        return image_file

    def gt_roidb(self):
        """
        return ground truth image regions database
        :return: imdb[image_index]['boxes', 'gt_classes', 'ins_id', 'ins_seg', 'gt_overlaps', 'flipped']
        """
        cache_file = os.path.join(self.cache_path, self.name + '_gt_roidb.pkl')
        if os.path.exists(cache_file):
            with open(cache_file, 'rb') as f:
                roidb = cPickle.load(f)
            print('{} gt roidb loaded from {}'.format(self.name, cache_file))
            return roidb
        gt_roidb = self.load_cityscape_annotations()
        with open(cache_file, 'wb') as f:
            cPickle.dump(gt_roidb, f, cPickle.HIGHEST_PROTOCOL)

        return gt_roidb

    def load_from_seg(self, ins_seg_gt_path):
        ins_seg_gt_path = os.path.join(self.data_path, ins_seg_gt_path)
        print(ins_seg_gt_path)
        assert os.path.exists(ins_seg_gt_path), 'Path does not exist: {}'.format(ins_seg_gt_path)
        ins_seg_mask = cv2.imread(ins_seg_gt_path, -1)
        boxes = []
        gt_classes = []
        ins_id = []
        gt_overlaps = []
        for c in range(1, len(self.class_id)):
            px = np.where((ins_seg_mask >= self.class_id[c] * 1000) & (ins_seg_mask < (self.class_id[c] + 1) * 1000))
            if len(px[0]) == 0:
                continue
            ids = np.unique(ins_seg_mask[px])
            for id in ids:
                px = np.where(ins_seg_mask == id)
                x_min = np.min(px[1])
                y_min = np.min(px[0])
                x_max = np.max(px[1])
                y_max = np.max(px[0])
                if x_max - x_min <= 1 or y_max - y_min <= 1:
                    continue
                boxes.append([x_min, y_min, x_max, y_max])
                gt_classes.append(c)
                ins_id.append(id % 1000)
                overlaps = np.zeros(self.num_classes)
                overlaps[c] = 1
                gt_overlaps.append(overlaps)
        return np.asarray(boxes), np.asarray(gt_classes), np.asarray(ins_id), ins_seg_mask, np.asarray(gt_overlaps)

    def load_cityscape_annotations(self):
        """
        for a given index, load image and bounding boxes info from a single image list
        :return: list of record['boxes', 'gt_classes', 'ins_id', 'ins_seg', 'gt_overlaps', 'flipped']
        """
        imglist_file = os.path.join(self.data_path, 'imglists', self.image_set + '.lst')
        assert os.path.exists(imglist_file), 'Path does not exist: {}'.format(imglist_file)
        imgfiles_list = []
        with open(imglist_file, 'r') as f:
            for line in f:
                file_list = dict()
                label = line.strip().split('\t')
                file_list['img_id'] = label[0]
                file_list['img_path'] = label[1]
                file_list['ins_seg_path'] = label[2].replace('labelTrainIds', 'instanceIds')
                if self.panoptic:
                    file_list['sem_seg_path'] = label[2]
                imgfiles_list.append(file_list)

        assert len(imgfiles_list) == self.num_images, 'number of boxes matrix must match number of images'

        pool = Pool(2 * multiprocessing.cpu_count())
        roidb = pool.map(self.process_image_list, imgfiles_list)

        return roidb

    def process_image_list(self, img_rec):
        roi_rec = dict()
        if self.load_memory:
            im_path = os.path.join(self.data_path, img_rec['img_path'])
            print('loading %s to memory' % im_path)
            im_array = cv2.imread(im_path)
            roi_rec['image'] = im_array
            size = im_array.shape
        else:
            roi_rec['image'] = os.path.join(self.data_path, img_rec['img_path'])
            size = cv2.imread(roi_rec['image']).shape

        roi_rec['height'] = size[0]
        roi_rec['width'] = size[1]
        boxes, gt_classes, ins_id, ins_seg_mask, gt_overlaps = self.load_from_seg(img_rec['ins_seg_path'])
        if boxes.size == 0:
            total_num_objs = 0
            boxes = np.zeros((total_num_objs, 4), dtype=np.uint16)
            gt_overlaps = np.zeros((total_num_objs, self.num_classes), dtype=np.float32)
            gt_classes = np.zeros((total_num_objs,), dtype=np.int32)
        roi_rec['boxes'] = boxes
        roi_rec['gt_classes'] = gt_classes
        roi_rec['ins_id'] = ins_id

        if self.load_memory:
            roi_rec['ins_seg'] = ins_seg_mask
            if self.panoptic:
                roi_rec['sem_seg'] = cv2.imread(img_rec['sem_seg_path'], -1)
        else:
            roi_rec['ins_seg'] = os.path.join(self.data_path, img_rec['ins_seg_path'])
            if self.panoptic:
                roi_rec['sem_seg'] = os.path.join(self.data_path, img_rec['sem_seg_path'])

        roi_rec['max_classes'] = gt_overlaps.argmax(axis=1)  # this field is not for gt bbox
        roi_rec['max_overlaps'] = gt_overlaps.max(axis=1)  # this field is not for gt bbox
        roi_rec['flipped'] = False
        return roi_rec

    def append_flipped_images(self, roidb):
        """
        append flipped images to an roidb
        flip boxes coordinates, images will be actually flipped when loading into network
        :param roidb: [image_index]['boxes', 'gt_classes', 'gt_overlaps', 'flipped']
        :return: roidb: [image_index]['boxes', 'gt_classes', 'gt_overlaps', 'flipped']
        """
        print('append flipped images to roidb')
        assert self.num_images == len(roidb)
        for i in range(self.num_images):
            roi_rec = roidb[i]
            boxes = roi_rec['boxes'].copy()
            if boxes.shape[0] != 0:
                oldx1 = boxes[:, 0].copy()
                oldx2 = boxes[:, 2].copy()
                boxes[:, 0] = roi_rec['width'] - oldx2 - 1
                boxes[:, 2] = roi_rec['width'] - oldx1 - 1
                assert (boxes[:, 2] >= boxes[:, 0]).all(), \
                    'img_name %s, width %d\n' % (roi_rec['image'], roi_rec['width']) + \
                    np.array_str(roi_rec['boxes'], precision=3, suppress_small=True)
            entry = {'image': roi_rec['image'],
                     'height': roi_rec['height'],
                     'width': roi_rec['width'],
                     'boxes': boxes,
                     'gt_classes': roidb[i]['gt_classes'],
                     'max_classes': roidb[i]['max_classes'],
                     'max_overlaps': roidb[i]['max_overlaps'],
                     'ins_seg': roidb[i]['ins_seg'],
                     'ins_id': roidb[i]['ins_id'],
                     'flipped': True}
            if self.panoptic:
                entry.update({'sem_seg': roidb[i]['sem_seg']})
            roidb.append(entry)

        self.image_set_index *= 2
        return roidb

    def evaluate_maskfcn(self, results_pack):
        for result_rec in results_pack['results_list']:
            image_path = result_rec['image']
            im_info = result_rec['im_info']
            detections = result_rec['boxes']
            seg_masks = result_rec['masks']

            filename = image_path.split("/")[-1]
            filename = filename.replace('.png', '')

            result_path = 'eval_ins/'

            print('writing results for: %s' % filename)
            result_txt = os.path.join(result_path, filename)
            result_txt = result_txt + '.txt'
            count = 0
            f = open(result_txt, 'w')

            for j, labelID in enumerate(self.class_id):
                if labelID == 0:
                    continue
                dets = detections[j]
                masks = seg_masks[j]
                for i in range(len(dets)):
                    bbox = dets[i, :4]
                    score = dets[i, -1]
                    bbox = map(int, bbox)
                    mask_image = np.zeros((int(im_info[0, 0]), int(im_info[0, 1])))
                    mask = masks[i, :, :]

                    mask_scale_factor = get_scale_factor(max(bbox[2] - bbox[0], bbox[3] - bbox[1]))
                    mask = cv2.resize(mask, dsize=None, fx=mask_scale_factor, fy=mask_scale_factor)
                    mask = mask[0:(bbox[3]-bbox[1]), 0:(bbox[2]-bbox[0])]

                    mask[mask > 0.5] = 200
                    mask[mask <= 0.5] = 0
                    mask_image[bbox[1]: bbox[3], bbox[0]: bbox[2]] = mask
                    cv2.imwrite(os.path.join(result_path, filename) + '_' + str(count) + '.png', mask_image)
                    f.write('{:s} {:s} {:.8f}\n'.format(filename + '_' + str(count) + '.png', str(labelID), score))
                    count += 1
            f.flush()
            f.close()

    def evaluate_mask(self, results_pack):
        for result_rec in results_pack['results_list']:
            image_path = result_rec['image']
            im_info = result_rec['im_info']
            detections = result_rec['boxes']
            seg_masks = result_rec['masks']

            filename = image_path.split("/")[-1]
            filename = filename.replace('.png', '')

            result_path = 'eval_ins/'

            print('writing results for: %s' % filename)
            result_txt = os.path.join(result_path, filename)
            result_txt = result_txt + '.txt'
            count = 0
            f = open(result_txt, 'w')

            for j, labelID in enumerate(self.class_id):
                if labelID == 0:
                    continue
                dets = detections[j]
                masks = seg_masks[j]
                for i in range(len(dets)):
                    bbox = dets[i, :4]
                    score = dets[i, -1]
                    bbox = map(int, bbox)
                    mask_image = np.zeros((int(im_info[0, 0]), int(im_info[0, 1])))
                    mask = masks[i, :, :]
                    mask = cv2.resize(mask, (bbox[2] - bbox[0], (bbox[3] - bbox[1])), interpolation=cv2.INTER_LINEAR)
                    mask[mask > 0.5] = 200
                    mask[mask <= 0.5] = 0
                    mask_image[bbox[1]: bbox[3], bbox[0]: bbox[2]] = mask
                    cv2.imwrite(os.path.join(result_path, filename) + '_' + str(count) + '.png', mask_image)
                    f.write('{:s} {:s} {:.8f}\n'.format(filename + '_' + str(count) + '.png', str(labelID), score))
                    count += 1
            f.flush()
            f.close()

    def evaluate_detections(self, results_pack):
        """
        :param results_pack: {
            all_boxes: [class_id][image_id][box_id]
            all_masks: [class_id][image_id][box_id]
            results_list: #img {
                image:
                im_info:
                boxes:
                masks:
                }
            }
        :return:
        """
        all_boxes_and_scores = results_pack["all_boxes"]
        all_boxes = [[None for _ in range(self.num_images)] for _ in range(self.num_classes)]
        all_scores = [[None for _ in range(self.num_images)] for _ in range(self.num_classes)]
        for i in range(1, self.num_classes):
            for j in range(self.num_images):
                if all_boxes_and_scores[i][j] is not None:
                    all_boxes[i][j] = all_boxes_and_scores[i][j][:, :4]
                    all_scores[i][j] = all_boxes_and_scores[i][j][:, 4]

        gt_roidb = self.gt_roidb()
        gt_boxes = [roirec["boxes"] for roirec in gt_roidb]
        gt_classes = [roirec["gt_classes"] for roirec in gt_roidb]

        all_gt_boxes = [[None for _ in range(self.num_images)] for _ in range(self.num_classes)]
        for i in range(self.num_images):
            for cls in range(1, self.num_classes):
                all_gt_boxes[cls][i] = gt_boxes[i][np.where(gt_classes[i] == cls)[0]]

        class_names = self.classes

        # the following codes are dataset-agnostic
        # generalized evaluation of detection takes 4 inputs
        # all_boxes, all_scores, all_gt_boxes, class names(not mandatory)
        from rcnn.core.helper import single_class_ap
        threshs = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95]
        aps = np.empty(shape=(len(threshs), len(all_gt_boxes) - 1))
        for i_thrs, thresh in enumerate(threshs):
            for i_cls, cls in enumerate(range(1, len(all_gt_boxes))):
                aps[i_thrs, i_cls] = single_class_ap(all_gt_boxes[cls], all_boxes[cls], all_scores[cls], thresh)
        for i_ap, (ap50, ap) in enumerate(zip(aps[0], aps.mean(axis=0))):
            print("%-15s: AP50=%.3f, AP=%.3f" % (class_names[i_ap + 1], ap50, ap))
        print("%-15s: AP50=%.3f, AP=%.3f" % ("average", aps[0].mean(), aps.mean()))

    def evaluate_segmentation(self, results_pack):
        """
        :param results_pack: {
            all_boxes: [class_id][image_id][box_id]
            all_masks: [class_id][image_id][box_id]
            results_list: #img {
                image:
                im_info:
                boxes:
                masks:
                seg:
                }
            }
        :return:
        """
        result_path = 'eval_seg/'
        trainId2id = [7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33]

        for result in results_pack["results_list"]:
            image_path = result['image']
            print("processing %s" % image_path)
            filename = image_path.split("/")[-1]
            segmentation = result["seg"].argmax(axis=0)
            assert len(segmentation.shape) == 2, segmentation.shape
            id_segmentation = segmentation.copy()

            for train_id in range(19):
                id_segmentation[segmentation == train_id] = trainId2id[train_id]

            cv2.imwrite(result_path + filename, id_segmentation)

    def create_roidb_from_box_list(self, proposal_list, gt_roidb):
        """
        given ground truth, prepare roidb
        :param proposal_list: (#img, #proposals, xyxy)
        :param gt_roidb: (#img), each item is a roirec
        :return: roidb: [image_index]['boxes', 'gt_classes', 'gt_overlaps', 'flipped']
        """
        assert len(proposal_list) == self.num_images, 'number of boxes matrix must match number of images'
        roidb = []
        for i in range(self.num_images):
            roi_rec = dict()
            roi_rec['image'] = gt_roidb[i]['image']
            roi_rec['height'] = gt_roidb[i]['height']
            roi_rec['width'] = gt_roidb[i]['width']
            if self.panoptic:
                roi_rec['sem_seg'] = gt_roidb[i]['sem_seg']

            proposals = proposal_list[i]
            if proposals.shape[1] == 5:
                proposals = proposals[:, :4]
            num_proposals = proposals.shape[0]
            if gt_roidb is not None and gt_roidb[i]['boxes'].size > 0:
                gt_boxes = gt_roidb[i]['boxes']
                gt_classes = gt_roidb[i]['gt_classes']
                gt_instances = gt_roidb[i]['ins_id']
                # n proposals and k gt_boxes => n * k overlap
                gt_overlaps = bbox_overlaps(proposals.astype(np.float32, copy=False),
                                            gt_boxes.astype(np.float32, copy=False))
                max_overlaps = gt_overlaps.max(axis=1)
                max_overlap_box_ids = gt_overlaps.argmax(axis=1)
                max_instances = gt_instances[max_overlap_box_ids]
                max_classes = gt_classes[max_overlap_box_ids]
                max_classes[max_overlaps == 0] = 0

                roi_rec.update({'boxes': proposals,
                                'gt_classes': np.zeros((num_proposals,), dtype=np.int32),
                                'max_classes': max_classes,
                                'max_overlaps': max_overlaps,
                                'ins_seg': gt_roidb[i]['ins_seg'],
                                'ins_id': max_instances,
                                'flipped': False})
            else:
                roi_rec.update({'boxes': proposals,
                                'gt_classes': np.zeros((num_proposals,), dtype=np.int32),
                                'max_classes': np.zeros((num_proposals,), dtype=np.int32),
                                'max_overlaps': np.zeros((num_proposals,), dtype=np.float32),
                                'ins_seg': gt_roidb[i]['ins_seg'],
                                'ins_id': np.zeros((num_proposals,), dtype=np.int32),
                                'flipped': False})
            # background roi => background class
            zero_indexes = np.where(roi_rec['max_overlaps'] == 0)[0]
            assert all(roi_rec['max_classes'][zero_indexes] == 0)
            # foreground roi => foreground class
            nonzero_indexes = np.where(roi_rec['max_overlaps'] > 0)[0]
            assert all(roi_rec['max_classes'][nonzero_indexes] != 0)

            roidb.append(roi_rec)

        return roidb
