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)

[((225, 235), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (233, 235), True, 'import matplotlib.pyplot as plt\n'), ((363, 375), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (373, 375), True, 'import matplotlib.pyplot as plt\n'), ((405, 438), 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', (['tree'], {'prog': '"""dot"""'}), "(tree, prog='dot')\n", (420, 438), False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), ((443, 504), 'networkx.draw', 'nx.draw', (['tree', 'nodes_pos'], {'node_size': '(150)', 'edge_color': '"""#7d7d7d"""'}), "(tree, nodes_pos, node_size=150, edge_color='#7d7d7d')\n", (450, 504), True, 'import networkx as nx\n'), ((845, 946), 'networkx.draw_networkx_labels', 'nx.draw_networkx_labels', (['tree', 'pos_attrs'], {'labels': 'custom_node_attrs', 'font_size': '(13)', 'bbox': 'nodes_bbox'}), '(tree, pos_attrs, labels=custom_node_attrs,\n    font_size=13, bbox=nodes_bbox)\n', (868, 946), True, 'import networkx as nx\n'), ((1018, 1051), 'networkx.drawing.nx_agraph.graphviz_layout', 'graphviz_layout', (['tree'], {'prog': '"""dot"""'}), "(tree, prog='dot')\n", (1033, 1051), False, 'from networkx.drawing.nx_agraph import graphviz_layout\n'), ((1074, 1120), 'networkx.get_edge_attributes', 'nx.get_edge_attributes', (['tree', '"""transformation"""'], {}), "(tree, 'transformation')\n", (1096, 1120), True, 'import networkx as nx\n'), ((1129, 1248), 'networkx.draw_networkx_edge_labels', 'nx.draw_networkx_edge_labels', (['tree'], {'pos': 'edges_pos', 'edge_labels': 'edge_labels', 'font_size': '(13)', 'rotate': 'rotate_edge_labels'}), '(tree, pos=edges_pos, edge_labels=edge_labels,\n    font_size=13, rotate=rotate_edge_labels)\n', (1157, 1248), True, 'import networkx as nx\n')]