import numpy as np
import pandas as pd
from scipy import spatial
from sko.IA import IA_TSP
import matplotlib.pyplot as plt

"""
numpy: 1.24.3
pandas: 1.5.3
scipy: 1.10.1
scikit-opt 0.6.6
sko 0.5.7
matplotlib: 3.7.5
"""

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 参数设置
num_points = 100
size_pop = 100
max_iter = 5000

# 读取坐标信息
points = pd.read_excel('数据.xlsx')
points_coordinate = np.array(points.values.tolist())
distance_matrix = spatial.distance.cdist(points_coordinate, points_coordinate, metric='euclidean')

def cal_total_distance(routine):
    num_points, = routine.shape
    return sum([distance_matrix[routine[i % num_points], routine[(i + 1) % num_points]] for i in range(num_points)])

# 初始化免疫优化算法算法
ia_tsp = IA_TSP(func=cal_total_distance, n_dim=num_points, size_pop=size_pop, max_iter=max_iter, prob_mut=0.2,
                T=0.7, alpha=0.95)

# 创建图形
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
fig.suptitle('基于免疫优化算法求解旅行商问题')

# 路线图设置
ax1.set_title('旅行商路线图')
ax1.set_xlim(0, 100)
ax1.set_ylim(0, 100)
ax1.set_xlabel('X轴')
ax1.set_ylabel('Y轴')
ax1.set_aspect('equal')

# 收敛曲线图设置
ax2.set_title('收敛曲线图')
ax2.set_xlabel('迭代次数')
ax2.set_ylabel('总里程数')
ax2.grid(True)
ax2.set_xlim(0, max_iter)
ax2.set_aspect('auto')

# 初始化图形元素
line, = ax1.plot([], [], 'o-r', markersize=4, linewidth=1)
scatter = ax1.scatter(points_coordinate[:, 0], points_coordinate[:, 1], c='r', s=20)
convergence_line, = ax2.plot([], [], 'b-', linewidth=2)
best_distance_text = ax2.text(0.75, 0.75, '', transform=ax2.transAxes)

# 计算y轴范围
initial_distance = cal_total_distance(np.arange(num_points))
ax2.set_ylim(0, initial_distance * 1.1)

# 生成最后迭代结果
best_points, best_distance = ia_tsp.run()
best_points_ = np.concatenate([best_points, [best_points[0]]])
best_points_coordinate = points_coordinate[best_points_, :]
ax1.plot(best_points_coordinate[:, 0], best_points_coordinate[:, 1], 'o-r')
ax2.plot(ia_tsp.generation_best_Y)
best_distance_text.set_text(f'迭代次数：{max_iter}\n\n总里程数：{best_distance[-1]:.2f}')
plt.tight_layout()
plt.show()