/home/mark/model/software/ScrumPy/ScrumPy/Util/TreeLayouts.py

#
# Put this in ScrumPy.Util.Set
Intersect = lambda l1,l2: filter(lambda x: x in l1,l2)


from ete2  import faces, TreeImageProperties


Black = "#000000"
White = "#FFFFFF"
Red   = "#FF0000"
Green = "#00FF00"
Blue  = "#0000FF"
ETE2fg= "#0030c1"

Solid =0
Dashed=1



def AddTextFace(n,text,col=0):
    faces.add_face_to_node(faces.TextFace(text),n,col)

def AddAttrFace(n, attr, col=0):
    faces.add_face_to_node(faces.AttrFace(attr),n,col)



# NodeTypes

Undef, Root, Internal, Module, Leaf = -1,0,1,2,3



######################### Image Properties ######################

DefaultImgPropDic={
    "force_topology"            : False,
    "draw_branch_length "       : False,
    "align_leaf_faces "         : False,
    "orientation"               : 0,
    "general_font_type"         : "Verdana",
    "branch_length_font_color"  : "#222",
    "branch_length_font_size"   : 6,
    "branch_support_font_color" : "red",
    "branch_support_font_size"  : 9,
    "tree_width"                : 850, #200,  # This is used to scale the tree branches
    "self.min_branch_separation": 1
}


class ImgProps(TreeImageProperties):
    
    def __init__(self, PropDic=DefaultImgPropDic):
        TreeImageProperties.__init__(self)
        self.Update(PropDic)

    def Update(self, PropDic):
        for p in PropDic:
            setattr(self,p, PropDic[p])
       


#
##
######################### Styles ######################
##
#

######################### Default #####################

DefaultStyle = {
    'vlwidth'         : 3,
    'vt_line_color'   : Black,
    'draw_descendants': 1,
    'hz_line_color'   : Black,
    'hlwidth'         : 3,
    'bgcolor'         : White,
    'shape'           : 'sphere',
    'line_type'       : 0,
    'ymargin'         : 0,
    'fgcolor'         : ETE2fg,
    'size'            : 6
}

class Style(dict):

    def __init__(self, BaseStyle=DefaultStyle, styledic={}):

        for style in Undef, Root, Internal, Module, Leaf:
            self[style] = dict(BaseStyle)

        for style in styledic:
            self[style].update(styledic[style])

############################  Explore   #####################


ExplBaseStyle = dict(DefaultStyle)
ExplBaseStyle.update({
    'vlwidth': 3,
    'hlwidth': 3,
    'fgcolor': White,
    'size'   : 10 })

ExplStyles = {
    Root     : {'fgcolor': Red  },
    Module   : {'fgcolor': Blue },
    Internal : {'fgcolor': Black},
    Leaf     : {'fgcolor' :Green}
}

ExplStyle = Style(ExplBaseStyle, ExplStyles)





#
##
########################## Layouts #########################
##
#

########################## Default #########################


class Layout:

    def __init__(self, StyleDic=ExplStyle):

        self.StyleDic = StyleDic

    def Layout(self, n):

        self.AllNodes(n)

        NodeType = n.GetNodeType()  

        n.img_style.update(self.StyleDic[NodeType])
        
        {
            Root     : self.Root,
            Internal : self.Internal,
            Module   : self.Module,
            Leaf     : self.Leaf
        } [NodeType](n)

    def AllNodes(self, n):
        pass

    def Root(self,n):
        AddAttrFace(n,"name")

    def Internal(self,n):
        AddAttrFace(n,"name")

    def Module(self, n):
        AddAttrFace(n,"name")

    def Leaf(self, n):
        AddAttrFace(n,"name")
        
DefaultLayout = Layout()



########################## LastModule #########################

# Expand all but the terminal modules    
class _LastModule(Layout):

    def Module(self,n):
       
        if not n.HasModule():
            n.collapsed = True
            lenstr = n.nLeafStr()
        else:
            lenstr = ""
            
        AddTextFace(n, n.name+lenstr)
        

    def Internal(self, n):
       
        if not n.HasModule():
            n.collapsed = True
            lenstr = n.nLeafStr()
            AddTextFace(n, lenstr)

    def Leaf(self,n):
        pass
          
LastModule = _LastModule()



class _FirstModule(_LastModule):

    def Root(self,n):

        n.dist = n.MeanDist()

    def Module(self,n):

        n.collapsed = True
        AddTextFace(n, n.name + n.nLeafStr())

FirstModule = _FirstModule()


########################## InternalOnly #########################


# only show nodes with no leaf children  and no modules
class _InternalOnly(Layout):

    def Internal(self, n):
        
        #n.img_style["size"] = 0
        if not n.HasModule():
            n.collapsed = True
            AddTextFace(n, n.nLeafStr())
            

    def Module(self, n):
        
        if n.HasModule():  
            #n.img_style["size"] = 0
            AddAttrFace(n,"name")
        else:
            n.collapsed = True
            AddTextFace(n, n.name + n.nLeafStr())


        
InternalOnly = _InternalOnly()
        


######################### Collapse preserving target leaves ######

class _CollapsePreserveTargs(Layout):

    def Module(self, n):
        
        AddAttrFace(n,"name")
        LNames = map(lambda node: node.name, n.get_leaves())

        if len(Intersect(LNames, n.GetRoot().LayoutData)) == 0:
            n.collapsed = True
        else:
            n.collapsed = False

       
    def Internal(self,n):
        
        LNames = map(lambda node: node.name, n.get_leaves())
        if len(Intersect(LNames, n.GetRoot().LayoutData)) == 0:
            n.collapsed = True
        else:
            n.collapsed = False
            AddAttrFace(n,"name")


    def Leaf(self, n):
        
        if n.name in n.GetRoot().LayoutData:
            AddAttrFace(n,"name")
            
            


        
CollapsePreserveTargs = _CollapsePreserveTargs()
        
            




##########################  Edit layout  #########################

EditBaseStyle = dict(DefaultStyle)

EditBaseStyle['size'] = 0
EditStyles = {
    Root     : {'fgcolor': Red  },
    Module   : {'fgcolor': Blue },
    Internal : {'fgcolor': Black, "size":5},
}

EditStyle = Style(EditBaseStyle, EditStyles)


class _Edit(Layout):

    def Module(self,n):
        AddAttrFace(n,"name")
        n.dist=0

    def Leaf(self, n):
        AddAttrFace(n,"name")

Edit = _Edit(EditStyle)
Explore = _Edit(ExplStyle)


class _LeavesAndModules(_Edit):

    def Internal(self, n):
        pass

LeavesAndModules = _LeavesAndModules(ExplStyle)

    

##########################  Skeleton layout  #########################
            

SkelBaseStyle = dict(DefaultStyle)  
SkelBaseStyle["size"] = 0
SkelBaseStyle['fgcolor'] =  Black
SkelStyle = Style(SkelBaseStyle)


class _Skeleton (Layout):
    def __init__(self):
        Layout.__init__(self,SkelStyle)

    def Root(self, n):
        n.dist = n.MeanDist()
        n.SetAllLeafMean()
        

    def Module(self, n):
        pass

    def Internal(self, n):
        pass

    def Leaf(self, n):
        pass

Skeleton = _Skeleton()

class _LeafNameOnly(_Skeleton):

    def __init__(self):
        _Skeleton.__init__(self)

    def Leaf(self, n):
        AddAttrFace(n,"name")

LeafNameOnly = _LeafNameOnly()
        

########################################################################
        


Layouts = {
    "LastModule"   : lambda n:LastModule.Layout(n),
    "FirstModule"  : lambda n:FirstModule.Layout(n),
    "InternalOnly" : lambda n:InternalOnly.Layout(n),
    "Default"      : lambda n:Default.Layout(n),
    "Edit"         : lambda n:Edit.Layout(n),
    "Explore"      : lambda n:Explore.Layout(n),
    "Skeleton"     : lambda n:Skeleton.Layout(n),
    "LeafNameOnly" : lambda n:LeafNameOnly.Layout(n),
    "LeavesAndModules":lambda n:LeavesAndModules.Layout(n),
    "CollapsePreserveTargs":  lambda n:CollapsePreserveTargs.Layout(n)
}