#!/usr/bin/env python
# ============================================================================
# name     : graph-usca.py
# author   : parke godfrey
# create   : 2017-11-28
# platform : Python3

import sys
import matplotlib.pyplot as plot
import networkx as nx

N      = 312 # number of cities
degree = 5;

# read in the list of cities
cities = []
with open('usca312_name.txt') as nameFD:
    i = 0
    for line in nameFD.read().splitlines():
        line = line.strip()
        if line.startswith('#'): # skip comments
            continue
        cname, province = line.split(',', 1)
        city = {}
        city['name']     = cname.strip()
        city['province'] = province.strip()
        city['nearby']   = []
        city['num']      = i
        cities.append(city)
        i += 1

# read in (x, y) coordinates of the cities
with open('usca312_xy.txt') as xyFD:
    i = 0
    for line in xyFD.read().splitlines():
        line = line.strip()
        if line.startswith('#'): # skip comments
            continue
        x, y = line.split(None, 1)
        cities[i]['x'] = float(x)
        cities[i]['y'] = float(y)
        i += 1

# read the distance matrix (N x N)
distMat = []
rowVec  = []
i = 0
with open('usca312_dist.txt') as distFD:
    for line in distFD.read().splitlines():
        line = line.strip()
        if line.startswith('#'): # skip comments
            continue
        for dist in line.split():
            if ((i % N == 0) and (i > 0)):
                distMat.append(rowVec)
                rowVec = []
            rowVec.append(int(dist))
            i += 1
if (i % N == 0):
    distMat.append(rowVec)

# add the nearest 'degree' neighbours of each city
# add reciprocals, too, since graph is undirected
for city in cities:
    i = city['num']
    dists = []
    for j in range(len(distMat[i])):
        dists.append((j, distMat[i][j]))
    d = degree
    for neighRec in sorted(dists, key=lambda n: n[1]):
        j    = neighRec[0]
        dist = neighRec[1]
        if (dist <= 0):
            continue
        if (d <= 0):
            break
        if (neighRec not in city['nearby']):
            city['nearby'].append(neighRec)
        neighbour = cities[j]
        revRec    = (i, distMat[i][j])
        if (revRec not in neighbour['nearby']):
            neighbour['nearby'].append(revRec)
        d -= 1

G = nx.Graph()
positions = {} # a position map of the nodes (cities)
for city in cities:
    cityID = city['name'] + ',' + city['province']
    G.add_node(cityID)
    positions[cityID] = (city['x'], city['y'])
    for neighbour in city['nearby']:
        if (city['num'] < neighbour[0]):
            near   = cities[neighbour[0]]
            nearID = near['name'] + ',' + near['province']
            G.add_edge(cityID, nearID)

# plot the city graph
plot.figure(figsize=(8,8))
nx.draw(G,
        positions,
        node_size   = [4 for v in G],
        with_labels = False)

#plot.savefig('city-plot.png')
plot.show()