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, 4, 8, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((12, 4, 12, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((15, 16, 15, 49), 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', (), '', False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), ((16, 4, 17, 48), 'networkx.draw', 'nx.draw', (), '', True, 'import networkx as nx\n'), ((29, 4, 30, 81), 'networkx.draw_networkx_labels', 'nx.draw_networkx_labels', (), '', True, 'import networkx as nx\n'), ((34, 20, 34, 53), 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', (), '', False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), ((35, 22, 35, 68), 'networkx.get_edge_attributes', 'nx.get_edge_attributes', ({(35, 45, 35, 49): 'tree', (35, 51, 35, 67): '"""transformation"""'}, {}), "(tree, 'transformation')", True, 'import networkx as nx\n'), ((36, 8, 37, 98), 'networkx.draw_networkx_edge_labels', 'nx.draw_networkx_edge_labels', (), '', True, 'import networkx as nx\n')]