"""Interval and IntervalSet

Represents an interval of time, and a set of such intervals."""

# First implementation kept a sorted list of intervals and used
# biesct() to optimize some operations, but this was too slow.

# This version is based on the quicksect implementation from python-bx,
# modified slightly to handle floating point intervals.

import pyximport
pyximport.install()
import bxintersect

import bisect

class IntervalError(Exception):
    """Error due to interval overlap, etc"""
    pass

class Interval(bxintersect.Interval):
    """Represents an interval of time."""

    def __init__(self, start, end):
        """
        'start' and 'end' are arbitrary floats that represent time
        """
        if start > end:
            raise IntervalError("start must precede end")
        bxintersect.Interval.__init__(self, start, end)

    def __repr__(self):
        s = repr(self.start) + ", " + repr(self.end)
        return self.__class__.__name__ + "(" + s + ")"

    def __str__(self):
        return "[" + str(self.start) + " -> " + str(self.end) + "]"

    def intersects(self, other):
        """Return True if two Interval objects intersect"""
        if (self.end <= other.start or self.start >= other.end):
            return False
        return True

    def subset(self, start, end):
        """Return a new Interval that is a subset of this one"""
        # A subclass that tracks additional data might override this.
        if start < (self.start) or end > self.end:
            raise IntervalError("not a subset")
        return Interval(start, end)

class IntervalSet(object):
    """
    A non-intersecting set of intervals.
    """

    def __init__(self, source=None):
        """
        'source' is an Interval or IntervalSet to add.
        """
        self.tree = bxintersect.IntervalTree()
        if source is not None:
            self += source

    def __iter__(self):
        return self.tree.traverse()

    def __len__(self):
        return sum(1 for x in self)

    def __repr__(self):
        descs = [ repr(x) for x in self ]
        return self.__class__.__name__ + "([" + ", ".join(descs) + "])"

    def __str__(self):
        descs = [ str(x) for x in self ]
        return  "[" + ", ".join(descs) + "]"

    def __eq__(self, other):
        # This isn't particularly efficient, but it shouldn't get used in the
        # general case.
        """Test equality of two IntervalSets.

        Treats adjacent Intervals as equivalent to one long interval,
        so this function really tests whether the IntervalSets cover
        the same spans of time."""
        i = 0
        j = 0
        outside = True

        def is_adjacent(a, b):
            """Return True if two Intervals are adjacent (same end or start)"""
            if a.end == b.start or b.end == a.start:
                return True
            else:
                return False

        this = [ x for x in self ]
        that = [ x for x in other ]

        try:
            while True:
                if (outside):
                    # To match, we need to be finished both sets
                    if (i >= len(this) and j >= len(that)):
                        return True
                    # Or the starts need to match
                    if (this[i].start != that[j].start):
                        return False
                    outside = False
                else:
                    # We can move on if the two interval ends match
                    if (this[i].end == that[j].end):
                        i += 1
                        j += 1
                        outside = True
                    else:
                        # Whichever ends first needs to be adjacent to the next
                        if (this[i].end < that[j].end):
                            if (not is_adjacent(this[i],this[i+1])):
                                return False
                            i += 1
                        else:
                            if (not is_adjacent(that[j],that[j+1])):
                                return False
                            j += 1
        except IndexError:
            return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __iadd__(self, other):
        """Inplace add -- modifies self

        This throws an exception if the regions being added intersect."""
        if isinstance(other, Interval):
            if self.intersects(other):
                raise IntervalError("Tried to add overlapping interval "
                                    "to this set")
            self.tree.insert_interval(other)
        else:
            for x in other:
                self.__iadd__(x)
        return self

    def __add__(self, other):
        """Add -- returns a new object"""
        new = IntervalSet(self)
        new += IntervalSet(other)
        return new

    def __and__(self, other):
        """
        Compute a new IntervalSet from the intersection of two others

        Output intervals are built as subsets of the intervals in the
        first argument (self).
        """
        out = IntervalSet()

        if not isinstance(other, IntervalSet):
            other = [ other ]

        for x in other:
            # Intersecting with a just a single interval.
            all_intersecting = self.tree.find(x.start, x.end)
            for i in all_intersecting:
                out += i.subset(max(i.start, x.start),
                                min(i.end, x.end))
        return out

    def intersects(self, other):
        """Return True if this IntervalSet intersects another"""
        if isinstance(other, Interval):
            return len(self.tree.find(other.start, other.end)) > 0
        else:
            for x in other:
                if self.intersects(x):
                    return True
            else:
                return False

