
""" Tree class. """

from function import Function
from node import Node

class TreeMethods(object):
    """ A static class that has methods for generating trees. """
    @staticmethod
    def create_full_tree(depth):
        """ Creates a tree using the full method with depth `depth`.
            Returns the root node.

        Args:
            depth: int

        Returns:
            Node instance
        """
        if depth == 0:
            # Generate a leaf node
            terminal = Function.random_terminal()
            node = Node(terminal)
            return node
        else:
            # Generate an intermediate node
            func = Function.random_function()
            node = Node(func)

            for _ in range(func.arity):
                node.add_child(TreeMethods.create_full_tree(depth - 1))
            return node

    @staticmethod
    def create_grow_tree(depth):
        """ Creates a tree using the grow method with depth `depth`. Note
        that since grow ends when all leaf nodes are 0-arity, it is
        possible that the tree generated by this method has a depth
        less than `depth`.

        Args:
            depth: int

        Returns:
            Node instance
        """
        func = Function.random_function()
        root = Node(func)

        while True:
            if root.grow(depth) is None:
                break

        return root

if __name__ == "__main__":
    for i in range(10):
        tree = TreeMethods.create_full_tree(3)
        tree.ete_draw("demo/tree{0}.png".format(i))
