import matplotlib.pyplot as plt
import networkx as nx
from networkx.drawing.nx_agraph import graphviz_layout


def display_resolution_tree(resolution_tree: nx.classes.DiGraph):
    _draw_resolution_tree_(resolution_tree)
    plt.show()


def _draw_resolution_tree_(tree: nx.classes.DiGraph, enable_edge_labels: bool = True, rotate_edge_labels: bool = False):
    plt.figure()

    # graph
    nodes_pos = graphviz_layout(tree, prog='dot')
    nx.draw(tree, nodes_pos,
            node_size=150, edge_color='#7d7d7d')

    # nodes labels
    pos_attrs = {}
    for node, coords in nodes_pos.items():
        pos_attrs[node] = (coords[0], coords[1] - 10)

    custom_node_attrs = {}
    for node, attr in tree.nodes.items():
        custom_node_attrs[node] = str(node)

    nodes_bbox = dict(facecolor="w", edgecolor="#d3d3d3", pad=6, lw=0.1)
    nx.draw_networkx_labels(
        tree, pos_attrs, labels=custom_node_attrs, font_size=13, bbox=nodes_bbox)

    # edge labels
    if enable_edge_labels:
        edges_pos = graphviz_layout(tree, prog='dot')
        edge_labels = nx.get_edge_attributes(tree, 'transformation')
        nx.draw_networkx_edge_labels(
            tree, pos=edges_pos, edge_labels=edge_labels, font_size=13, rotate=rotate_edge_labels)

[(8, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (12, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (15, 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), (16, 'networkx.draw', 'nx.draw', True, 'import networkx as nx\n'), (29, 'networkx.draw_networkx_labels', 'nx.draw_networkx_labels', True, 'import networkx as nx\n'), (34, 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), (35, 'networkx.get_edge_attributes', 'nx.get_edge_attributes', True, 'import networkx as nx\n'), (36, 'networkx.draw_networkx_edge_labels', 'nx.draw_networkx_edge_labels', True, 'import networkx as nx\n')]