from pyvis.network import Network


def plot_networkx(graph, save_path=None, front_end_metric=None, root_cause_path=None):
    if root_cause_path:
        for node in root_cause_path:
            if node in list(graph.nodes()):
                graph.nodes.get(node).update({"color": "#dd4b39"})

    if front_end_metric and front_end_metric in list(graph.nodes()):
        graph.nodes.get(front_end_metric).update({"color": "#ffa500"})

    # pyvis画图
    pyvis_g = Network(width=1920, height=800, select_menu=True,
                      filter_menu=True, directed=True)
    pyvis_g.from_nx(graph)
    pyvis_g.toggle_physics(True)  #允许更流畅的图形交互默认=True
    # pyvis_g.show_buttons(filter_=['nodes', 'edges', 'physics']) #页面布局按钮
    if save_path:
        pyvis_g.show(save_path)
    else:
        pyvis_g.show("figure.html")

    # # networkx画图
    # import matplotlib.pyplot as plt
    # pos = nx.circular_layout(graph)  # 圆形布局，起到美化作用
    # nx.draw(graph, pos, with_labels=True, font_weight="bold")
    # plt.title("DiGraph")
    # if save_path:
    #     plt.savefig(save_path.replace(".html", ".png"))
    # else:
    #     plt.savefig("demo.png")

    return graph