#!/usr/bin/env python

import argparse
import csv
import math
import sys

try:
    from colour import Color
except:
    print('Colour module not installed, please install it')
    sys.exit(1)

try:
    import pydot
except:
    print('Pydot module not installed, please install it')
    sys.exit(1)

color_range = list(Color("DarkGreen").range_to(Color("DarkRed"), 20))
multiplier = float(len(color_range))

def get_color(value, max):
    if value < max:
        val = math.floor(multiplier * value / (max))
    else:
        val = math.floor(multiplier - 1)

    if value < 0:
        val = math.floor(multiplier - 1)

    color = color_range[val]

    return color.get_hex()

def get_penwidth(value, max):
    if max == 0:
        return "1.0"

    r = (4.0 * float(abs(value)) / (0.5 * max)) + 1
    return str(r)

description = "CSV to bandwidth graph script"
epilog = "2018 Sander Vrijders <sander.vrijders@ugent.be>"

parser = argparse.ArgumentParser(description=description,
                                 epilog=epilog)

parser.add_argument('-o', '--output', type=str, required=True,
                    help='Output PDF filename.')

parser.add_argument('-i', '--input', type=str, required=True,
                    help='Input CSV filename.')

parser.add_argument('-m', '--max', type=int, default=10000, required=False,
                    help='Maximum bandwidth in the network.')

args = parser.parse_args()

gvizg = pydot.Dot(graph_type='graph', layout='fdp')

csvfile = open(args.input, 'r')
reader = csv.reader(csvfile, delimiter=';', quotechar='|')

bandwidth = {}
nodes = set()
for row in reader:
    nodes.add(row[0])
    nodes.add(row[1])
    bandwidth[row[0], row[1]] = int(row[2])

node_bandwidth = {}
for key, value in bandwidth.items():
    n0, n1 = key
    if n0 in node_bandwidth:
        node_bandwidth[n0] += value
    else:
        node_bandwidth[n0] = value
    if n1 in node_bandwidth:
        node_bandwidth[n1] += value
    else:
        node_bandwidth[n1] = value

max_b = args.max
for key, value in bandwidth.items():
    max_b = max(value, max_b)

for node in nodes:
    try:
        color = get_color(value, 2 * max_b)
        g_node = pydot.Node(node,
                            label=node,
                            style="filled",
                            fillcolor=color,
                            color='white',
                            fontcolor='white')
        gvizg.add_node(g_node)
    except Exception as e:
        print('Failed to create pydot Node: ' + str(e))

for key, value in bandwidth.items():
    try:
        color = get_color(value, max_b)
        pw = get_penwidth(value, max_b)
        edge = pydot.Edge(key[0], key[1],
                          color = color,
                          penwidth=pw,
                          label=round(value * 8 / 1000000, 2))
        gvizg.add_edge(edge)
    except Exception as e:
        print('Failed to create pydot Edge: ' + str(e))

try:
    gvizg.write_pdf(args.output)
    print('Generated PDF of bandwidth graph')
except Exception as e:
    print('Failed to write PDF: ' + str(e))

csvfile.close()