from string import ascii_uppercase
from graphviz import Digraph, Graph

# Returns a function to generate node names with an optional prefix
def make_name_generator (length = 1, prefix = ''):
  wheels = [{ 'position': None, 'max': len(ascii_uppercase), 'values': list(ascii_uppercase)} for _ in range(length)]

  def name_generator ():
    for wheel in wheels:
      if wheel['position'] is None:
        wheel['position'] = 0
      else:
        wheel['position'] += 1
        if wheel['position'] < wheel['max']:
          break
        else:
          wheel['position'] = 0

    return prefix + ''.join(reversed([wheel['values'][wheel['position']] for wheel in wheels]))
  
  return name_generator

def visualize (tree, graphname):
  generate_node_name = make_name_generator(length=3)
  
  def visualize_node (graph, node, previous_node_name = None, tailport = None):
    if node:
      node_name = generate_node_name() 
      graph.node(node_name, label=str(node.value))
      
      if previous_node_name:
        graph.edge(previous_node_name, node_name, tailport=tailport, headport='s')

    
      visualize_node(graph, node.left, node_name, 'nw')
      # Hack to have a more beautiful layout
      visualize_node(graph, None, node_name, 'n')
      visualize_node(graph, node.right, node_name, 'ne')

    else:
      node_name = generate_node_name()
      graph.node(node_name, label=node_name, style='invis')
      graph.edge(previous_node_name, node_name, style='invis', tailport=tailport, headport='s')

  graph = Graph(name=graphname, format='svg', engine='dot')
  graph.attr('graph', splines='line', rankdir='BT')
  visualize_node(graph, tree)
  graph.render(graphname)

if __name__ == '__main__':
  from treesort import make_tree
  from random import randint

  values = [randint(0, 250) for _ in range(15)]
  tree = make_tree(values)
  visualize(tree, 'tree')