
class BinaryTreesOnAParticleSet(object):

    def __init__(self, particles_set, name_of_firstchild_attribute, name_of_secondchild_attribute):
        self.particles_set = particles_set
        self.name_of_firstchild_attribute = name_of_firstchild_attribute
        self.name_of_secondchild_attribute = name_of_secondchild_attribute
        
        
    def iter_roots(self):
        return self.iter_binary_trees()
        
    def iter_binary_trees(self):
        binaries = self._binaries()
        if len(binaries) == 0:
            return
        binaries_children1 = self._get_inner_nodes(binaries, self.name_of_firstchild_attribute)
        binaries_children2 = self._get_inner_nodes(binaries, self.name_of_secondchild_attribute)
            
        roots = (binaries - (binaries_children1 + binaries_children2))
        
        for particle in roots:
            yield BinaryTreeOnParticle(particle, self.name_of_firstchild_attribute, self.name_of_secondchild_attribute)
            
    def particles_not_in_a_multiple(self):
        binaries = self._binaries()
        binaries_children1 = self._get_descendant_nodes(self.particles_set, self.name_of_firstchild_attribute)
        binaries_children2 = self._get_descendant_nodes(self.particles_set, self.name_of_secondchild_attribute)
        if len(binaries) == 0:
            return self.particles_set
        else:
            singles = (self.particles_set - (self.roots() + binaries_children1 + binaries_children2))
            return singles
    
    def roots(self):
        binaries = self._binaries()
        if len(binaries) == 0:
            return binaries
        binaries_children1 = self._get_inner_nodes(binaries, self.name_of_firstchild_attribute)
        binaries_children2 = self._get_inner_nodes(binaries, self.name_of_secondchild_attribute)
        return (binaries - (binaries_children1 + binaries_children2))
        
    def _binaries(self):
        return self.particles_set.select_array(lambda x : x != [None], [self.name_of_firstchild_attribute,])

    def _get_inner_nodes(self, set, name_of_attribute):
        descendants = self._get_descendant_nodes(set, name_of_attribute)
        return descendants.select_array(lambda x : x != [None], [name_of_attribute,])

    def _get_descendant_nodes(self, set, name_of_attribute):
        return getattr(set, name_of_attribute).as_set().compressed()
        
        
class BinaryTreeOnParticle(object):

    def __init__(self, particle, name_of_firstchild_attribute = "child1" , name_of_secondchild_attribute = "child2"):
        self.particle = particle
        self.name_of_firstchild_attribute = name_of_firstchild_attribute
        self.name_of_secondchild_attribute = name_of_secondchild_attribute
        
        
    def iter_descendants(self):
        stack = [self.particle]
        while len(stack) > 0:
            current = stack.pop()
            children = []
            child1 = getattr(current, self.name_of_firstchild_attribute)
            if not child1 is None:
                yield child1
                children.append(child1)
                
            child2 = getattr(current, self.name_of_secondchild_attribute)
            if not child2 is None:
                yield child2
                children.append(child2)
            
            stack.extend(reversed(children))
            
    def iter_leafs(self):
        stack = [self.particle]
        while len(stack) > 0:
            current = stack.pop()
        
            children = []
            child1 = getattr(current, self.name_of_firstchild_attribute)
            if not child1 is None:
                children.append(child1)
            
            child2 = getattr(current, self.name_of_secondchild_attribute)
            if not child2 is None:
                children.append(child2)
            
            stack.extend(reversed(children))
        
            if len(children) == 0:
                yield current


    def iter_inner_nodes(self):
        stack = [self.particle]
        while len(stack) > 0:
            current = stack.pop()
        
            children = []
            child1 = getattr(current, self.name_of_firstchild_attribute)
            if not child1 is None:
                children.append(child1)
            
            child2 = getattr(current, self.name_of_secondchild_attribute)
            if not child2 is None:
                children.append(child2)
            
            stack.extend(reversed(children))
        
            if len(children) > 0:
                yield current

    def get_inner_nodes_subset(self):    
        keys = [x.key for x in self.iter_inner_nodes()]
        return self.particle.particles_set._subset(keys)

    def get_descendants_subset(self):    
        keys = [x.key for x in self.iter_descendants()]
        return self.particle.particles_set._subset(keys)

    def get_leafs_subset(self):    
        keys = [x.key for x in self.iter_leafs()]
        return self.particle.particles_set._subset(keys)
    
    def copy(self):
        copy_of_set = self.get_tree_subset().copy()
        root = copy_of_set[0]
        return BinaryTreeOnParticle(
            root,
            name_of_firstchild_attribute = self.name_of_firstchild_attribute,
            name_of_secondchild_attribute = self.name_of_secondchild_attribute
        )
        
    def get_tree_subset(self):    
        keys = [x.key for x in iter(self)]
        return self.particle.particles_set._subset(keys)

    def iter_events(self):
        stack = [('start', self.particle)]
        while len(stack) > 0:
            event, current = stack.pop()
            yield event,current
            if event == 'end':
                continue
            stack.append( ('end', current, ) ) 
        
            children = []
            child1 = getattr(current, self.name_of_firstchild_attribute)
            if not child1 is None:
                children.append( ('start', child1, ) )
                
            child2 = getattr(current, self.name_of_secondchild_attribute)
            if not child2 is None:
                children.append( ('start', child2, ) )
            
            stack.extend(reversed(children))
            
    def iter_levels(self):
        level = -1
        
        for event, particle in self.iter_events():
            if event == 'start':
                level += 1
                yield level, particle
            else:
                level -= 1
    
    
    def __iter__(self):
        stack = [self.particle]
        while len(stack) > 0:
            current = stack.pop()
            yield current
            
            children = []
            child1 = getattr(current, self.name_of_firstchild_attribute)
            if not child1 is None:
                children.append(child1)
                
            child2 = getattr(current, self.name_of_secondchild_attribute)
            if not child2 is None:
                children.append(child2)
            
            stack.extend(reversed(children))
            

class ChildTreeOnParticleSet(object):

    def __init__(self, particles_set, names_of_child_attributes = ["child1" , "child2"]):
        self.particles_set = particles_set
        self.names_of_child_attributes = names_of_child_attributes
        
    @property
    def particle(self):
        return self.particle_set #a particle set associated, maybe upgrade to some aggregate particle
        
            
    def is_leaf(self):
        return False
            
    def iter_leafs(self):
        for node in self.iter_children():
            if node.is_leaf():
                yield node
            
    def iter_branches(self):
        for node in self.iter_children():
            if not node.is_leaf():
                yield node
                
            
    def iter_descendants(self):
        stack = list(reversed(list(self.iter_children())))
        while len(stack) > 0:
            current = stack.pop()
            
            yield current
            
            children = list(current.iter_children())
            
            stack.extend(reversed(children))
            
    
    def iter_events(self):
        stack = [('start', x) for x in (reversed(list(self.iter_children())))]
        while len(stack) > 0:
            event, current = stack.pop()
            yield event,current
            if event == 'end':
                continue
            stack.append( ('end', current, ) ) 
            
            children = list(current.iter_children())
            
            stack.extend([('start', x) for x in (reversed(list(self.iter_children())))])
            
            
    def iter_descendant_leafs(self):
        stack = list(reversed(list(self.iter_children())))
        while len(stack) > 0:
            current = stack.pop()

            if current.is_leaf():
                yield current
            else:
                children = list(current.iter_children())
                stack.extend(reversed(children))


    def iter_descendant_branches(self):
        
        stack = list(reversed(list(self.iter_children())))
        while len(stack) > 0:
            current = stack.pop()
            
            if not current.is_leaf():
                yield current
                children = list(current.iter_children())
                stack.extend(reversed(children))

    def iter_children(self):
        for particle in self._children():
            yield ChildTreeOnParticle(particle, self.names_of_child_attributes)
            
        
    def iter_levels(self):
        level = -1
        
        for event, particle in self.iter_events():
            if event == 'start':
                level += 1
                yield level, particle
            else:
                level -= 1
    
    
    
    def get_children(self):
        return list(self.iter_children())
        
    def __iter__(self):
        return self.iter_children()

    def _branches(self):
        binaries = self._binaries()
        result = binaries
        
        for name in self.names_of_child_attributes:
            result -= self._get_inner_nodes(binaries, name)
        
        return result
    
    def _inner_particles(self):
        result = None
        
        for name in self.names_of_child_attributes:
            descendants = self._get_descendant_nodes(self.particles_set, name)
            if result is None:
                result = descendants
            else:
                result += descendants
                
        return result
        
    def _children(self):
        return self.particles_set - self._inner_particles()
        
    def _get_descendant_nodes(self, set, name_of_attribute):
        return getattr(set, name_of_attribute).as_set().compressed()
        

    def get_children_subset(self):    
        return self._children()

class ChildTreeOnParticle(object):

    def __init__(self, particle, names_of_child_attributes = ["child1" , "child2"]):
        self.particle = particle
        self.names_of_child_attributes = names_of_child_attributes
        
        
    def iter_descendants(self):
        stack = [self.particle]
        while len(stack) > 0:
            current = stack.pop()
            children = []
            for name in self.names_of_child_attributes:
                child  = getattr(current, name)
                if not child is None:
                    yield ChildTreeOnParticle(child, self.names_of_child_attributes)
                    children.append(child)
            
            stack.extend(reversed(children))
            
    def is_leaf(self):
        for name in self.names_of_child_attributes:
            child  = getattr(self.particle, name)
            if not child is None:
                return False
        return True
        
    def is_binary(self):
        if not len(self.names_of_child_attributes) == 2:
            return False
            
        for child in self.iter_children():
            if not child.is_leaf():
                return False
        return True
            
    def iter_leafs(self):
        for node in self.iter_children():
            if node.is_leaf():
                yield node
                
    def iter_branches(self):
        for node in self.iter_children():
            if not node.is_leaf():
                yield node
            
    def iter_descendant_leafs(self):
        stack = list(reversed(list(self.iter_children())))
        while len(stack) > 0:
            current = stack.pop()

            if current.is_leaf():
                yield current
            else:
                children = list(current.iter_children())
                stack.extend(reversed(children))

    def iter_descendant_branches(self):
        
        stack = list(reversed(list(self.iter_children())))
        while len(stack) > 0:
            current = stack.pop()
            
            if not current.is_leaf():
                yield current
                children = list(current.iter_children())
                stack.extend(reversed(children))

    def get_children_subset(self):    
        keys = [x.particle.key for x in self.iter_children()]
        return self.particle.particles_set._subset(keys)
        
    def iter_children(self):
        current = self.particle
        
        for name in self.names_of_child_attributes:
            child  = getattr(current, name)
            if not child is None:
                yield ChildTreeOnParticle(child, self.names_of_child_attributes)
    
    def get_children(self):
        return list(self.iter_children())
        
    def get_children_particles(self):
        return [x.particle for x in self.iter_children()]
        
    def __iter__(self):
        return self.iter_children()
        
    def iter_levels(self):
        level = -1
        
        for event, particle in self.iter_events():
            if event == 'start':
                level += 1
                yield level, particle
            else:
                level -= 1
    
    def iter_events(self):
        stack = [('start', self,), ]
        while len(stack) > 0:
            event, current = stack.pop()
            yield event,current
            if event == 'end':
                continue
            stack.append( ('end', current, ) ) 
            
            children = list(current.iter_children())
            
            stack.extend([('start', x) for x in reversed(children)])
            
      
