########################################
# The contents of this file are subject to the MLX PUBLIC LICENSE version
# 1.0 (the "License"); you may not use this file except in
# compliance with the License.
# 
# Software distributed under the License is distributed on an "AS IS"
# basis, WITHOUT WARRANTY OF ANY KIND, either express or implied.  See
# the License for the specific language governing rights and limitations
# under the License.
# 
# The Original Source Code is "compClust", released 2003 September 03.
# 
# The Original Source Code was developed by the California Institute of
# Technology (Caltech).  Portions created by Caltech are Copyright (C)
# 2002-2003 California Institute of Technology. All Rights Reserved.
########################################
#
#       Author: Lucas J. Scharenbroich
# 
# Implementation:  July 23 by Lucas Scharenbroich
#
##########################################################################

"""
A standard n-way tree inherited from a DirectedGraph

The multi-way tree is a specialized case of a directed graph.  Each
node is allowed to have only a single edge incident to it, though it may
have any number of edges leave each node 
"""

from Graph import Graph

class MultiWayTree(Graph):
  def __init__(self):
    Graph.__init__(self)
    self.__root  = None
    self.__links = {}

    
  def append(self, key, node):
    """
    Add a node as a child of an existing node.
    """

    self.addNode(node)
    nkey = node.key()

    if not self.edgeExists(key, nkey):
      self.addEdge(key, nkey)

      self.insertAt(key, nkey, len(self.children(key)) + 1)
      self.insertAt(nkey, key, 0)

      
  def children(self, key):
    """
    Returns a list of keys which are the children of the node with the
    passed key.  The list may contain None elements.
    """

    links = self.__links.get(key, [])
    return links[1:]

    
  def clear(self):
    Graph.clear(self)
    self.__root = None
    self.__links.clear()

    
  def depth(self, key):
    """
    Returns the depth of a node.
    """
      
    depth = -1
    while key != None:
      key = self.parent(key)
      depth += 1
            
    return depth


  def insert(self, node):
    """
    Adds a node to the tree.  If it is the first node added, then it
    is set as the root.  Otherwise it is a disconnected node.
    """

    Graph.addNode(self, node)
    
    if self.__root is None:
      self.__root = node.key()


  def insertAt(self, key1, key2, pos):

    links = self.__links.get(key1, [])

    n = len(links)
    if n <= pos:
      links += [None] * (pos - n + 1)

    links[pos] = key2

    self.__links[key1] = links

    
  def isLeaf(self, key):
    return len(filter(lambda x : x is not None, self.children(key))) == 0

        
  def iterator(self, type='LCR'):
    return MultiWayTreeIterator(self, type)

    
  def join(self, key, children=[]):

    for child in children:
      self.append(key, child)


  def parent(self, key):
    return self.__links[key][0]


  def prettyPrint(self):
    iter = self.iterator('RCL')

    node = iter.next()
    while node != None:
      depth = self.depth(node.key())
      print "   "*depth + str(node.value())
      node = iter.next()

            
  def removeNode(self, key):

    Graph.removeNode(self, key)

    del self.__links[key]

    if self.order() == 0:
      self.__root = None

  def setRoot(self, root):
    self.__root = root
    
  def root(self):
    return self.__root


##############################################################################
#
# Iterator
#
##############################################################################

LEFT      = 0
RIGHT     = -1
INORDER   = 1
PREORDER  = 2
POSTORDER = 3
    
class MultiWayTreeIterator:
    """
    MultiWayTreeIterator

    Provides for all basic types of traversals through a multiway tree.
    Inorder, preorder and postorder are all supported with biasing of the
    primary branch.  The type of traversal is specified in the constructor
    and defaults to 'LCR'.  The types of traversals are:

      * LCR - Left branch/Center node/Right branch (inorder)

      * RCL - Right branch/Center node/Left branch

      * CLR - Center Node/Left branch/Right branch (preorder)

      * CRL - Center Node/Right branch/Left branch

      * LRC - Left branch/Right branch/Center node (postorder)

      * RLC - Right branch/Left branch/Center node

    If there are N child nodes, 'left branch' refers to all the child
    nodes from [0, N/2-1] and 'right branch' refers to all the child nodes
    from [N/2, N-1].  Note that if there are and odd number of child nodes,
    then the LCR and RCL traversals will give a slightly different ordering
    due to their respevtive choices for a center node.
    """
    
    def __init__(self, tree, type='LCR'):
        """ construct a new iterator """

        #
        # These are the valid ways to iterate through a n-ary tree, the first
        # two are in-order (left-right-center) and second two are pre-order
        # and the last two are post-order
        
        self.mode = 0

        self.type = INORDER
        if type == "CLR" or type == "CRL":
            self.type = PREORDER
        elif type == "LRC" or type == "RLC":
            self.type = POSTORDER
            
        #
        # Determine what kind of iteration to perform, the legal values
        # are 0 = inorder, 1 = preorder, 2 = postorder.  The bias tell
        # whether to take (1) left branches or (-1) right branches first

        if type == "RCL" or type == "CRL" or type == "RLC":
            self.primary   = RIGHT
            self.secondary = LEFT
            self.bias      = -1
        else:  
            self.primary   = LEFT
            self.secondary = RIGHT
            self.bias      = 1

        #
        # The preorder traversal needs to process the root node first,
        # so introduce a small hack to keep it from missing it on the
        # first call to .next()

        self.preorderHack = 1
        self.tree         = tree
        self.setRoot(tree.root())
        
    def setRoot(self, root):
        self.current      = root
        self.root         = root
      
    def lastLeaf(self):
        """
        lastLeaf():

        Return the last leaf of an iterator, this will always be the leaf
        furthest down the secondary side.
        """
        
        tree  = self.tree

        index = 0
        if self.secondary == RIGHT:
            index = -1
                
        key = self.root
        while not tree.isLeaf(key):
            children = tree.children(key)
            key = filter(lambda x : x is not None, children)[index]
            
        return tree.find(key)

    def firstLeaf(self):
        """
        firstLeaf():

        Return the first leaf of an iterator, this will always be the leaf
        furthest down the primary side.
        """

        tree  = self.tree

        index = 0
        if self.primary == RIGHT:
            index = -1

        key = self.root
        while not tree.isLeaf(key):
            children = tree.children(key)
            key = filter(lambda x : x is not None, children)[index]

        return tree.find(key)

    def reverse(self):
        """
        reverse():
        
        Stop the iteration at the current node and reverses direction.
        The effect is that the current node is treated as a leaf. """

        self.mode = 1
    
    def next(self):
        """
        next(): Returns the next element for a given iterator.  If the end
                of the iteration is reached, None is returned."""
        
        nextNode = None
        tree     = self.tree
        
        if self.type == INORDER:
            nextNode = self.__next_inorder()
        
        elif self.type == PREORDER:
            nextNode = self.__next_preorder()

            #
            # If this is the first pass, reset the iterator and return the
            # root node.  Afterwards, used the next_preorder() method as
            # usual
            #
            
            if self.preorderHack == 1:
                self.current      = self.root
                nextNode          = tree.find(self.current)
                self.preorderHack = 0
        
        elif self.type == POSTORDER:
            nextNode = self.__next_postorder()

        return nextNode

    def __next_inorder(self):
        
        # mode 0:  scanning down the tree
        # mode 1:  scanning up the tree

        key  = self.current
        tree = self.tree
        
        #
        # At a leaf node, so now move up the tree and search toward the
        # secondary side, stopping is we cross the midpoint

        if self.mode == 1:

            while (1):
                if key is self.root:
                    return None

                pkey = key
                key  = tree.parent(key)

                children = tree.children(key)
                orig  = children.index(pkey)
                next  = self.__nextIndex(children, orig)

                numChildren = len(children) - 1

                if next is not None:
                    break

                endpoint = 0
                if self.primary < 0:
                    endpoint = numChildren
                    
                if not self.__straddle(endpoint, orig, numChildren):
                    self.current = key
                    return tree.find(key)

            if self.__straddle(orig, next, numChildren):
                self.current = children[next]
                self.mode    = 0
                return tree.find(key)

            key = children[next]

        #
        # To seach and burrow mode mean to keep going down the primary
        # side of a tree until either a leaf is reached, or the first
        # index is in the right half of the children

        self.mode = 0

        while not tree.isLeaf(key):
            children = tree.children(key)
            next = self.__nextIndex(children, None)

            numChildren = len(children) - 1

            endpoint = 0
            if self.primary < 0:
                endpoint = numChildren

            if self.__straddle(endpoint, next, numChildren):
                self.current = children[next]
                return tree.find(key)
                
            key = children[next]

        self.mode    = 1
        self.current = key
        return tree.find(key)

    def __next_postorder(self):
        
        # mode 0:  scanning to the primary side
        # mode 1:  scanning up the tree
        
        key  = self.current
        tree = self.tree
        
        #
        # First scan down the most primary child.  Once a leaf is hit
        # begin moving up the tree and shifting towards the secondary
        # side.  If we come up the last secondary link, return the node
        # but stay in mode 1

        if self.mode == 1:

            pkey = key
            key  = tree.parent(key)

            #
            # if we hit the root going up, then stop
                
            #if pkey is tree.root():
            if pkey is self.root:
                return None

            children = tree.children(key)
            index = children.index(pkey)
            index = self.__nextIndex(children, index)

            if index is not None:
                key = children[index]
                self.mode = 0
            else:
                self.current = key
                return tree.find(key)
        
        while not tree.isLeaf(key):
            children = tree.children(key)
            index = self.__nextIndex(children, None)
            key   = children[index]

        #
        # Hit a leaf node, now go to mode 1
                
        self.mode    = 1
        self.current = key
        return tree.find(key)
    
    def __next_preorder(self):

        tree = self.tree
        key  = self.current

        #
        # Return the first child node

        if self.mode == 0:
            if not tree.isLeaf(key):
                children     = filter(lambda x : x is not None,
                                      tree.children(key))
                self.current = children[self.primary]
                return tree.find(self.current)

        #
        # No children, so move up the tree and look to the secondary side
        # until we either hit the root or fins something

        self.mode = 0
        root      = self.root
        
        while key != root:

            pkey = key
            key  = tree.parent(key)

            children = tree.children(key)
            index = children.index(pkey)
            index = self.__nextIndex(children, index)

            if index != None:
                self.current = children[index]
                return tree.find(self.current)

        #
        # Hit the root, so we are done
        
        return None

    def __straddle(self, x, y, z):

        """__straddle(x,y,z): Returns true if the numbers x and y 'straddle'
                              the midpoint of z.  i.e. 1 and 6 straddle 10,
                              but 7 and 9 do not. """

        return ((z - (x << 1)) < 0) != ((z - (y << 1)) < 0)

    def __nextIndex(self, keys, index=None):

        bias = self.bias
        size = len(keys)
        val  = self.primary

        if val < 0:
            val += size

        if index != None:
            val = index + bias

        while val >= 0  and val < size:
            if keys[val] != None:
                break
            val += bias
        else:
            val = None
        
        return val
