# API for Google Refexp.

from __future__ import print_function
from pycocotools.coco import COCO
import cPickle as pkl
import os.path
import json
import warnings
from common_utils import draw_bbox


class Refexp:
    def __init__(self, dataset_filename, coco_filename, cache_filename=None):
        """
        Constructor for Refexp class for reading and visualizing referent expressions dataset.
        :param dataset_filename  (str): location of the dataset file
        :param coco_filename     (str): location of the COCO instances JSON file (e.g. `instances_train2014.json`)
        :param cache_filename    (str): if specified, writes/reads to this cache file. Should be specific to each (`dataset_filename`, `coco_filename`).
        :return:
        """
        
        # Load from cache:
        if cache_filename is not None and os.path.isfile(cache_filename):
            print('Loading from cache...')
            with open(cache_filename, 'r') as f:
                obj = pkl.load(f)
            self.coco = obj['coco']
            self.dataset = obj['dataset']
            self.imgIds = obj['imgIds']
            self.annIds = obj['annIds']
            self.refexpIds = obj['refexpIds']
            self.refexpToAnnId = obj['refexpToAnnId']
            self.annToImgId = obj['annToImgId']
            self.catIds = obj['catIds']
        else:
            # Instantiate COCO class to use its methods.
            self.coco = COCO(coco_filename)
            
            # Load dataset form the _aligned_coco.json files generated by setup.py.
            with open(dataset_filename, 'r') as f:
                self.dataset = json.load(f)
            # Convert string keys to integers for images, annotations, refexps.
            def _convert(dataset_key):
                converted = {}
                for k, v in self.dataset[dataset_key].iteritems():
                    converted[int(k)] = v
                self.dataset[dataset_key] = converted
            _convert('images')
            _convert('annotations')
            _convert('refexps')
            print('Dataset loaded.')

            # List of all image ids in the dataset.
            self.imgIds = self.dataset['images'].keys()
            
            # List of all annotation ids in the dataset.
            self.annIds = self.dataset['annotations'].keys()
            
            # List of all referring expression ids in the dataset.
            self.refexpIds = self.dataset['refexps'].keys()
            
            # Map from referring expression id to corresponding annotation id (1:1).
            refexpToAnnId = {}
            # Map annotation id to image id (1:1).
            annToImgId = {}
            # Set of all category ids in the dataset.
            catIds = set([])
            for ann_id, ann in self.dataset['annotations'].iteritems():
                if ann['category_id'] not in catIds:
                    catIds.add(ann['category_id'])
                # Check 1:1 relation.
                assert ann_id not in annToImgId 
                annToImgId[ann_id] = ann['image_id']
                for ref_id in ann['refexp_ids']:
                    # Check 1:1 relation.
                    assert ref_id not in refexpToAnnId
                    refexpToAnnId[ref_id] = ann_id
                    
            self.refexpToAnnId = refexpToAnnId
            self.annToImgId = annToImgId
            self.catIds = list(catIds)
            
            # And save to cache if cache_filename given:
            if cache_filename is not None:
                print('Writing to cache...')
                obj = {}
                obj['coco'] = self.coco
                obj['dataset'] = self.dataset
                obj['imgIds'] = self.imgIds
                obj['annIds'] = self.annIds
                obj['refexpIds'] = self.refexpIds
                obj['refexpToAnnId'] = self.refexpToAnnId
                obj['annToImgId'] = self.annToImgId
                obj['catIds'] = self.catIds
                with open(cache_filename, 'w') as f:
                    pkl.dump(obj, f)
        
    # Helper functions for the rest of the COCO get functions.
    def _filterImgIds(self, imgIds):
        """
        Sets a list of image ids to the intersection of itself with the list of GoogleRefexp image ids.
        Raises warning if the list contains images outside GoogleRefexp dataset.
        """
        imgIds = imgIds if isinstance(imgIds, list) else [imgIds]
        if len(imgIds) == 0:
            return self.imgIds
        imgIdsSet = set(imgIds)
        refexpImgsSet = set(self.imgIds)
        notRefexpImages = imgIdsSet - refexpImgsSet
        if len(notRefexpImages) > 0:
            warnings.warn('Images ' + str(notRefexpImages) + ' are not part of the GoogleRefexp dataset and will be ignored from the answer.', RuntimeWarning)
        imgIds = imgIdsSet.intersection(refexpImgsSet)
        return list(imgIds)
    
    def _filterAnnIds(self, annIds, warn=True):
        """
        Sets a list of annotation ids to the intersection of itself with the list of GoogleRefexp annotation ids.
        Raises warning if the input annotation id list contains ids outside GoogleRefexp dataset.
        """
        annIds = annIds if isinstance(annIds, list) else [annIds]
        if len(annIds) == 0:
            return self.annIds
        annIdsSet = set(annIds)
        refexpAnnsSet = set(self.annIds)
        notRefexpAnns = annIdsSet - refexpAnnsSet
        if len(notRefexpAnns) > 0 and warn:
            warnings.warn('Annotations ' + str(notRefexpAnns) + ' are not part of the GoogleRefexp dataset and will be ignored from the answer.', RuntimeWarning)
        
        return list(annIdsSet.intersection(refexpAnnsSet))
    
    def _filterCatIds(self, catIds):
        """
        Sets a list of categories ids to the intersection of itself with the list of GoogleRefexp category ids.
        Raises warning if the category id list contains ids outside GoogleRefexp dataset.
        """
        catIds = catIds if isinstance(catIds, list) else [catIds]
        if len(catIds) == 0:
            return self.catIds
        catIdsSet = set(catIds)
        refexpCatsSet = set(self.catIds)
        notRefexpCats = catIdsSet - refexpCatsSet
        if len(notRefexpCats) > 0:
            warnings.warn('Category ids ' + str(notRefexpCats) + ' are not part of the GoogleRefexp dataset and will be ignored from the answer.', RuntimeWarning)
        catIds = catIdsSet.intersection(refexpCatsSet)
        return list(catIds)
    
    # Below wrappers of the COCO toolkit methods, adding referring expression ids as additional filter. 
    # Filters default to the set of GoogleRefexp images, annotations and categories(subset of COCO).
    def getAnnIds(self, imgIds=[], catIds=[], refexpIds=[], areaRng=[], iscrowd=None):
        """
        Same as COCO.getAnnIds(). Gets only annotations that appear in GoogleRefexp.
        Adds filtering by referring expression ids.
        """
        refexpIds = refexpIds if isinstance(refexpIds, list) else [refexpIds]
        imgIds = self._filterImgIds(imgIds)
        if len(refexpIds) == 0:
            return self._filterAnnIds(self.coco.getAnnIds(imgIds, catIds, areaRng, iscrowd), warn=False)
        annsForRefexps = set([])
        for refexp_id in refexpIds:
            annsForRefexps.add(self.refexpToAnnId[refexp_id])
        return list(annsForRefexps.intersection(set(self.coco.getAnnIds(imgIds, catIds, areaRng, iscrowd))))
    
    def getCatIds(self, catNms=[], supNms=[], catIds=[]):
        """
        Same as COCO.getAnnIds(). Gets only categories that appear in GoogleRefexp.
        """
        catIds = self._filterCatIds(catIds)
        return self.coco.getCatIds(catNms, supNms, catIds)
        
    def getImgIds(self, imgIds=[], catIds=[], refexpIds=[]):
        """
        Same as COCO.getImgIds(). Gets only images that appear in GoogleRefexp.
        Adds filtering by referring expression ids.
        """
        refexpIds = refexpIds if isinstance(refexpIds, list) else [refexpIds]
        imgIds = self._filterImgIds(imgIds)
        if len(refexpIds) == 0:
            return self.coco.getImgIds(imgIds, catIds)
        imgsForRefexps = set([])
        for refexp_id in refexpIds:
            annForRefexp = self.refexpToAnnId[refexp_id]
            imgsForRefexps.add(self.annToImgId[annForRefexp])
        cocoImgs = set(self.coco.getImgIds(imgIds, catIds))
        return list(cocoImgs.intersection(imgsForRefexps))
    
    def loadAnns(self, ids=[]):
        """
        Load annotations with the specified ids.
        :param   ids  (int array)    : integer ids specifying anns
        :return: anns (object array) : loaded ann objects
        """
        ids = self._filterAnnIds(ids)
        annotations = self.dataset['annotations']
        if type(ids) == list:
            return [annotations[id] for id in ids]
        elif type(ids) == int:
            return [annotations[ids]]
    
    def loadCats(self, ids=[]):
        """
        Load categories with the specified ids.
        :param   ids        (int array)    : integer ids specifying categories
        :return: categories (object array) : loaded ann objects
        """
        ids = self._filterCatIds(ids)
        return self.coco.loadCats(ids)
    
    def loadImgs(self, ids=[]):
        """
        Load images with the specified ids.
        :param   ids        (int array)    : integer ids specifying images
        :return: images     (object array) : loaded images objects
        """
        ids = ids if isinstance(ids, list) else [ids]
        ids = self._filterImgIds(ids)
        if len(ids) == 0:
            return []
        images = self.dataset['images']
        return [images[id] for id in ids]

    
    def loadRefexps(self, ids=[]):
        """
        Load referring expressions with the specified ids.
        :param   ids        (int array)    : integer ids specifying referring expressions
        :return: images     (object array) : loaded referring expressions objects
        """
        refexps = self.dataset['refexps']
        if type(ids) == list:
            return [refexps[id] for id in ids]
        elif type(ids) == int:
            return [refexps[ids]]
    
    def download(self, tarDir=None, imgIds=[]):
        """
        Wrapper for COCO.download(). Defaults to GoogleRefexp images only.
        """
        imgIds = self._filterImgIds(imgIds)
        return self.coco.download(tarDir, imgIds)
    
    def showAnn(self, ann, ax=None, printRefexps=True):
        """
        Display the bbox and referring expressions of the given annotation.
        For segmentation display refer to COCO toolkit.
        :param anns (array of object): annotations to display
        :return: None
        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        bbox = ann['bbox']
        draw_bbox(ax, bbox, edge_color='green')
        if printRefexps:
            print('Referring expressions for the object in the bounding box: ')
            for ref_id in ann['refexp_ids']:
                print(self.dataset['refexps'][ref_id]['raw'])
        
    def showRegionCandidates(self, image, ax=None):
        """
        Display the region candidates for the given image.
        :param imgId (int): image id for which we display region candidates
        :param ax: matplotlib axes to draw on, or if unspecified, the current axis from `plt.gca()`
        :return: None
        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()

        candidates = image['region_candidates']
        for candidate in candidates:
            bbox = candidate['bounding_box']
            draw_bbox(ax, bbox, edge_color='blue')
    
    # Methods specific to Refexp.
    def getRefexpIds(self, tokens=[], len_min=0, len_max=-1, referent=None, referent_has_attributes=False):
        """
        Get description ids that contain all the tokens in the given list, 
        legnth in the specified range and in which the referent has attributes.
        :param:  contains_tokens            (str array) : get refexps containing all given tokens. Defaults to all referring expressions.
        :param:  len_min                    (int)       : minimum number of tokens in the referring expression
        :param:  len_max                    (int)       : maximum number of tokens in the referring expression
        :param:  referent                   (str)       : the referent word in the referring expression
        :param:  referent_has_attributes   (bool)      : if True return referring expressions in which the referent has attributes
        :return: ids                        (int array) : integer array of refexp 
        """
        tokens = tokens if isinstance(tokens, list) else [tokens]
        tokensSet = set(tokens)
        refexps_ids = set([])
        for ref_id, ref in self.dataset['refexps'].iteritems():
            if (set(ref['tokens']).issuperset(tokensSet) 
                    and len(ref['tokens']) >= len_min
                    and (len_max == -1 or len(ref['tokens']) <= len_max)
                    and (referent == None or ('word' in ref['parse']['referent'] and ref['parse']['referent']['word'] == referent))
                    and (not referent_has_attributes or 
                        ('amods' in ref['parse']['referent'] and len(ref['parse']['referent']['amods']) > 0))):
                refexps_ids.add(ref_id)
        return list(refexps_ids)
