# Python3 version of tree.py - much sleeker

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):
        self.firstSubhalo=firstSubhalo
        self.neighbour=neighbour
        self.firstProgenitor=firstProgenitor
        self.nextSibling=nextSibling

    # 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
        yield from self.nextSpatial()
        yield from self.nextTemporal('temporal')
    # More general iterator method with optional argument specifying tree type:
    def next(self,treeType=None):
        yield self.node
        if treeType==None:
            yield from self.nextSpatial()
            yield from self.nextTemporal('temporal')
        elif treeType=='spatial': yield from self.nextSpatial('spatial')
        elif treeType=='temporal': yield from self.nextTemporal('temporal')
        else: raise valueError('Invalid treeType for tree')
    # Spatial tree search
    def nextSpatial(self,treeType=None):
        if self.firstSubhalo != None: yield from self.firstSubhalo.next(treeType)
        if self.neighbour != None: yield from self.neighbour.next(treeType)
        #raise StopIteration()
    # Temporal tree search
    def nextTemporal(self,treeType=None):
        if self.firstProgenitor != None: yield from self.firstProgenitor.next(treeType)
        if self.nextSibling != None: yield from self.nextSibling.next(treeType)
        #raise StopIteration()

    # Printing
    def __str__(self):
        print ('node,firstSubhalo,neighbour,firstProgenitor,nextSibling=',
               self.node,self.firstSubhalo,self.neighbour,self.firstProgenitor,self.nextSibling)

    # Tracing
    def trace(self):
        for halo in self.next(): print(halo)
