import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import time


# Ackley函數(shù)
defackley(x):
    a = 20
    b = 0.2
    c = 2 * np.pi
    d = len(x)
    sum_sq = sum([xi**2for xi in x])
    sum_cos = sum([np.cos(c * xi) for xi in x])
    term1 = -a * np.exp(-b * np.sqrt(sum_sq / d))
    term2 = -np.exp(sum_cos / d)
return term1 + term2 + a + np.exp(1)

# 麻雀搜索算法
defsparrow_search_algorithm_with_history(obj_func, dim, lb, ub, max_iter, n_sparrows):
# 初始化
    positions = np.random.uniform(lb, ub, (n_sparrows, dim))
    fitness = np.array([obj_func(p) for p in positions])


# 歷史記錄
    history = {
'positions': [positions.copy()],
'best_positions': [],
'best_fitness': [],
'worst_positions': []
    }


# 找出當(dāng)前最佳和最差
    best_index = np.argmin(fitness)
    worst_index = np.argmax(fitness)
    best_position = positions[best_index].copy()
    best_fitness = fitness[best_index]
    worst_position = positions[worst_index].copy()


    history['best_positions'].append(best_position.copy())
    history['best_fitness'].append(best_fitness)
    history['worst_positions'].append(worst_position.copy())


# 迭代參數(shù)
    ST = 0.6# 安全閾值
    PD = 0.7# 發(fā)現(xiàn)者比例
    SD = 0.2# 警戒者比例


    n_p = int(n_sparrows * PD)
    n_s = int(n_sparrows * SD)


# 迭代過程
for t inrange(max_iter):
if t % 10 == 0or t == max_iter-1or t == 0:
            print(f"迭代 {t+1:3d}/{max_iter} | 當(dāng)前最佳適應(yīng)度: {best_fitness:.6f}")


# 更新發(fā)現(xiàn)者位置
        R2 = np.random.rand()
for i inrange(n_p):
if R2 < ST:
                alpha = np.random.randn()
for d inrange(dim):
                    positions[i,d] *= np.exp(-alpha * t / max_iter)
else:
                Q = np.random.randn()
for d inrange(dim):
                    positions[i,d] += Q * np.random.rand()
            positions[i] = np.clip(positions[i], lb, ub)


# 更新跟隨者位置
for i inrange(n_p, n_sparrows):
if i > n_sparrows / 2:
                Q = np.random.randn()
for d inrange(dim):
                    positions[i,d] = Q * np.exp((worst_position[d] - positions[i,d]) / i**2)
else:
                A = np.random.randint(-1, 2, dim)
                A_plus = A.T @ np.linalg.pinv(A @ A.T)
for d inrange(dim):
                    positions[i,d] = best_position[d] + np.abs(positions[i,d] - best_position[d]) * A_plus[d]
            positions[i] = np.clip(positions[i], lb, ub)


# 更新警戒者位置
for i inrange(n_s):
if fitness[i] > best_fitness:
                beta = np.random.randn()
for d inrange(dim):
                    positions[i,d] = best_position[d] + beta * np.abs(positions[i,d] - best_position[d])
elif fitness[i] == best_fitness:
                K = np.random.rand() * 2 - 1
for d inrange(dim):
                    positions[i,d] = positions[i,d] + K * (np.abs(positions[i,d] - worst_position[d]) / 
                                    (fitness[i] - worst_position[d] + 1e-50))
            positions[i] = np.clip(positions[i], lb, ub)


# 重新計算適應(yīng)度
        fitness = np.array([obj_func(p) for p in positions])


# 更新全局最佳和最差
        current_best_index = np.argmin(fitness)
        current_worst_index = np.argmax(fitness)


if fitness[current_best_index] < best_fitness:
            best_position = positions[current_best_index].copy()
            best_fitness = fitness[current_best_index]


        worst_position = positions[current_worst_index].copy()


# 記錄歷史
        history['positions'].append(positions.copy())
        history['best_positions'].append(best_position.copy())
        history['best_fitness'].append(best_fitness)
        history['worst_positions'].append(worst_position.copy())


return history

# ==================================================
print("準(zhǔn)備運(yùn)行麻雀搜索算法...")
print("="*50)
print("麻雀搜索算法 (SSA) 開始運(yùn)行")
print(f"麻雀數(shù)量: {30}")
print(f"最大迭代次數(shù): {50}")
print(f"搜索空間維度: {2}")
print(f"搜索范圍: [-5, 5]")
print("="*50)
time.sleep(1)


# 參數(shù)設(shè)置
dim = 2
lb = -5
ub = 5
max_iter = 50
n_sparrows = 30


# 運(yùn)行算法并記錄歷史
history = sparrow_search_algorithm_with_history(ackley, dim, lb, ub, max_iter, n_sparrows)


# 最終結(jié)果
print("="*50)
print("優(yōu)化完成!")
print(f"找到的最佳解: [{history['best_positions'][-1][0]:.4f}{history['best_positions'][-1][1]:.4f}]")
print(f"最佳適應(yīng)度值: {history['best_fitness'][-1]:.6f}")
print("="*50)
time.sleep(1)
print("生成可視化結(jié)果...")
time.sleep(1)
# ==================================================


# 準(zhǔn)備可視化數(shù)據(jù)
x = np.linspace(lb, ub, 100)
y = np.linspace(lb, ub, 100)
X, Y = np.meshgrid(x, y)
Z = np.zeros_like(X)
for i inrange(X.shape[0]):
for j inrange(X.shape[1]):
        Z[i,j] = ackley([X[i,j], Y[i,j]])


# 創(chuàng)建圖形
plt.rcParams['font.sans-serif'] = ['SimHei']  # 設(shè)置中文顯示
plt.rcParams['axes.unicode_minus'] = False# 解決負(fù)號顯示問題


fig = plt.figure(figsize=(18, 6), dpi=100)
fig.suptitle('麻雀搜索算法動態(tài)可視化', fnotallow=16)


# 統(tǒng)一子圖尺寸
gs = fig.add_gridspec(2, 3, width_ratios=[1, 1, 1], height_ratios=[1, 1])


# 3D曲面圖
ax1 = fig.add_subplot(gs[:, 0], projectinotallow='3d')
surf = ax1.plot_surface(X, Y, Z, cmap='viridis', alpha=0.6)
fig.colorbar(surf, ax=ax1, shrink=0.6, aspect=10, label='函數(shù)值')
scatter = ax1.scatter([], [], [], c='red', s=50, label='麻雀位置')
best_scatter = ax1.scatter([], [], [], c='blue', marker='*', s=200, label='最優(yōu)解')
ax1.set_title('3D函數(shù)曲面與種群分布', fnotallow=12)
ax1.set_xlabel('x1', fnotallow=10)
ax1.set_ylabel('x2', fnotallow=10)
ax1.set_zlabel('f(x)', fnotallow=10)
ax1.legend(loc='upper right', fnotallow=8)


# 2D等高線圖
ax2 = fig.add_subplot(gs[:, 1])
contour = ax2.contourf(X, Y, Z, levels=50, cmap='viridis')
fig.colorbar(contour, ax=ax2, shrink=0.6, aspect=10, label='函數(shù)值')
scatter2d = ax2.scatter([], [], c='red', s=30, label='麻雀位置')
best_scatter2d = ax2.scatter([], [], c='blue', marker='*', s=100, label='最優(yōu)解')
ax2.set_title('2D等高線與種群分布', fnotallow=12)
ax2.set_xlabel('x1', fnotallow=10)
ax2.set_ylabel('x2', fnotallow=10)
ax2.legend(loc='upper right', fnotallow=8)


# 收斂曲線
ax3 = fig.add_subplot(gs[0, 2])
convergence_line, = ax3.plot([], [], 'b-', linewidth=2, label='最佳適應(yīng)度')
current_point = ax3.scatter([], [], c='red', s=50, label='當(dāng)前值')
ax3.set_title('適應(yīng)度收斂曲線', fnotallow=12)
ax3.set_xlabel('迭代次數(shù)', fnotallow=10)
ax3.set_ylabel('適應(yīng)度值', fnotallow=10)
ax3.grid(True, linestyle='--', alpha=0.6)
ax3.set_xlim(0, max_iter)
ax3.set_ylim(0, max(history['best_fitness']))
ax3.legend(loc='upper right', fnotallow=8)


# 參數(shù)顯示
ax4 = fig.add_subplot(gs[1, 2])
ax4.axis('off')
info_text = ax4.text(0.1, 0.5, '', fnotallow=10, bbox=dict(facecolor='white', alpha=0.8))


plt.tight_layout()


# 動畫更新函數(shù)
defupdate(frame):
# 更新3D圖
    current_pos = history['positions'][frame]
    current_z = np.array([ackley(p) for p in current_pos])
    scatter._offsets3d = (current_pos[:,0], current_pos[:,1], current_z)


    best_pos = history['best_positions'][frame]
    best_z = ackley(best_pos)
    best_scatter._offsets3d = ([best_pos[0]], [best_pos[1]], [best_z])


# 更新2D圖
    scatter2d.set_offsets(current_pos)
    best_scatter2d.set_offsets([best_pos])


# 更新收斂曲線
    x_data = range(frame+1)
    y_data = history['best_fitness'][:frame+1]
    convergence_line.set_data(x_data, y_data)
    current_point.set_offsets([[frame, history['best_fitness'][frame]]])


# 更新文本信息
    info = f"迭代次數(shù): {frame}\n"
    info += f"最佳適應(yīng)度: {history['best_fitness'][frame]:.6f}\n"
    info += f"最佳位置: [{best_pos[0]:.4f}, {best_pos[1]:.4f}]\n"
    info += f"麻雀數(shù)量: {len(current_pos)}\n"
    info += f"發(fā)現(xiàn)者比例: 70%\n"
    info += f"警戒者比例: 20%"
    info_text.set_text(info)


return scatter, best_scatter, scatter2d, best_scatter2d, convergence_line, current_point, info_text


# 創(chuàng)建動畫
ani = FuncAnimation(fig, update, frames=max_iter+1, interval=500, blit=True)


# 顯示動畫
plt.close()
HTML(ani.to_jshtml())