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

import os

import ete2
import TreeLayouts
# TODO: TreeLayouts should go in GUI, but circular dependency ??



reload(TreeLayouts)

mean = lambda seq: float(sum(seq))/len(seq)

DefaultNames = "-", "NoName"

#FixName = lambda name: name.replace("(","_").replace(")","_").replace(".","_")
#FixNames = lambda names: map(FixName, names)


class Tree(ete2.Tree):

    def __init__(self, newick=None, format=1):

        ete2.Tree.__init__(self,newick=newick,format=format)

        if newick != None and newick[-2:] != ");":  # ie newick is a file name (probably)
            name = newick.split(os.sep)[-1].split(".")[0]
            self.name = name

        #self.ImgProps = TreeLayouts.ImgProps()

        self.IsRoot = self.is_root        # just to make consistent w/ ScrumPy idiom
        self.IsLeaf = self.is_leaf
        self.GetRoot = self.get_tree_root

        self.LayoutData = None  # arbitary data for the potential use of layout functions


    def Copy(self):

        return Tree(self.NWStr())


    def IsInternal(self):
        return True not in (self.is_leaf(), self.is_root(), not self.name in DefaultNames)


    def HasChild(self, node):

        for ch in self.traverse():
            if ch == node:
                return True
        return False


    def IsModule(self):
        """ A module is assumed to be an internal node
            with a non-default name """
        return True not in (self.is_leaf(), self.is_root(), self.name in DefaultNames)


    def HasModule(self):

        for n in self.traverse():
            if n.IsModule() and n != self:
                return True
        return False


    def IsTerminalModule(self):
        """  is a module that has no modules in descendants """

        return self.IsModule() and (not self.HasModule())


    def GetModules(self):

        rv =  {}
        for n in self.traverse():
            if n.IsModule():
                mod = Tree(n.NWStr())       # work with a COPY of the subtree
                mod.dist = n.MeanDist()
                mod.name = n.name           # better way to do this re self.name = in __init__
                mod.Rename()
                rv[n.name] = mod

        return rv




    def GetNodesByName(self,Names):
        """ pre: Names is a list of strings
           post: GetNodesByName(Names) == list of nodes with name in Names
        """

        searchres = map(lambda n: self.search_nodes(name=n),Names)  # list of lists
        searchits = filter(lambda L: len(L)>0, searchres)
        
        return map(lambda L:L[0], searchits)

    def GetLeafNamesInOrder(self):
        rv = []
        for ch in self.children:
            if ch.IsLeaf():
                rv.append(ch.name)
            else:
                rv.extend(ch.GetLeafNamesInOrder())
        return rv
        

        
    def GetCommonAncestor(self, Names):

        targs   = self.GetNodesByName(Names)
        current = targs[0].up
        hits    = current.GetNodesByName(Names)
        
        while len(hits) < len(targs):
            current = current.up
            hits = current.GetNodesByName(Names)
            
        return current



    def PruneNodesMissing(self, Names):

        for ch in self.children[:]:
            if len(ch.GetNodesByName(Names)) == 0:
                self.children.remove(ch)
            else:
                ch.PruneNodesMissing(Names)
                

    def GetMinTreeHolding(self, Names):

        rv = self.GetCommonAncestor(Names).Copy()
        rv.PruneNodesMissing(Names)
        rv.MergeSingles()
        rv.SetName("-")

        return rv


    def HasOnlyLeaves(self):

        return False not in map(lambda ch: ch.IsLeaf(), self.children)


    def DistToFurthest(self):
        return self.get_farthest_leaf()[1]

    def DistToFurthestChild(self):

        dists = map(lambda node: node.dist, self.get_children())
        return max(dists)


    def MeanDist(self):

        return mean(map(lambda n: n.dist, self.traverse()))


    def Consolidate(self, thresh=0.0):

        def RecCons(n, thresh):
            if n.dist <= thresh and not n.is_root():
                RecCons(n.up,thresh)
                n.delete()

        for leaf in self.iter_leaves():
            RecCons(leaf.up, thresh)


    def MergeSingles(self):

        for ch in self.children[:]:
            if len(ch.children) == 1:
                ch.MergeSingles()
                self.dist += ch.dist
                ch.delete()
                
        for ch in self.children[:]:
            ch.MergeSingles()
                
                
                


        
        




    ### v   output and display  v




    def Save(self, features=None, outfile=None, format=1):

        if outfile==None:
            outfile=self.name + ".nw"

        ignore = ete2.Tree.write(self, features, outfile, format)


    def NWStr(self, format=1):

        return self.write(format=format)


    def Rename(self):

        for n in self.traverse():
            if n.name == "NoName":
                n.name="-"

    def SetName(self, name):

        self.name = name


    def nLeafStr(self):
        """ num leaves as a string """

        return "(" + str(self.nLeaf) + ")"

    def GetNodeType(self):

        return {
            self.IsRoot()    : TreeLayouts.Root,
            self.IsInternal(): TreeLayouts.Internal,
            self.IsModule()  : TreeLayouts.Module,
            self.IsLeaf()    : TreeLayouts.Leaf
        }[True]


    def _ExpandAll(self):

        for node in self.traverse():
            node.collapsed = False


    def _SetnLeaf(self):

        self._ExpandAll()

        for node in self.traverse():
            node.nLeaf = len(node.get_leaves())


    def _CopyDists(self):
        """ copy node distances, so that they can be restored if modified by layout functions """

        for n in self.traverse():
            n._Dist = n.dist


    def _RestoreDists(self):
        """ restore dists modified by layout functions """

        for n in self.traverse():
            n.dist = n._Dist



    def SetZeroLeafMean(self):

        dist = self.MeanDist()
        for leaf in self.get_leaves():
            if leaf.dist == 0:
                leaf.dist = dist

    def SetAllLeafMean(self):

        dist = self.MeanDist()
        for leaf in self.get_leaves():
            leaf.dist = dist

    def SetAllDists(self,dist):

        for n in self.traverse():
            n.dist = dist


    def SetAllDistsMean(self):
        self.SetAllDists(self.MeanDist())




    def SetImgProps(self,Props=None):

        pass


    def _Render(self, Layout=None, KeepCollapsed=False, ImgProps=None, FileName=None, Width=None, Height=None):

        Header = None

        self._CopyDists()
        self._SetnLeaf()

        self.SetImgProps(ImgProps)
        layout = TreeLayouts.Layouts[Layout]
        if FileName == None:
            self.show(layout = layout) #, image_properties = self.ImgProps)
        else:
            self.render(file_name = FileName,
                        layout    = layout,
                        header    = Header,
                        w         = Width,
                        h         = Height,
                        img_properties = TreeLayouts.ImgProps())

        if not KeepCollapsed:
            self._ExpandAll()
        self._RestoreDists()


    def Show(self, Layout="Explore",KeepCollapsed=False):

        self._Render(Layout=Layout,KeepCollapsed=KeepCollapsed)


    def SavePDF(self, fname=None, Layout="Explore",KeepCollapsed=False):

        if fname == None:
            fname = self.name

        fname += "_" + Layout + ".pdf"
        fname = fname.replace(" ","_")
        self._Render(Layout=Layout, FileName=fname,KeepCollapsed=KeepCollapsed)