# -*- Mode: Python -*-

include "python.pxi"

cdef extern from "stdint.h":
    ctypedef int int32_t
    ctypedef unsigned int uint32_t

# abstract class
cdef class avl_node:
    cdef avl_node left, right, parent
    cdef uint32_t rank_and_balance

    cdef uint32_t get_rank (self):
        return (self.rank_and_balance >> 2)

    cdef int32_t get_balance (self):
        return (self.rank_and_balance & 3) - 1

    cdef set_rank (self, uint32_t rank):
        self.rank_and_balance = (self.rank_and_balance & 3) | (rank << 2)
        
    cdef set_balance (self, int32_t balance):
        self.rank_and_balance = (self.rank_and_balance & (~3)) | (balance + 1)

    cdef set_rank_and_balance (self, int32_t balance, uint32_t rank):
        self.rank_and_balance = ((rank << 2) | ((balance + 1) & 3))

cdef class object_key_node (avl_node):

    cdef public object key

    cdef int compare (self, other):
        return PyObject_Compare (self, other)

    def __repr__ (self):
        return '<node key=%r at %x>' % (self.key, <unsigned int> self)

    def __init__ (self, key, left=None, right=None):
        self.left = None
        self.right = None
        self.rank_and_balance = 1<<2

# based directly on Knuth, Sorting and Searching, Section 6.2.3

# I prefer beautiful recursive algorithms, but in
# the interest of speed & efficiency I'll take advantage
# of code designed to run on a circa 1968 machine.
#
# The base algorithm is modified to include the use
# of the RANK field as described by Knuth, in order to
# use the tree as an ordered list data structure, where
# elements can be referenced by index.
#
# The key comparisons were also changed to permit duplicate
# keys. duplicates are inserted and deleted at the lowest
# possible position in the tree
#
# Note: I think that the <RANK> field could be used to simplify
# the rebalancing code

class Unbalanced (Exception):
    'Unbalanced AVL Tree Error'

class Invalid_Balance (Exception):
    'Invalid Balance attribute Error'

class Invalid_Rank (Exception):
    'Invalid RANK attribute Error'

class Invalid_Parent (Exception):
    'Invalid PARENT attribute Error'

cdef class avl_tree:

    cdef uint32_t height, length
    cdef object node_class
    cdef avl_node tree
    cdef readonly object compare

    def __init__ (self, node_class, compare):
        # following Knuth, there's a special 'fake'
        # node at the top.  the real tree's root is
        # in self.tree.right.
        self.tree = avl_node (0)
        self.height = 0
        self.length = 0
        self.node_class = node_class
        self.compare = compare
        
    def insert (self, key):
        cdef int32_t ii, a
        cdef avl_node t, s, p, q, r
        if not self.tree.right:
            self.tree.right = self.node_class (key, self.tree)
        else:
            t = self.tree
            s = p = t.right
            ii = 0
            while 1:
                if self.compare (key, p.key) < 0:
                    # move left
                    p.set_rank (p.get_rank() + 1)
                    q = p.left
                    if not q:
                        # insert
                        q = self.node_class (key, p)
                        p.left = q
                        break
                    elif q.balance:
                        t = p
                        s = q
                    p = q
                else:   # (key > p.key)
                    # move right
                    q = p.right
                    ii = ii + p.get_rank()
                    if not q:
                        # insert
                        q = self.node_class (key, p)
                        p.right = q
                        break
                    elif q.get_balance():
                        t = p
                        s = q
                    p = q
            self.length = self.length + 1
            # adjust balance factors
            if self.compare (key, s.key) < 0:
                r = p = s.left
            else:
                r = p = s.right
            while p is not q:
                if self.compare (key, p.key) < 0:
                    p.set_balance(-1)
                    p = p.left
                else: #  key >= p.key:
                    p.set_balance(1)
                    p = p.right
            # balancing act
            if self.compare (key, s.key) < 0:
                a = -1
            else:
                a = +1
            if not s.get_balance():
                s.set_balance (a)
                self.height = self.height + 1
                return ii
            elif s.get_balance() == -a:
                s.set_balance (0)
                return ii
            elif s.get_balance() == a:
                if r.get_balance() == a:
                    # single rotation
                    # -1 == left
                    # +1 == right
                    p = r
                    if a == -1:
                        s.left = r.right
                        if r.right:
                            r.right.parent = s
                        r.right = s
                        s.parent = r
                        s.set_rank (s.get_rank() - r.get_rank())
                    else:
                        s.right = r.left
                        if r.left:
                            r.left.parent = s
                        r.left = s
                        s.parent = r
                        r.set_rank (r.get_rank() + s.get_rank())
                    s.set_balance(0)
                    r.set_balance(0)
                elif r.get_balance() == -a:
                    # double rotation
                    # -1 == left
                    # +1 == right
                    if a == -1:
                        p = r.right
                        r.right = p.left
                        if p.left:
                            p.left.parent = r
                        p.left = r
                        r.parent = p
                        s.left = p.right
                        if p.right:
                            p.right.parent = s
                        p.right = s
                        s.parent = p
                        p.set_rank (p.get_rank() + r.get_rank())
                        s.set_rank (s.get_rank() - p.get_rank())
                    else:
                        p = r.left
                        r.left = p.right
                        if p.right:
                            p.right.parent = r
                        p.right = r
                        r.parent = p
                        s.right = p.left
                        if p.left:
                            p.left.parent = s
                        p.left = s
                        s.parent = p
                        r.set_rank (r.get_rank() - p.get_rank())
                        p.set_rank (p.get_rank() + s.get_rank())
                    if p.get_balance() == a:
                        s.set_balance (-a)
                        r.set_balance (0)
                    elif p.get_balance() == -a:
                        s.set_balance (0)
                        r.set_balance (a)
                    else:
                        s.set_balance (0)
                        r.set_balance (0)
                    p.set_balance(0)
                # finishing touch
                if s == t.right:
                    t.right = p
                else:
                    t.left = p
                p.parent = t
                return ii

    def __len__ (self):
        return self.length

    def verify (self):
        self._verify_balance (self.tree.right)
        self._verify_rank    (self.tree.right)
        self._verify_parent  (self.tree.right, self.tree)

    # this verifies balance 'manually', and also
    # double-checks each node's <balance> member.
    cdef _verify_balance (self, avl_node node):
        if node == None:
            return 0
        lh = self._verify_balance (node.left)
        rh = self._verify_balance (node.right)
        if (rh - lh) != node.get_balance():
            raise Invalid_Balance, 'at node %r' % (node,)
        if abs (lh - rh) > 1:
            raise Unbalanced, 'at node %r' % (node,)
        else:
            return (1 + max (lh, rh))

    cdef _verify_rank (self, avl_node node):
        if not node:
            return 0
        num_left = num_right = 0
        if node.left:
            num_left = self._verify_rank (node.left)
        if node.right:
            num_right = self._verify_rank (node.right)
        if node.rank != num_left + 1:
            raise Invalid_Rank, 'at node %r' % (node,)
        return num_left + num_right + 1

    cdef _verify_parent (self, avl_node node, avl_node parent):
        if node.parent != parent:
            raise Invalid_Parent, 'at node %r' % (node,)
        if node.left:
            self._verify_parent (node.left, node)
        if node.right:
            self._verify_parent (node.right, node)