'''
Created on Apr 5, 2016
@author: Brett Paufler
Copyright Brett Paufler


Graph of Tree exported to here as it takes too long to import:
    Networkx
    matplotlib
    
'''

import networkx as nx
import matplotlib.pyplot as plt

from maze_tree import Node, Tree

def graph_by_num(a_tree, sN="./output/graph_by_num.png", show=False, graph_type='spring'):
    '''Simplistic NetworkX graph of Tree.
    
    sN = saveName complete with relative path
    show = True, print to screen
        True or False, image is saved to file
    
    graph_type = NetworkX graph type (there are only four)
        circular: everything in a circle
        spring: random connections, default
        shell
        spectral
        NOTE: at a small enough size (number of Nodes),
            this selection means nothing
                as NetworkX defaults to spring.'''
    
    graph = {'circular': nx.draw_circular,
                  'shell': nx.draw_shell,
                  'spectral': nx.draw_spectral, 
                  'spring': nx.draw_spring}
    
    #Initialize NetworkX graph and add Tree edges
    G = nx.DiGraph()
    G.add_edges_from(a_tree.all_edges())
    
    #Call graph dictionary, complicated switch control structure, heavy lifting
    graph[graph_type](G, with_labels=True)

    plt.savefig(sN)
    if show:
        plt.show()
    plt.close()
    
if __name__ == '__main__':
    
    nodes = [Node(links=[1, 2]) for _ in range(10)]
    for graph_type in ['circular', 'spectral', 'shell', 'spring']:
        
        
        t = Tree(nodes=nodes)
        graph_by_num(t, sN='./output/graph_five_%s.png' % graph_type, graph_type=graph_type)

        #print t.all_edges()

        t.invert()
        graph_by_num(t, sN='./output/graph_five_%s_invert.png' % graph_type, graph_type=graph_type)