ufes-20191-redes-mininet/proj-impl/topoviewer.py

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

import sys
import time
import json
import tkinter
import networkx
import threading
from pathlib import Path
from graphtools import Dijkstra, graph_from_topo
from sortundirectednodepair import _sort_pair
from networkx.drawing.layout import spring_layout
from networkx.drawing.nx_pylab import draw_networkx
from networkx.drawing.nx_pylab import draw_networkx_labels
from networkx.drawing.nx_pylab import draw_networkx_nodes
from networkx.drawing.nx_pylab import draw_networkx_edges
from networkx.drawing.nx_pylab import draw_networkx_edge_labels
import matplotlib
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

matplotlib.use("TkAgg")


def nx_graph_from_topo(topo):
    g = networkx.Graph()
    g.add_nodes_from(topo[0])
    g.add_nodes_from(topo[1])
    g.add_edges_from([(x, y) for x, y, z in topo[2]])
    return g


def parse_state(s):
    pathb, loadb = [], []
    current = 0
    for line in s.splitlines():
        if line.startswith('@'):
            if line.startswith('@start.'):
                if line.endswith('path'):
                    current = 1
                elif line.endswith('load'):
                    current = 2
                else:
                    raise ValueError(f"Start of unexpected block: {line}")
            elif line.startswith('@end'):
                current = 0
            else:
                raise ValueError(f"Unexpected command: {line}")
        elif current == 1:
            pathb.append(eval(line))
        elif current == 2:
            loadb.append(eval(line))
    loads, path_segments = dict(), list()
    for path in pathb:
        for i in range(len(path)-1):
            path_segments.append(_sort_pair(*path[i:i+2]))
    for load in loadb:
        loads[_sort_pair(load[0], load[1])] = load[2]
    return pathb, set(path_segments), loads


CLR = {True: 'FF', False: '00'}


def plot(nxg, topo, statepath, djkt, pos, ax):
    state_txt = ""
    if statepath.exists():
        state_txt = statepath.read_text()
    paths, path_segments, loads = parse_state(state_txt)
    ax.clear()
    ax.axis('off')
    labels = dict()
    for unsortededge in nxg.edges:
        edge = _sort_pair(*list(map(str, unsortededge)))
        labels[unsortededge] = "%0.4f" % (loads.get(edge, 0.0),)
        # labels[unsortededge] = repr(loads.get(edge, 0.0),)
    edlab = draw_networkx_edge_labels(
        nxg, pos, ax=ax,
        edge_labels=labels,
        bbox=dict(
            facecolor='none',
            edgecolor='none'
        )
    )
    for unsortededge in nxg.edges:
        edge = _sort_pair(*list(map(str, unsortededge)))
        in_djkt = edge in djkt
        in_sgmt = edge in path_segments
        color = '#'+CLR[in_djkt]+'00'+CLR[in_sgmt]
        if in_djkt and in_sgmt:
            color = 'green'
        draw_networkx_edges(
            nxg, pos, ax=ax,
            edgelist=[edge], edge_color=color
        )
    draw_networkx_nodes(
        nxg, pos, ax=ax,
        nodelist=topo[0], node_color='lightgreen'
    )
    draw_networkx_nodes(
        nxg, pos, ax=ax,
        nodelist=topo[1], node_color='cyan'
    )
    draw_networkx_labels(nxg, pos, ax=ax)


continuePlotting = True


def to_closing_state():
    global continuePlotting
    continuePlotting = False


def show_window(nxg, topo, statepath, djkt):
    root = tkinter.Tk()
    root.title("Network graph viewer")
    root.geometry("1000x600")

    lab = tkinter.Label(root, text="Plotting network")
    lab.pack()

    fig = Figure()

    ax = fig.add_subplot(111)
    ax.axis('off')
    fig.patch.set_facecolor("#D9D9D9")

    graph = FigureCanvasTkAgg(fig, master=root)
    graph.get_tk_widget().pack(side="top", fill='both', expand=True)

    pos = spring_layout(nxg, iterations=2000)

    b = tkinter.Button(root, text="Update GUI", command=graph.draw)
    b.pack()

    def plotter():
        while continuePlotting:
            plot(nxg, topo, statepath, djkt, pos, ax)
            b.invoke()
            time.sleep(0.1)


    def close_cmd():
        to_closing_state()
        root.destroy()

    thread = threading.Thread(target=plotter)
    root.protocol("WM_DELETE_WINDOW", close_cmd)
    thread.start()

    root.mainloop()


def main(topo, statepath):
    djkt = Dijkstra(graph_from_topo(topo))
    statepath = Path(statepath)
    nxg = nx_graph_from_topo(topo)
    djkt_pairs = list()
    for h1 in topo[0]:
        for h2 in topo[0]:
            if h1 != h2:
                h1, h2 = _sort_pair(h1, h2)
                djkt_seq = djkt(h1)(h2)[0]
                for i in range(len(djkt_seq)-1):
                    djkt_pairs.append(_sort_pair(*djkt_seq[i:i+2]))
    show_window(nxg, topo, statepath, set(djkt_pairs))


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print('Usage:')
        print(f'  {sys.argv[0]} <toponame>')
        print()
        print('   Where toponame is will be resolved to')
        print('   toponame.json and toponame.state')
    else:
        topopath = Path(f'{sys.argv[1]}.json')
        if not topopath.exists():
            print(f'Topology {topopath} does not exist.')
        else:
            topo = json.loads(topopath.read_text())
            statepath = Path(f'{sys.argv[1]}.state')
            main(topo, statepath)