# -*- coding: utf-8 -*-
"""
Created on Wed Jun 08 15:21:03 2016

@author: SYSTEM
"""
import os
os.getcwd()
import numpy as np
import matplotlib.pyplot as plt
%pylab
coordinates = np.array([[565.0,575.0],[25.0,185.0],[345.0,750.0],[945.0,685.0],[845.0,655.0],
                        [880.0,660.0],[25.0,230.0],[525.0,1000.0],[580.0,1175.0],[650.0,1130.0],
                        [1605.0,620.0],[1220.0,580.0],[1465.0,200.0],[1530.0,  5.0],[845.0,680.0],
                        [725.0,370.0],[145.0,665.0],[415.0,635.0],[510.0,875.0],[560.0,365.0],
                        [300.0,465.0],[520.0,585.0],[480.0,415.0],[835.0,625.0],[975.0,580.0],
                        [1215.0,245.0],[1320.0,315.0],[1250.0,400.0],[660.0,180.0],[410.0,250.0],
                        [420.0,555.0],[575.0,665.0],[1150.0,1160.0],[700.0,580.0],[685.0,595.0],
                        [685.0,610.0],[770.0,610.0],[795.0,645.0],[720.0,635.0],[760.0,650.0],
                        [475.0,960.0],[95.0,260.0],[875.0,920.0],[700.0,500.0],[555.0,815.0],
                        [830.0,485.0],[1170.0, 65.0],[830.0,610.0],[605.0,625.0],[595.0,360.0],
                        [1340.0,725.0],[1740.0,245.0]])

def getdistmat(coordinates):
    num = coordinates.shape[0]
    distmat = np.zeros((52,52))
    for i in range(num):
        for j in range(i,num):
            distmat[i][j] = distmat[j][i]=np.linalg.norm(coordinates[i]-coordinates[j])
    return distmat

distmat = getdistmat(coordinates)

numant = 40 #螞蟻個數
numcity = coordinates.shape[0] #城市個數
alpha = 1   #資訊素重要程度因子
beta = 5    #啟發函式重要程度因子
rho = 0.1   #資訊素的揮發速度
Q = 1

iter = 0
itermax = 250

etatable = 1.0/(distmat+np.diag([1e10]*numcity)) #啟發函式矩陣，表示螞蟻從城市i轉移到矩陣j的期望程度
pheromonetable  = np.ones((numcity,numcity)) # 資訊素矩陣
pathtable = np.zeros((numant,numcity)).astype(int) #路徑記錄表

distmat = getdistmat(coordinates) #城市的距離矩陣

lengthaver = np.zeros(itermax) #各代路徑的平均長度
lengthbest = np.zeros(itermax) #各代及其之前遇到的最佳路徑長度
pathbest = np.zeros((itermax,numcity)) # 各代及其之前遇到的最佳路徑長度


while iter < itermax:


    # 隨機產生各個螞蟻的起點城市
    if numant <= numcity:#城市數比螞蟻數多
        pathtable[:,0] = np.random.permutation(range(0,numcity))[:numant]
    else: #螞蟻數比城市數多，需要補足
        pathtable[:numcity,0] = np.random.permutation(range(0,numcity))[:]
        pathtable[numcity:,0] = np.random.permutation(range(0,numcity))[:numant-numcity]

    length = np.zeros(numant) #計算各個螞蟻的路徑距離

    for i in range(numant):

        #i=0
        visiting = pathtable[i,0] # 當前所在的城市

        #visited = set() #已訪問過的城市，防止重複
        #visited.add(visiting) #增加元素
        unvisited = set(range(numcity))#未訪問的城市
        unvisited.remove(visiting) #刪除元素


        for j in range(1,numcity):#迴圈numcity-1次，訪問剩餘的numcity-1個城市
            #j=1
            #每次用輪盤法選擇下一個要訪問的城市
            listunvisited = list(unvisited)

            probtrans = np.zeros(len(listunvisited))

            for k in range(len(listunvisited)):
                probtrans[k] = np.power(pheromonetable[visiting][listunvisited[k]],alpha)\
                        *np.power(etatable[visiting][listunvisited[k]],alpha)
            cumsumprobtrans = (probtrans/sum(probtrans)).cumsum()

            cumsumprobtrans -= np.random.rand()

            k = listunvisited[find(cumsumprobtrans>0)[0]] #下一個要訪問的城市

            pathtable[i,j] = k

            unvisited.remove(k)
            #visited.add(k)

            length[i] += distmat[visiting][k]

            visiting = k

        length[i] += distmat[visiting][pathtable[i,0]] #螞蟻的路徑距離包括最後一個城市和第一個城市的距離


    #print length
    # 包含所有螞蟻的一個迭代結束後，統計本次迭代的若干統計引數

    lengthaver[iter] = length.mean()

    if iter == 0:
        lengthbest[iter] = length.min()
        pathbest[iter] = pathtable[length.argmin()].copy()      
    else:
        if length.min() > lengthbest[iter-1]:
            lengthbest[iter] = lengthbest[iter-1]
            pathbest[iter] = pathbest[iter-1].copy()

        else:
            lengthbest[iter] = length.min()
            pathbest[iter] = pathtable[length.argmin()].copy()    


    # 更新資訊素
    changepheromonetable = np.zeros((numcity,numcity))
    for i in range(numant):
        for j in range(numcity-1):
            changepheromonetable[pathtable[i,j]][pathtable[i,j+1]] += Q/distmat[pathtable[i,j]][pathtable[i,j+1]]

        changepheromonetable[pathtable[i,j+1]][pathtable[i,0]] += Q/distmat[pathtable[i,j+1]][pathtable[i,0]]

    pheromonetable = (1-rho)*pheromonetable + changepheromonetable


    iter += 1 #迭代次數指示器+1

    #觀察程式執行進度，該功能是非必須的
    if (iter-1)%20==0: 
        print iter-1


# 做出平均路徑長度和最優路徑長度        
fig,axes = plt.subplots(nrows=2,ncols=1,figsize=(12,10))
axes[0].plot(lengthaver,'k',marker = u'')
axes[0].set_title('Average Length')
axes[0].set_xlabel(u'iteration')

axes[1].plot(lengthbest,'k',marker = u'')
axes[1].set_title('Best Length')
axes[1].set_xlabel(u'iteration')
fig.savefig('Average_Best.png',dpi=500,bbox_inches='tight')
plt.close()


#作出找到的最優路徑圖
bestpath = pathbest[-1]

plt.plot(coordinates[:,0],coordinates[:,1],'r.',marker=u'$\cdot$')
plt.xlim([-100,2000])
plt.ylim([-100,1500])

for i in range(numcity-1):#
    m,n = bestpath[i],bestpath[i+1]
    print m,n
    plt.plot([coordinates[m][0],coordinates[n][0]],[coordinates[m][1],coordinates[n][1]],'k')
plt.plot([coordinates[bestpath[0]][0],coordinates[n][0]],[coordinates[bestpath[0]][1],coordinates[n][1]],'b')

ax=plt.gca()
ax.set_title("Best Path")
ax.set_xlabel('X axis')
ax.set_ylabel('Y_axis')

plt.savefig('Best Path.png',dpi=500,bbox_inches='tight')
plt.close()