# --------------------------------------------------------
# Deformable Convolutional Networks
# Copyright (c) 2017 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Zheng Zhang
# --------------------------------------------------------

import cPickle
import os
import cv2
import numpy as np
import itertools

from imdb import IMDB
from PIL import Image

class CityScape(IMDB):
    def __init__(self, image_set, root_path, dataset_path, result_path=None):
        """
        fill basic information to initialize imdb
        :param image_set: leftImg8bit_train, etc
        :param root_path: 'selective_search_data' and 'cache'
        :param dataset_path: data and results
        :return: imdb object
        """
        image_set_main_folder, image_set_sub_folder= image_set.split('_', 1)
        super(CityScape, self).__init__('cityscape', image_set, root_path, dataset_path, result_path)  # set self.name

        self.image_set_main_folder = image_set_main_folder
        self.image_set_sub_folder = image_set_sub_folder
        self.root_path = root_path
        self.data_path = dataset_path
        self.num_classes = 19
        self.image_set_index = self.load_image_set_index()
        self.num_images = len(self.image_set_index)
        print 'num_images', self.num_images

        self.config = {'comp_id': 'comp4',
                       'use_diff': False,
                       'min_size': 2}


    def load_image_set_index(self):
        """
        find out which indexes correspond to given image set
        :return: the indexes of given image set
        """

        #Collection all subfolders
        image_set_main_folder_path = os.path.join(self.data_path, self.image_set_main_folder, self.image_set_sub_folder)
        image_name_set = [filename for parent, dirname, filename in os.walk(image_set_main_folder_path)]
        image_name_set = list(itertools.chain.from_iterable(image_name_set))
        index_set = ['' for x in range(len(image_name_set))]
        valid_index_count = 0
        for i, image_name in enumerate(image_name_set):
            splited_name_set = image_name.split('_')
            ext_split = splited_name_set[len(splited_name_set) - 1].split('.')
            ext = ext_split[len(ext_split)-1]
            if splited_name_set[len(splited_name_set) - 1] != 'flip.png' and ext == 'png':
                index_set[valid_index_count] = "_".join(splited_name_set[:len(splited_name_set)-1])
                valid_index_count += 1

        return index_set[:valid_index_count]

    def image_path_from_index(self, index):
        """
        find the image path from given index
        :param index: the given index
        :return: the image path
        """
        index_folder = index.split('_')[0]
        image_file = os.path.join(self.data_path, self.image_set_main_folder, self.image_set_sub_folder, index_folder, index + '_' + self.image_set_main_folder + '.png')
        assert os.path.exists(image_file), 'Path does not exist: {}'.format(image_file)
        return image_file

    def annotation_path_from_index(self, index):
        """
        find the gt path from given index
        :param index: the given index
        :return: the image path
        """
        index_folder = index.split('_')[0]
        image_file = os.path.join(self.data_path, 'gtFine', self.image_set_sub_folder, index_folder, index + '_gtFine_labelTrainIds.png')
        assert os.path.exists(image_file), 'Path does not exist: {}'.format(image_file)
        return image_file

    def load_segdb_from_index(self, index):
        """
        load segdb from given index
        :param index: given index
        :return: segdb
        """
        seg_rec = dict()
        seg_rec['image'] = self.image_path_from_index(index)
        size = cv2.imread(seg_rec['image']).shape
        seg_rec['height'] = size[0]
        seg_rec['width'] = size[1]

        seg_rec['seg_cls_path'] = self.annotation_path_from_index(index)
        seg_rec['flipped'] = False

        return seg_rec

    def gt_segdb(self):
        """
        return ground truth image regions database
        :return: imdb[image_index]['', 'flipped']
        """
        cache_file = os.path.join(self.cache_path, self.name + '_gt_segdb.pkl')
        if os.path.exists(cache_file):
            with open(cache_file, 'rb') as fid:
                segdb = cPickle.load(fid)
            print '{} gt segdb loaded from {}'.format(self.name, cache_file)
            return segdb

        gt_segdb = [self.load_segdb_from_index(index) for index in self.image_set_index]