from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import mst
 
cities, edge = mst.get_mst()
v1 = [[cities[x].x, cities[x].y] for x, y in edge]
v1 = np.array(v1)
v2 = [[cities[y].x, cities[y].y] for x, y in edge]
v2 = np.array(v2)
m = Basemap(width=6000000, height=6000000, projection='lcc',
            resolution='h', lat_0=30, lon_0=105)
m.readshapefile('bou2_4p', 'conties')
m.drawmapboundary(fill_color='#ffffff')
m.fillcontinents(color='#eeeeee', 
            lake_color='#689CD2', zorder=0)
x1, y1 = m(v1[:, 0], v1[:, 1])
x2, y2 = m(v2[:, 0], v2[:, 1])
v = [[c.x, c.y] for c in cities]
v = np.array(v)
x, y = m(v[:, 0], v[:, 1])
names = [c.name for c in cities]
m.scatter(x, y, marker='o')
y_offset    = 15.0
mpl.rcParams['font.sans-serif'] = ['SimHei']
for i in range(len(cities)):
    plt.text(x[i], y[i]+y_offset, names[i].decode('utf8'),
            rotation=20,fontsize=6)
for i in range(len(x1)):
    plt.plot([x1[i], x2[i]], [y1[i], y2[i]])
plt.show()

import math
import heapq
distMat = {}
marked = []
pq = []
class City:
    def __init__(self, id, name, longitude, latitude):
        self.id = id
        self.name = name
        self.x = longitude
        self.y = latitude
        
def create_graph(filename):
    with open(filename, 'r') as f:
        id = 0
        cities = []
        for line in f.readlines():
            arr = line.split(',')
            city = City(id, arr[0], float(arr[1]), float(arr[2]))
            cities.append(city)
            id += 1
    return cities
def create_edges(cities):
    adj = []
    N = len(cities)
    for i in range(N):
        for j in range(i):
            adj.append((j, i, dist(cities[i], cities[j])))
    return adj
def visit(adj, v):
    global marked, pq
    marked[v] = True
    for a, b, dist in adj:
        if a == v:
            if not marked[b]:
                heapq.heappush(pq, (dist, (a, b)))
        if b == v:
            if not marked[a]:
                heapq.heappush(pq, (dist, (a, b)))
def dist(city1, city2):
    from math import cos
    from math import sin
    from math import acos
	global distMat
    if city1.id < city2.id:
        pairs = (city1.id, city2.id)
    else:
        pairs = (city2.id, city1.id)
    if not pairs in distMat:
        log1 = math.radians(city1.x)
        lat1 = math.radians(city1.y)
        log2 = math.radians(city2.x)
        lat2 = math.radians(city2.y)
        x1 = cos(lat1) * cos(log1)
        y1 = cos(lat1) * sin(log1)
        z1 = sin(lat1)
        x2 = cos(lat2) * cos(log2)
        y2 = cos(lat2) * sin(log2)
        z2 = sin(lat2)
        distMat[pairs] = acos(x1 * x2 + y1 * y2 + z1 * z2)
        # distMat[pairs] = math.sqrt((city1.x - city2.x) ** 2 +\
        # (city1.y - city2.y) ** 2)
        return distMat[pairs]
def mst(cities, adj):
    global marked, pq, degree
    collect = []
    N = len(cities)
    degree = [0] * N
    marked = [False] * N
    visit(adj, 0)
    while len(pq) > 0:
        dist, (a, b) = heapq.heappop(pq)
        if marked[a] and marked[b]:
            continue
        collect.append((a, b))
        degree[a] += 1
        degree[b] += 1
        if not marked[a]: visit(adj, a)
        if not marked[b]: visit(adj, b)
    return collect
def get_mst():
    cities = create_graph('cities')
    adj = create_edges(cities)
    edge = mst(cities, adj)
    return cities, edge