"""tree.py
   Defines a Tree class for dealing with combined temporal and merger trees
"""

#-------------------------------------------------------------------------------
# The following imports the python3 print function into python2;
# It seesm to be quietly ignored in python3
from __future__ import print_function

# In Python 3 the loops of the form
#    for __ in self.something():
#       yield __
# can all be replaced by
#    yield from self.something()
# which makes for much cleaner code.

#-------------------------------------------------------------------------------
class Tree:
    """Tree class for dealing with combined temporal and merger trees"""

    # Constructor: creates an empty tree node.
    def __init__(self,node):
        self.node=node

    # Load pointer arrays.
    def load(self,firstSubhalo,neighbour,firstProgenitor,nextSibling,hostHalo=-1,descendant=-1):
        self.firstSubhalo=firstSubhalo
        self.neighbour=neighbour
        self.firstProgenitor=firstProgenitor
        self.nextSibling=nextSibling
        self.hostHalo=hostHalo
        self.descendant=descendant

    # Default iterator that takes no arguments:
    # Note the asymmetry between spatial and temporal: because there is more
    # than one way to reach a node, can only iterate over one of them.
    # Because of the nature of the trees, space has to be the outer loop.
    def __next__(self):
        yield self.node
        for __ in self.nextSpatial():
            yield __
        for __ in self.nextTemporal('temporal'):
            yield __
    # More general iterator method with optional argument specifying tree type:
    def next(self,treeType=None):
        yield self.node
        if treeType is None:
            for __ in self.nextSpatial():
                yield __
            for __ in self.nextTemporal('temporal'):
                yield __
        elif treeType=='spatial':
            for __ in self.nextSpatial('spatial'):
                yield __
        elif treeType=='temporal':
            for __ in self.nextTemporal('temporal'):
                yield __
        else:
            raise valueError('Invalid treeType for tree')
    # Spatial tree search
    def nextSpatial(self,treeType=None):
        if self.firstSubhalo is not None:
            for __ in self.firstSubhalo.next(treeType):
                yield __
        if self.neighbour is not None:
            for __ in self.neighbour.next(treeType):
                yield __
        #raise StopIteration()
    # Temporal tree search
    def nextTemporal(self,treeType=None):
        if self.firstProgenitor is not None:
            for __ in self.firstProgenitor.next(treeType):
                yield __
        if self.nextSibling is not None:
            for __ in self.nextSibling.next(treeType):
                yield __
        #raise StopIteration()

    # Printing
    def printheader(self):
        print ('node,firstSubhalo,neighbour,firstProgenitor,nextSibling,hostHalo,descendant=')
    def __str__(self):
        s1='-1'
        if self.firstSubhalo: s1=str(self.firstSubhalo.node)
        s2='-1'
        if self.neighbour: s2=str(self.neighbour.node)
        s3='-1'
        if self.firstProgenitor: s3=str(self.firstProgenitor.node)
        s4='-1'
        if self.nextSibling: s4=str(self.nextSibling.node) 
        s5='-1'
        if self.hostHalo: s5=str(self.hostHalo.node)
        s6='-1'
        if self.descendant: s6=str(self.descendant.node)
        return  str(self.node)+', '+s1+', '+s2+', '+s3+', '+s4+', '+s5+', '+s6

#-------------------------------------------------------------------------------
