# cython: profile=False
# cython: cdivision=True

"""
Jim Paris <jim@jtan.com>

Red-black tree, where keys are stored as start/end timestamps.
This is a basic interval tree that holds half-open intervals:
  [start, end)
Intervals must not overlap.  Fixing that would involve making this
into an augmented interval tree as described in CLRS 14.3.

Code that assumes non-overlapping intervals is marked with the
string 'non-overlapping'.
"""

import sys

cdef class RBNode:
    """One node of the Red/Black tree, containing a key (start, end)
    and value (obj)"""
    cdef public object obj
    cdef public double start, end
    cdef public int red
    cdef public RBNode left, right, parent

    def __cinit__(RBNode self, double start, double end, object obj = None):
        self.obj = obj
        self.start = start
        self.end = end
        self.red = False
        self.left = None
        self.right = None

    def __str__(self):
        if self.red:
            color = "R"
        else:
            color = "B"
        if self.start == sys.float_info.min:
            return "[node nil]"
        return ("[node ("
                + str(self.obj) + ") "
                + str(self.start) + " -> " + str(self.end) + " "
                + color + "]")

cdef class RBTree:
    """Red/Black tree"""

    cdef public RBNode nil, root

    # Init
    def __cinit__(RBTree self):
        self.nil = RBNode(start = sys.float_info.min,
                          end = sys.float_info.min)
        self.nil.left = self.nil
        self.nil.right = self.nil
        self.nil.parent = self.nil

        self.root = RBNode(start = sys.float_info.max,
                           end = sys.float_info.max)
        self.root.left = self.nil
        self.root.right = self.nil
        self.root.parent = self.nil

    # We have a dummy root node to simplify operations, so from an
    # external point of view, its left child is the real root.
    cpdef getroot(RBTree self):
        return self.root.left

    # Rotations and basic operations
    cdef void __rotate_left(RBTree self, RBNode x):
        """Rotate left:
        #   x           y
        #  / \   -->   / \
        # z   y       x   w
        #    / \     / \
        #   v   w   z   v
        """
        cdef RBNode y = x.right
        x.right = y.left
        if y.left is not self.nil:
            y.left.parent = x
        y.parent = x.parent
        if x is x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    cdef void __rotate_right(RBTree self, RBNode y):
        """Rotate right:
        #     y           x
        #    / \   -->   / \
        #   x   w       z   y
        #  / \             / \
        # z   v           v   w
        """
        cdef RBNode x = y.left
        y.left = x.right
        if x.right is not self.nil:
            x.right.parent = y
        x.parent = y.parent
        if y is y.parent.left:
            y.parent.left = x
        else:
            y.parent.right = x
        x.right = y
        y.parent = x

    cdef RBNode __successor(RBTree self, RBNode x):
        """Returns the successor of RBNode x"""
        cdef RBNode y = x.right
        if y is not self.nil:
            while y.left is not self.nil:
                y = y.left
        else:
            y = x.parent
            while x is y.right:
                x = y
                y = y.parent
            if y is self.root:
                return self.nil
        return y
    cpdef RBNode successor(RBTree self, RBNode x):
        """Returns the successor of RBNode x, or None"""
        cdef RBNode y = self.__successor(x)
        return y if y is not self.nil else None

    cdef RBNode __predecessor(RBTree self, RBNode x):
        """Returns the predecessor of RBNode x"""
        cdef RBNode y = x.left
        if y is not self.nil:
            while y.right is not self.nil:
                y = y.right
        else:
            y = x.parent
            while x is y.left:
                if y is self.root:
                    y = self.nil
                    break
                x = y
                y = y.parent
        return y
    cpdef RBNode predecessor(RBTree self, RBNode x):
        """Returns the predecessor of RBNode x, or None"""
        cdef RBNode y = self.__predecessor(x)
        return y if y is not self.nil else None

    # Insertion
    cpdef insert(RBTree self, RBNode z):
        """Insert RBNode z into RBTree and rebalance as necessary"""
        z.left = self.nil
        z.right = self.nil
        cdef RBNode y = self.root
        cdef RBNode x = self.root.left
        while x is not self.nil:
            y = x
            if (x.start > z.start or (x.start == z.start and x.end > z.end)):
                x = x.left
            else:
                x = x.right
        z.parent = y
        if (y is self.root or
            (y.start > z.start or (y.start == z.start and y.end > z.end))):
            y.left = z
        else:
            y.right = z
        # relabel/rebalance
        self.__insert_fixup(z)

    cdef void __insert_fixup(RBTree self, RBNode x):
        """Rebalance/fix RBTree after a simple insertion of RBNode x"""
        x.red = True
        while x.parent.red:
            if x.parent is x.parent.parent.left:
                y = x.parent.parent.right
                if y.red:
                    x.parent.red = False
                    y.red = False
                    x.parent.parent.red = True
                    x = x.parent.parent
                else:
                    if x is x.parent.right:
                        x = x.parent
                        self.__rotate_left(x)
                    x.parent.red = False
                    x.parent.parent.red = True
                    self.__rotate_right(x.parent.parent)
            else: # same as above, left/right switched
                y = x.parent.parent.left
                if y.red:
                    x.parent.red = False
                    y.red = False
                    x.parent.parent.red = True
                    x = x.parent.parent
                else:
                    if x is x.parent.left:
                        x = x.parent
                        self.__rotate_right(x)
                    x.parent.red = False
                    x.parent.parent.red = True
                    self.__rotate_left(x.parent.parent)
        self.root.left.red = False

    # Deletion
    cpdef delete(RBTree self, RBNode z):
        if z.left is None or z.right is None:
            raise AttributeError("you can only delete a node object "
                                 + "from the tree; use find() to get one")
        cdef RBNode x, y
        if z.left is self.nil or z.right is self.nil:
            y = z
        else:
            y = self.__successor(z)
        if y.left is self.nil:
            x = y.right
        else:
            x = y.left
        x.parent = y.parent
        if x.parent is self.root:
            self.root.left = x
        else:
            if y is y.parent.left:
                y.parent.left = x
            else:
                y.parent.right = x
        if y is not z:
            # y is the node to splice out, x is its child
            y.left = z.left
            y.right = z.right
            y.parent = z.parent
            z.left.parent = y
            z.right.parent = y
            if z is z.parent.left:
                z.parent.left = y
            else:
                z.parent.right = y
            if not y.red:
                y.red = z.red
                self.__delete_fixup(x)
            else:
                y.red = z.red
        else:
            if not y.red:
                self.__delete_fixup(x)

    cdef void __delete_fixup(RBTree self, RBNode x):
        """Rebalance/fix RBTree after a deletion.  RBNode x is the
        child of the spliced out node."""
        cdef RBNode rootLeft = self.root.left
        while not x.red and x is not rootLeft:
            if x is x.parent.left:
                w = x.parent.right
                if w.red:
                    w.red = False
                    x.parent.red = True
                    self.__rotate_left(x.parent)
                    w = x.parent.right
                if not w.right.red and not w.left.red:
                    w.red = True
                    x = x.parent
                else:
                    if not w.right.red:
                        w.left.red = False
                        w.red = True
                        self.__rotate_right(w)
                        w = x.parent.right
                    w.red = x.parent.red
                    x.parent.red = False
                    w.right.red = False
                    self.__rotate_left(x.parent)
                    x = rootLeft # exit loop
            else: # same as above, left/right switched
                w = x.parent.left
                if w.red:
                    w.red = False
                    x.parent.red = True
                    self.__rotate_right(x.parent)
                    w = x.parent.left
                if not w.left.red and not w.right.red:
                    w.red = True
                    x = x.parent
                else:
                    if not w.left.red:
                        w.right.red = False
                        w.red = True
                        self.__rotate_left(w)
                        w = x.parent.left
                    w.red = x.parent.red
                    x.parent.red = False
                    w.left.red = False
                    self.__rotate_right(x.parent)
                    x = rootLeft # exit loop
        x.red = False

    # Walking, searching
    def __iter__(RBTree self):
        return self.inorder()

    def inorder(RBTree self, RBNode x = None):
        """Generator that performs an inorder walk for the tree
        rooted at RBNode x"""
        if x is None:
            x = self.getroot()
        while x.left is not self.nil:
            x = x.left
        while x is not self.nil:
            yield x
            x = self.__successor(x)

    cpdef RBNode find(RBTree self, double start, double end):
        """Return the node with exactly the given start and end."""
        cdef RBNode x = self.getroot()
        while x is not self.nil:
            if start < x.start:
                x = x.left
            elif start == x.start:
                if end == x.end:
                    break # found it
                elif end < x.end:
                    x = x.left
                else:
                    x = x.right
            else:
                x = x.right
        return x if x is not self.nil else None

    cpdef RBNode find_left_end(RBTree self, double t):
        """Find the leftmode node with end >= t.  With non-overlapping
        intervals, this is the first node that might overlap time t.

        Note that this relies on non-overlapping intervals, since
        it assumes that we can use the endpoints to traverse the
        tree even though it was created using the start points."""
        cdef RBNode x = self.getroot()
        while x is not self.nil:
            if t < x.end:
                if x.left is self.nil:
                    break
                x = x.left
            elif t == x.end:
                break
            else:
                if x.right is self.nil:
                    x = self.__successor(x)
                    break
                x = x.right
        return x if x is not self.nil else None

    cpdef RBNode find_right_start(RBTree self, double t):
        """Find the rightmode node with start <= t.  With non-overlapping
        intervals, this is the last node that might overlap time t."""
        cdef RBNode x = self.getroot()
        while x is not self.nil:
            if t < x.start:
                if x.left is self.nil:
                    x = self.__predecessor(x)
                    break
                x = x.left
            elif t == x.start:
                break
            else:
                if x.right is self.nil:
                    break
                x = x.right
        return x if x is not self.nil else None

    # Intersections
    def intersect(RBTree self, double start, double end):
        """Generator that returns nodes that overlap the given
        (start,end) range.  Assumes non-overlapping intervals."""
        # Start with the leftmode node that ends after start
        cdef RBNode n = self.find_left_end(start)
        while n is not None:
            if n.start >= end:
                # this node starts after the requested end; we're done
                break
            if start < n.end:
                # this node overlaps our requested area
                yield n
            n = self.successor(n)
