import itertools


def make_nodenames_balanced(BFs, leaf_nodes=False, root = True):
    """ creates a list of node names as in aircond example
    Args:
        BFs (list): branching factor for each stage
        leaf_nodes (bool): if False, exclude leaf node names
        root (bool): if False, no "ROOT_" string on non-root nodes
    
    Returns: 
        nodenames (list): list of all nodenames, "ROOT" is last, ordered as:
        0, 1, 2, ..., BFs[0], 0_0, 0_1, ..., 0_BFs[1], 1_0, 1_1, ... ,
        1_BFs[1], ... , BFs[0]_BFs[1], ... , BFs[0]...BFs[-2]
    """
    if leaf_nodes == False:
        BFs = BFs[:-1] # exclude leaf nodes

    # Constructs all nodenames
    # 0, 1, 2, ..., BFs[0], 00, 01, ..., 0BFs[1], 10, 11, ... ,
    # 1BFs[1], ... , BFs[0]BFs[1], ... , BFs[0]\cdotsBFs[-2]

    max_str_len = len(BFs)
    lists = [[str(k) for k in range(BFs[0])]]
    for i in range(max_str_len-1):
        the_list = []
        for j in range(len(lists[i])):
            the_list.append([lists[i][j]+"_"+str(a) for a in range(BFs[i+1])])
        lists.append(list(itertools.chain(*the_list)))

    strings = list(itertools.chain(*lists))

    nodenames = []

    if root:
        for string in strings:
            rt_string = "ROOT_" + string
            nodenames.append(rt_string)
    else: 
        nodenames = strings
    # add root node (it is nicer to keep it at the end)
    nodenames.append("ROOT")
        
    return nodenames