import math
import random
import time
import matplotlib.pyplot as plt
import networkx as nx
from amplpy import AMPL, Environment, DataFrame

# prepare data
def make_data(n,m):
    I = range(1,n+1)
    J = range(1,m+1)
    x = {i:random.random() for i in range(max(m,n)+1)}    # positions of the points in the plane
    y = {i:random.random() for i in range(max(m,n)+1)}
    c = {}
    for i in I:
        for j in J:
            c[i,j] = math.sqrt((x[i]-x[j])**2 + (y[i]-y[j])**2)

    return I,J,c,x,y



random.seed(1)
n = 200
m = n
I,J,c,x_pos,y_pos = make_data(n,m)
k = 3

ampl = AMPL()
ampl.setOption('solver', 'gurobi')
ampl.read("kmedian.mod")
ampl.param['n'] = n
ampl.param['m'] = n
ampl.param['k'] = k
ampl.param['c'] = c

print("solving")
start = time.time()
ampl.solve()
cpu = time.time() - start
print("Used", cpu, "s")
cost = ampl.obj['cost']
print("Objective is:", cost.value())

facilities,edges = [],[]
x_ = ampl.var['x']
y_ = ampl.var['y']
for j in J:
    v = y_[j].value()
    if v > 0:
        facilities.append(j)
    for i in I:
        v = x_[i,j].value()
        if v > 0:
            edges.append((i,j))

# plot figure
plt.figure(figsize=(7, 7))
G = nx.Graph()
facilities = set(facilities)
unused = set(j for j in J if j not in facilities)
client = set(i for i in I if i not in facilities and i not in unused)
G.add_nodes_from(facilities)
G.add_nodes_from(client)
G.add_nodes_from(unused)
for (i,j) in edges:
    G.add_edge(i,j)

position = {}
for i in range(len(x_pos)):
    position[i] = (x_pos[i],y_pos[i])

nx.draw(G,position,with_labels=False,node_color="green",nodelist=facilities)
nx.draw(G,position,with_labels=False,node_color="c",nodelist=unused,node_size=50)
nx.draw(G,position,with_labels=False,node_color="g",nodelist=client,node_size=50)
plt.savefig("kmedian-{}-{}-{}.pdf".format(n,m,k))
plt.show()