###############################################################################
##
## Copyright (C) 2014-2015, New York University.
## Copyright (C) 2011-2014, NYU-Poly.
## Copyright (C) 2006-2011, University of Utah.
## All rights reserved.
## Contact: contact@vistrails.org
##
## This file is part of VisTrails.
##
## "Redistribution and use in source and binary forms, with or without
## modification, are permitted provided that the following conditions are met:
##
##  - Redistributions of source code must retain the above copyright notice,
##    this list of conditions and the following disclaimer.
##  - Redistributions in binary form must reproduce the above copyright
##    notice, this list of conditions and the following disclaimer in the
##    documentation and/or other materials provided with the distribution.
##  - Neither the name of the New York University nor the names of its
##    contributors may be used to endorse or promote products derived from
##    this software without specific prior written permission.
##
## THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
## AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
## THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
## PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
## CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
## EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
## PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
## OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
## WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
## OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
## ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
##
###############################################################################
""" This file has the implementation of an algorithm to layout general
rooted trees in a nice way. The "lw" in the file name "tree_layout_lw.py"
stands for (L)inear (W)alker. This code is based on the paper:

    Christoph Buchheim, Michael Junger, and Sebastian Leipert.
    Improving walker's algorithm to run in linear time.
    In Stephen G. Kobourov and Michael T. Goodrich, editors, Graph
    Drawing, volume 2528 of Lecture Notes in Computer Science, pages
    344-353. Springer, 2002.

which is a faster (it is linear!) way to compute the tree layout
proposed by Walker than the algorithm described by him.
The original paper is:

    John Q. Walker II.
    A node-positioning algorithm for general trees.
    Softw., Pract. Exper., 20(7):685-705, 1990.

"""

from __future__ import division

class TreeLW(object):
    """
    The input to the algorithm must be a tree
    in this format.

    """
    def __init__(self):
        self.nodes = []
        self.maxLevel = 0

    def root(self):
        return self.nodes[0]

    def addNode(self, parentNode, width, height, object = None):
        newNode = NodeLW(width,height,object)        
        self.nodes.append(newNode)

        # add
        if parentNode is not None:
            parentNode.addChild(newNode)

        # update max level
        self.maxLevel = max(self.maxLevel,newNode.level)
        return newNode

    def changeParentOfNodeWithNoParent(self, parentNode, childNode):
        if childNode.parent is not None:
            raise ValueError("Node already has a parent")

        parentNode.addChild(childNode)
        maxLevel = self.__dfsUpdateLevel(childNode)

        # update max level
        self.maxLevel = max(self.maxLevel, maxLevel)

    def __dfsUpdateLevel(self, node):
        if node.parent is None:
            node.level = 0
        else:
            node.level = node.parent.level + 1
        maxLevel = node.level
        for child in node.children:
            maxLevel = max(maxLevel, self.__dfsUpdateLevel(child))
        return maxLevel

    def boundingBox(self):
        kbb = KeepBoundingBox()
        for w in self.nodes:
            kbb.addPoint(w.x-w.width/2.0, w.y-w.height/2.0)
            kbb.addPoint(w.x+w.width/2.0, w.y+w.height/2.0)
        return kbb.getBoundingBox()            

    def getMaxNodeHeightPerLevel(self):
        result = [0] * (self.maxLevel+1)
        for w in self.nodes:
            level = w.level
            result[level] = max(result[level],w.height)
        return result

    @staticmethod
    def randomTree(n,  k=10000000):
        p = [0] * n
        import random
        for i in xrange(1, n):
            minIndex = max(i-k, 0)
            index = random.randint(minIndex, i-1) # random number in {0,1,2,...,i-1}
            p[i] = index
        t = TreeLW()
        nodes= []
        for i in xrange(n):
            if i==0:
                parent= None
            else:
                parent=nodes[p[i]]
            width = 5 + 10*random.random()
            height = 5 + 10*random.random()
            nodes.append(t.addNode(parent,width,height,i))
        return t            

class KeepBoundingBox(object):
    def __init__(self):
        self.minx = None
        self.miny = None
        self.maxx = None
        self.maxy = None
        self.size = 0

    def addPoint(self,x,y):
        if self.minx is None or self.minx > x:
            self.minx = x
        if self.miny is None or self.miny > y:
            self.miny = y
        if self.maxx is None or self.maxx < x:
            self.maxx = x
        if self.maxy is None or self.maxy < y:
            self.maxy = y
        self.size = self.size + 1

    def getBoundingBox(self):
        return [self.minx, self.miny, self.maxx-self.minx, self.maxy - self.miny]

class NodeLW(object):
    """
    Node of the tree with all the auxiliar
    variables needed to the LW algorithm.
    The fields width, height and object
    are given as input. The first two are
    used to layout while the last one might
    be used by the user of this class.

    """    
    def __init__(self, width, height, object = None):
        self.width = width
        self.height = height
        self.object = object

        self.children = []

        self.parent = None
        self.index = 0

        # level of the node
        self.level = 0
        
        # intermediate variables for 
        # layout algorithm
        self.mod = 0
        self.prelim = 0
        self.ancestor = None
        self.thread = None
        self.change = 0
        self.shift = 0

        # final center position
        self.x = 0
        self.y = 0
        
    def getNumChilds(self):
        return len(self.children)
        
    def hasChild(self):
        return bool(self.children)

    def addChild(self, node):
        self.children.append(node)
        node.index = len(self.children) - 1
        node.parent = self
        node.level = self.level + 1

    def isLeaf(self):
        return not self.children

    def leftChild(self):
        return self.children[0]

    def rightChild(self):
        return self.children[-1]

    def leftSibling(self):
        if self.index > 0:
            return self.parent.children[self.index-1]
        else:
            return None

    def leftMostSibling(self):
        if self.parent is not None:
            return self.parent.children[0]
        else:
            return self

    def isSiblingOf(self, v):
        return self.parent == v.parent and self.parent is not None

class TreeLayoutLW(object):

    """
    TreeLayoutLW: the LW stands for Linear Walker.

    This code is based on the paper:

    Christoph Buchheim, Michael Junger, and Sebastian Leipert.
    Improving walker's algorithm to run in linear time.
    In Stephen G. Kobourov and Michael T. Goodrich, editors, Graph
    Drawing, volume 2528 of Lecture Notes in Computer Science, pages
    344-353. Springer, 2002.

    which is a faster (linear) way to compute the tree layout
    proposed by Walker than the algorithm described by him.
    The original paper is:

    John Q. Walker II.
    A node-positioning algorithm for general trees.
    Softw., Pract. Exper., 20(7):685-705, 1990.
    """

    # vertical alignment of nodes in the 
    # level band: TOP, MIDDLE, BOTTOM
    TOP = 0
    MIDDLE = 1
    BOTTOM = 2

    def __init__(self, tree, vertical_alignment=1, xdistance=10, ydistance=10):
        self.xdistance = xdistance
        self.ydistance = ydistance
        self.tree = tree
        self.vertical_alignment = vertical_alignment
        self.treeLayout()

    def treeLayout(self):
        for v in self.tree.nodes:
            v.mod = 0
            v.thread = None
            v.ancestor = v
            
        r = self.tree.root()
        self.firstWalk(r)
        self.secondWalk(r, -r.prelim)
        self.setVerticalPositions()

    def setVerticalPositions(self):

        # set y position
        maxNodeHeightPerLevel = self.tree.getMaxNodeHeightPerLevel()
        info_level = []
        position_level = 0
        for level in xrange(len(maxNodeHeightPerLevel)):
            height_level = maxNodeHeightPerLevel[level]
            info_level.append((position_level,height_level))
            position_level += self.ydistance + height_level
            
        #
        for w in self.tree.nodes:
            level = w.level
            position_level, height_level = info_level[level]
            if self.vertical_alignment == TreeLayoutLW.TOP:
                w.y = position_level + w.height/2.0
            elif self.vertical_alignment == TreeLayoutLW.MIDDLE:
                w.y = position_level + height_level/2.0
            else: # bottom
                w.y = position_level + height_level - w.height/2.0
                
        
    def gap(self, v1, v2):

        return self.xdistance + (v1.width + v2.width)/2.0        


    def firstWalk(self, v):
        
        if v.isLeaf():
            v.prelim = 0
            w = v.leftSibling()
            if w is not None:
                v.prelim = w.prelim + self.gap(w,v)

        else:
            
            defaultAncestor = v.leftChild()
            for w in v.children:
                self.firstWalk(w)
                defaultAncestor = self.apportion(w, defaultAncestor)
            self.executeShifts(v)
            
            midpoint = (v.leftChild().prelim + v.rightChild().prelim) / 2.0

            w = v.leftSibling()
            if w is not None:
                v.prelim = w.prelim + self.gap(w,v)
                v.mod = v.prelim - midpoint
            else:
                v.prelim = midpoint


    def apportion(self,  v,  defaultAncestor):

        """
        Apportion: to divide and assign proportionally.

        Suppose the the left siblings of "v" are all aligned. 
        Now align the subtree with root "v" (note: the correct
        alignment of the left siblings of "v" are encoded 
        in the auxiliar variables, the x and y 
        are not correct; only in the end the correct 
        values are assigned to x  and y). 
        By property (*) in Section 4 the gratest 
        distinct ancestor of a node "w" in the 
        subtrees at rooted at the left siblings of "v"
        is w.ancestor if this value is a left sibling
        of v otherwise it is "defaultAncestor".

        """
        w = v.leftSibling()
        if w is not None:
            # p stands for + or plus (right subtree)
            # m stands for - or minus (left subtree)
            # i stands for inside
            # o stands for outside
            # v stands for vertex
            # s stands for shift
            vip = vop = v
            vim = w
            vom = vip.leftMostSibling()
            sip = vip.mod
            sop = vop.mod
            sim = vim.mod
            som = vom.mod
            while self.nextRight(vim) is not None and self.nextLeft(vip) is not None:
                
                vim = self.nextRight(vim)
                vip = self.nextLeft(vip)
                vom = self.nextLeft(vom)
                vop = self.nextRight(vop)

                vop.ancestor = v
                
                shift = (vim.prelim + sim) - (vip.prelim + sip) + self.gap(vim,vip)
                
                if shift > 0:
                    self.moveSubtree(self.ancestor(vim,v,defaultAncestor),v,shift)
                    sip += shift
                    sop += shift

                sim += vim.mod
                sip += vip.mod
                som += vom.mod
                sop += vop.mod

            if self.nextRight(vim) is not None and self.nextRight(vop) is None:
                vop.thread = self.nextRight(vim)
                vop.mod += sim - sop

            if self.nextLeft(vip) is not None and self.nextLeft(vom) is None:
                vom.thread = self.nextLeft(vip)
                vom.mod += sip - som
                defaultAncestor = v
            
        return defaultAncestor

    def nextLeft(self, v):
        if v.hasChild():
            return v.leftChild()
        else:
            return v.thread

    def nextRight(self, v):
        if v.hasChild():
            return v.rightChild()
        else:
            return v.thread

    def moveSubtree(self, wm, wp, shift):
        subtrees = float(wp.index - wm.index)
        wp.change += -shift/subtrees
        wp.shift += shift
        wm.change += shift/subtrees
        wp.prelim += shift
        wp.mod += shift

    def executeShifts(self, v):
        shift = 0
        change = 0
        for i in xrange(v.getNumChilds()-1,-1,-1):
            w = v.children[i]
            w.prelim += shift
            w.mod += shift
            change += w.change
            shift += w.shift + change

    def ancestor(self, vim, v, defaultAncestor):
        if vim.ancestor.isSiblingOf(v):
            return vim.ancestor
        else:
            return defaultAncestor

    def secondWalk(self,  v, m):
        v.x = v.prelim + m
        for w in v.children:
            self.secondWalk(w, m + v.mod)

# graph
if __name__ == "__main__":

    t = TreeLW()

    a = t.addNode(None,1,1,"a")

    b = t.addNode(a,1,1,"b")
    c = t.addNode(a,1,1,"c")
    d = t.addNode(a,1,1,"d")

    e = t.addNode(b,1,1,"e")
    f = t.addNode(b,1,1,"f")
    g = t.addNode(d,1,1,"g")
    h = t.addNode(d,1,1,"h")

    i = t.addNode(f,1,1,"i")
    j = t.addNode(f,1,1,"j")
    k = t.addNode(h,1,1,"k")
    l = t.addNode(h,1,1,"l")
    m = t.addNode(h,1,1,"m")
    n = t.addNode(h,1,1,"n")
    o = t.addNode(h,1,1,"o")
    
    layout = TreeLayoutLW(t)
    
