import numpy as np
import matplotlib.pyplot as plt
# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 距离计算函数
def calculate_distance(citys):
    n = len(citys)
    D = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            D[i, j] = np.sqrt((citys[i, 0] - citys[j, 0])**2 + (citys[i, 1] - citys[j, 1])**2)
            D[j, i] = D[i, j]
        D[i, i] = 1e-4  # 对角线赋值为一个很小的数，避免0的倒数
    return D

# 参数与城市位置初始化
citys = np.array([
    [18.47, 95.1], [16.47, 94.64], [20.09, 94.54], [14.39, 93.37],
    [25.23, 97.24], [22.00, 93.05], [23.47, 92.02], [16.2, 96.29],
    [17.3, 97.38], [13.05, 98.12], [15.53, 97.38], [24.52, 95.59],
    [16.41, 97.13], [15.09, 92.55]
])
D = calculate_distance(citys)
n = len(D)

NC_max = 200  # 最大迭代次数
m = int(1.5 * n)  # 蚂蚁数量
alpha = 1  # 信息素重要程度参数
beta = 4  # 启发因子重要程度参数
rho = 0.2  # 信息素挥发因子
Q = 20  # 信息素释放系数

Eta = 1 / D  # 启发函数
Tau = np.ones((n, n))  # 信息素矩阵
Table = np.zeros((m, n), dtype=int)  # 路径记录表

rBest = np.zeros((NC_max, n), dtype=int)  # 各代最佳路线
lBest = np.full(NC_max, np.inf)  # 各代最佳路线的长度
lAverage = np.zeros(NC_max)  # 各代平均路线长度

# 迭代过程
for NC in range(NC_max):
    # 第1步：初始化每只蚂蚁的起点
    start = np.random.choice(n, m)
    Table[:, 0] = start

    # 第2步：路径选择
    citys_index = np.arange(n)
    for i in range(m):
        for j in range(1, n):
            tabu = Table[i, :j]
            allow_index = ~np.isin(citys_index, tabu)
            Allow = citys_index[allow_index]

            P = np.array([Tau[tabu[-1], k] ** alpha * Eta[tabu[-1], k] ** beta for k in Allow])
            P = P / P.sum()  # 归一化

            target = np.random.choice(Allow, p=P)
            Table[i, j] = target

    # 第3步：计算路径距离
    length = np.zeros(m)
    for i in range(m):
        route = Table[i]
        length[i] = sum(D[route[j], route[j + 1]] for j in range(n - 1)) + D[route[-1], route[0]]

    # 第4步：更新最佳路径
    min_length = length.min()
    min_index = length.argmin()
    lBest[NC] = min(min_length, lBest[NC - 1] if NC > 0 else np.inf)
    lAverage[NC] = length.mean()
    if lBest[NC] == min_length:
        rBest[NC] = Table[min_index]

    # 第5步：更新信息素矩阵
    Delta_tau = np.zeros((n, n))
    for i in range(m):
        for j in range(n - 1):
            Delta_tau[Table[i, j], Table[i, j + 1]] += Q / length[i]
        Delta_tau[Table[i, -1], Table[i, 0]] += Q / length[i]
    Tau = (1 - rho) * Tau + Delta_tau

    # 重置路径记录表
    Table = np.zeros((m, n), dtype=int)

# 结果显示
shortest_Length = lBest.min()
shortest_index = lBest.argmin()
shortest_Route = rBest[shortest_index]

print(f"最短距离: {shortest_Length}")
print(f"最短路径: {np.append(shortest_Route, shortest_Route[0])}")

# 绘图显示最优路径
plt.figure(1)
plt.plot(citys[shortest_Route, 0], citys[shortest_Route, 1], 'o-')
plt.plot([citys[shortest_Route[0], 0], citys[shortest_Route[-1], 0]],
         [citys[shortest_Route[0], 1], citys[shortest_Route[-1], 1]], 'o-')
plt.grid(True)
for i, city in enumerate(citys):
    plt.text(city[0], city[1], f" {i+1}")
plt.xlabel("城市位置横坐标")
plt.ylabel("城市位置纵坐标")
plt.title(f"蚁群算法优化路径(最短距离: {shortest_Length})")

# 绘制距离变化图
plt.figure(2)
plt.plot(range(1, NC_max + 1), lBest, 'b', label='最短距离')
plt.plot(range(1, NC_max + 1), lAverage, 'r', label='平均距离')
plt.legend()
plt.xlabel("迭代次数")
plt.ylabel("距离")
plt.title("各代最短距离与平均距离对比")

plt.show()