"""
 tree_april12.py
 
 recursively search a tree
 
 in class April 12 2021
 
 $ python tree_april12.py
 first way : data structure, functions to extract info
A
B
D
H
I
E
J
K
C
F
G
 second way: pack into objects, use object methods
A
B
D
H
I
E
J
K
C
F
G
The number of nodes is 11
 
"""


a='A'; b='B'; c='C'; d='D'; e='E'; f='F'; g='G'; h='H'; i='I'; j='J'; k='K'


# each pair is (parent, child) in the tree
edges_1 = ( (a, b),
            (a, c),
            (b, d),
            (b, e),
            (c, f),
            (c, g),
            (d, h),
            (d, i),
            (e, j),
            (e, k)
           )

# --- one way to do things : let each node be in a class 

class Node:
    """ each node has node.name (a string) and node.children (list of nodes) """
    def __init__(self, name):
        self.name = name
        self.children = []
    def add_child(self, child):
        self.children.append(child)
    def search(self, function):
        """ apply this function to each node recursively """
        function(self.name)
        for child in self.children:
            child.search(function)

def edges_to_nodes(edges):
    """ take a sequence of edges and create a dictionary of nodes (Node objects) """
    # nodes = { 'a' : Node() , ... }
    nodes = {}
    for (parent_name, child_name) in edges:
        try:
            parent = nodes[parent_name] 
        except KeyError:
            nodes[parent_name] = Node(parent_name)
            parent = nodes[parent_name]
        try:
            child = nodes[child_name]
        except KeyError:
            nodes[child_name] = Node(child_name)
            child = nodes[child_name]
        parent.add_child(child)
    return nodes

# --- another way to do things - just work with the data structures,
# defining functions to extract the needed information

def children(node, down_edges):
    """ return list of children of nodes given list of downward links """
    # children of things in edges_1
    result = []
    for (parent, child) in down_edges:
        if parent == node:
            result.append(child)
    return result

def tree_search(node, edges, function):
    """ apply function to each element of a tree recursively """
    function(node)
    for child in children(node, edges):
        tree_search(child, edges, function)

counter = 0
def count(name):
    """ increment a counter """
    global counter
    counter = counter + 1
        
def main():
    print(" first way : data structure, functions to extract info ")
    tree_search(a, edges_1, print)
    
    print(" second way: pack into objects, use object methods ")
    nodes = edges_to_nodes(edges_1)
    nodes[a].search(print)
    
    tree_search(a, edges_1, count)
    print("The number of nodes is", counter)
    
if __name__ == '__main__':
    main()