"""Transform a roidb into a trainable roidb by adding a bunch of metadata."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import pickle

import datasets
import numpy as np
from model.utils.config import cfg
from datasets.factory import get_imdb
from PIL import Image
import pdb

def prepare_roidb(imdb):
    """Enrich the imdb's roidb by adding some derived quantities that
    are useful for training. This function precomputes the maximum
    overlap, taken over ground-truth boxes, between each ROI and
    each ground-truth box. The class with maximum overlap is also
    recorded.
    """

    roidb = imdb.roidb
    # roidb handler is used and saves a dictionary of the roi
    if not "al_" in imdb.name:
        cache_file = os.path.join(imdb.cache_path, imdb.name + '_sizes.pkl')
        if os.path.exists(cache_file):
            print('Image sizes loaded from %s' % cache_file)
            with open(cache_file, 'rb') as f:
                sizes = pickle.load(f)
        else:
            print('Extracting image sizes... (It may take long time)')
            sizes = [Image.open(imdb.image_path_at(i)).size
                                for i in range(imdb.num_images)]
            with open(cache_file, 'wb') as f:
                pickle.dump(sizes, f)
            print('Done!!')
    else:
        sizes = [Image.open(imdb.image_path_at(i)).size for i in range(imdb.num_images)]
    #caches sizes of images(w x h)

    for i in range(len(imdb.image_index)):
        # i means the index of an image
        roidb[i]['image'] = imdb.image_path_at(i)
        if not (imdb.name.startswith('coco')):
            roidb[i]['width'] = sizes[i][0]
            roidb[i]['height'] = sizes[i][1]
        # need gt_overlaps as a dense array for argmax
        gt_overlaps = roidb[i]['gt_overlaps'].toarray()
        # max overlap with gt over classes (columns)
        max_overlaps = gt_overlaps.max(axis=1)
        # gt class that had the max overlap
        max_classes = gt_overlaps.argmax(axis=1)
        roidb[i]['max_classes'] = max_classes
        roidb[i]['max_overlaps'] = max_overlaps
        # sanity checks
        # max overlap of 0 => class should be zero (background)
        zero_inds = np.where(max_overlaps == 0)[0]
        assert all(max_classes[zero_inds] == 0)
        # max overlap > 0 => class should not be zero (must be a fg class)
        nonzero_inds = np.where(max_overlaps > 0)[0]
        assert all(max_classes[nonzero_inds] != 0)


def rank_roidb_ratio(roidb):
        # rank roidb based on the ratio between width and height.
        ratio_large = 2 # largest ratio to preserve.
        ratio_small = 0.5 # smallest ratio to preserve.        
        
        ratio_list = []
        for i in range(len(roidb)):
            width = roidb[i]['width']
            height = roidb[i]['height']
            ratio = width / float(height)

            if ratio > ratio_large:
                roidb[i]['need_crop'] = 1
                ratio = ratio_large
            elif ratio < ratio_small:
                roidb[i]['need_crop'] = 1
                ratio = ratio_small                
            else:
                roidb[i]['need_crop'] = 0

            ratio_list.append(ratio)

        ratio_list = np.array(ratio_list)
        ratio_index = np.argsort(ratio_list)
        return ratio_list[ratio_index], ratio_index

def filter_roidb(roidb):
        # filter the image without bounding box.
        print('before filtering, there are %d images...' % (len(roidb)))
        i = 0
        while i < len(roidb):
            if len(roidb[i]['boxes']) == 0:
                del roidb[i]
                i -= 1
            i += 1

        print('after filtering, there are %d images...' % (len(roidb)))
        return roidb

def combined_roidb(imdb_names, training=True):
    """
    Combine multiple roidbs
    """

    def get_training_roidb(imdb):
        """Returns a roidb (Region of Interest database) for use in training."""
        if cfg.TRAIN.USE_FLIPPED:
            print('Appending horizontally-flipped training examples...')
            imdb.append_flipped_images()
            print('done')

        print('Preparing training data...')

        prepare_roidb(imdb)
        #ratio_index = rank_roidb_ratio(imdb)
        print('done')

        return imdb.roidb
    
    def get_roidb(imdb_name):
        imdb = get_imdb(imdb_name)
        # imdb is the database instance, located @factory.py
        print('Loaded dataset `{:s}`'.format(imdb.name))
        imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
        # set roidb_handler to gt_roidb
        print('Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD))
        roidb = get_training_roidb(imdb)
        return roidb

    roidbs = [get_roidb(s) for s in imdb_names.split('+')]
    roidb = roidbs[0]
    # controversial

    if len(roidbs) > 1:
        # roidb with many datasets, might not work
        for r in roidbs[1:]:
            roidb.extend(r)
        tmp = get_imdb(imdb_names.split('+')[1])
        imdb = datasets.imdb.imdb(imdb_names, tmp.classes)
    else:
        imdb = get_imdb(imdb_names)

    # if training:
    #     roidb = filter_roidb(roidb)
    # we use ws_train, so we should not filter out roidbs without bboxes

    ratio_list, ratio_index = rank_roidb_ratio(roidb)

    return imdb, roidb, ratio_list, ratio_index
    # ratio_index is the order of original ratio_list
    # ratio_list is the sorted ration_list
    # imdb is the class instance of an imagenet class
    # roidb is a list of dictionary of rois
