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
import math


# 設置中文顯示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# Ackley函數
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)

# Logistic混沌映射
deflogistic_map(x, mu=4.0):
    return mu * x * (1 - x)


# 混沌優化算法
defchaos_optimization(obj_func, dim, lb, ub, max_iter, chaos_iter=1000):
    """
    參數:
    - obj_func: 目標函數
    - dim: 問題維度
    - lb: 下界
    - ub: 上界
    - max_iter: 最大迭代次數
    - chaos_iter: 混沌迭代次數
    """
    # 初始化混沌變量
    chaos_vars = np.random.rand(dim)
    
    # 生成混沌序列
    chaos_sequence = []
    for _ inrange(chaos_iter):
        chaos_vars = logistic_map(chaos_vars)
        chaos_sequence.append(chaos_vars.copy())
    chaos_sequence = np.array(chaos_sequence)
    
    # 將混沌序列映射到搜索空間
    search_points = lb + (ub - lb) * chaos_sequence
    
    # 評估初始解
    fitness = np.array([obj_func(p) for p in search_points])
    best_idx = np.argmin(fitness)
    best_position = search_points[best_idx].copy()
    best_fitness = fitness[best_idx]
    
    # 記錄歷史
    history = {
        'positions': [search_points.copy()],
        'best_position': [best_position.copy()],
        'best_fitness': [best_fitness],
        'current_iter': [0]
    }
    
    # 算法開始
    print("="*50)
    print("混沌優化算法(COA)開始運行")
    print(f"搜索空間維度: {dim}")
    print(f"搜索范圍: [{lb}, {ub}]")
    print(f"最大迭代次數: {max_iter}")
    print(f"混沌迭代次數: {chaos_iter}")
    print(f"初始最佳適應度: {best_fitness:.6f}")
    print("="*50)
    time.sleep(1)
    
    # 二次載波搜索
    foriterinrange(max_iter):
        # 縮小搜索范圍
        current_lb = np.maximum(lb, best_position - (ub - lb) * 0.9**(iter+1))
        current_ub = np.minimum(ub, best_position + (ub - lb) * 0.9**(iter+1))
        
        # 生成新的混沌序列
        new_chaos_vars = np.random.rand(dim)
        new_search_points = []
        for _ inrange(chaos_iter):
            new_chaos_vars = logistic_map(new_chaos_vars)
            new_point = current_lb + (current_ub - current_lb) * new_chaos_vars
            new_search_points.append(new_point)
        new_search_points = np.array(new_search_points)
        
        # 評估新解
        new_fitness = np.array([obj_func(p) for p in new_search_points])
        current_best_idx = np.argmin(new_fitness)
        current_best_position = new_search_points[current_best_idx].copy()
        current_best_fitness = new_fitness[current_best_idx]
        
        # 更新最優解
        if current_best_fitness < best_fitness:
            best_position = current_best_position.copy()
            best_fitness = current_best_fitness
        
        # 記錄歷史
        history['positions'].append(new_search_points.copy())
        history['best_position'].append(best_position.copy())
        history['best_fitness'].append(best_fitness)
        history['current_iter'].append(iter+1)
        
        # 打印進度
        ifiter % 10 == 0oriter == max_iter-1:
            print(f"迭代 {iter+1:3d}/{max_iter} | 當前最佳適應度: {best_fitness:.6f}")
    
    print("="*50)
    print("優化完成!")
    print(f"找到的最佳解: {best_position}")
    print(f"最佳適應度值: {best_fitness:.6f}")
    print("="*50)
    time.sleep(1)
    print("生成可視化結果...")
    time.sleep(1)
    
    return history

# 參數設置
dim = 2
lb = -5
ub = 5
max_iter = 50
chaos_iter = 500


# 運行算法
history = chaos_optimization(ackley, dim, lb, ub, max_iter, chaos_iter)


# 準備可視化數據
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]])


# 創建可視化圖形
fig = plt.figure(figsize=(18, 6), dpi=100)
fig.suptitle('混沌優化算法優化過程', fontsize=16)


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


# 3D曲面圖
ax1 = fig.add_subplot(gs[:, 0], projection='3d')
surf = ax1.plot_surface(X, Y, Z, cmap='viridis', alpha=0.6)
fig.colorbar(surf, ax=ax1, shrink=0.6, aspect=10, label='函數值')
scatter = ax1.scatter([], [], [], c='red', s=10, alpha=0.5, label='混沌搜索點')
best_scatter = ax1.scatter([], [], [], c='blue', marker='*', s=200, label='最優解')
ax1.set_title('3D函數曲面與混沌搜索', fontsize=12)
ax1.set_xlabel('x1', fontsize=10)
ax1.set_ylabel('x2', fontsize=10)
ax1.set_zlabel('f(x)', fontsize=10)
ax1.legend(loc='upper right', fontsize=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='函數值')
scatter2d = ax2.scatter([], [], c='red', s=10, alpha=0.5, label='混沌搜索點')
best_scatter2d = ax2.scatter([], [], c='blue', marker='*', s=100, label='最優解')
search_area = plt.Rectangle((0,0), 0, 0, color='yellow', alpha=0.3, label='當前搜索區域')
ax2.add_patch(search_area)
ax2.set_title('2D等高線與混沌搜索', fontsize=12)
ax2.set_xlabel('x1', fontsize=10)
ax2.set_ylabel('x2', fontsize=10)
ax2.legend(loc='upper right', fontsize=8)


# 收斂曲線
ax3 = fig.add_subplot(gs[0, 2])
convergence_line, = ax3.plot([], [], 'b-', linewidth=2, label='最佳適應度')
current_point = ax3.scatter([], [], c='red', s=50, label='當前值')
ax3.set_title('適應度收斂曲線', fontsize=12)
ax3.set_xlabel('迭代次數', fontsize=10)
ax3.set_ylabel('適應度值', fontsize=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', fontsize=8)


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


plt.tight_layout()


# 修改更新函數
defupdate(frame):
    # 只顯示部分點避免過于密集
    display_points = history['positions'][frame][::10]
    
    # 更新3D圖
    current_z = np.array([ackley(p) for p in display_points])
    scatter._offsets3d = (display_points[:,0], display_points[:,1], current_z)
    
    best_pos = history['best_position'][frame]
    best_z = ackley(best_pos)
    best_scatter._offsets3d = ([best_pos[0]], [best_pos[1]], [best_z])
    
    # 更新2D圖
    scatter2d.set_offsets(display_points)
    best_scatter2d.set_offsets([best_pos])
    
    # 初始化搜索區域
    current_iter = history['current_iter'][frame]
    if current_iter == 0:
        # 第一幀使用全局搜索范圍
        current_lb = np.array([lb, lb])
        current_ub = np.array([ub, ub])
    else:
        # 后續幀縮小搜索范圍
        current_lb = np.maximum(lb, best_pos - (ub - lb) * 0.9**current_iter)
        current_ub = np.minimum(ub, best_pos + (ub - lb) * 0.9**current_iter)
    
    # 更新搜索區域顯示
    search_area.set_xy((current_lb[0], current_lb[1]))
    search_area.set_width(current_ub[0] - current_lb[0])
    search_area.set_height(current_ub[1] - current_lb[1])
    
    # 更新收斂曲線
    x_data = range(current_iter+1)
    y_data = history['best_fitness'][:current_iter+1]
    convergence_line.set_data(x_data, y_data)
    current_point.set_offsets([[current_iter, history['best_fitness'][current_iter]]])
    
    # 更新文本信息
    info = f"迭代次數: {current_iter}\n"
    info += f"最佳適應度: {history['best_fitness'][current_iter]:.6f}\n"
    info += f"最佳位置: [{best_pos[0]:.4f}, {best_pos[1]:.4f}]\n"
    info += f"搜索區域: [{current_lb[0]:.2f}, {current_ub[0]:.2f}] x [{current_lb[1]:.2f}, {current_ub[1]:.2f}]\n"
    info += f"混沌迭代次數: {chaos_iter}\n"
    info += f"當前搜索點數: {len(history['positions'][frame])}"
    info_text.set_text(info)
    
    return scatter, best_scatter, scatter2d, best_scatter2d, search_area, convergence_line, current_point, info_text


# 創建動畫
ani = FuncAnimation(fig, update, frames=len(history['positions']), interval=500, blit=True)


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