from sets import Set as set
set.update = set.union_update

class AnnotationsException(Exception):
    pass

class Annotations:
    def __init__(self, *dimensions):
        """ Initializes a new Annotation with the given dimension labels.
        dimensions is a list of dimension labels.
        """ 
        if len(dimensions) != 2:
            msg = 'Annotations only supports two dimensions.'
            raise AnnotationsException(msg)

        self.dimensions = {}
        for d in dimensions:
            self.dimensions[d] = {}

    def add_annotations(self, dim, id, ann_dim, annotations):
        """ Adds new annotations.
        dim: the dimension in which the new data should be added.
        id: the identifier that should be annotated.
        ann_dim: the dimension of the annotations to id.
        annotations: the new annotations to id.
        Returns the total set of annotations to id.
        """
        if not self.has_dimension(dim):
            msg = 'Annotations object does not contain dimension %s' % dim
            raise AnnotationsException(msg)

        if not self.has_dimension(ann_dim):
            msg = 'Annotations object does not contain dimension %s' % ann_dim
            raise AnnotationsException(msg)

        for a in annotations:
            if not self.dimensions[ann_dim].has_key(a):
                self.dimensions[ann_dim][a] = set()
            self.dimensions[ann_dim][a].add(id)

        if not self.dimensions[dim].has_key(id):
            self.dimensions[dim][id] = set()
        self.dimensions[dim][id].update(annotations)

        return self.dimensions[dim][id]
        
    def get_annotations(self, dim, id, ann_dim):
        """Returns all annotations to id.
        dim: the dimension where id can be found.
        id: the id to retrieve annotations for.
        """
        if not self.has_dimension(dim):
            msg = 'Annotations object does not contain dimension %s' % dim
            raise AnnotationsException(msg)

        if self.dimensions[dim].has_key(id):
            return self.dimensions[dim][id]
        return set()

    def has_dimension(self, dim):
        """ Retuns true if the Annotations object indexes dim.
        """
        return self.dimensions.has_key(dim)