import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# --- 1. 定义能量函数 (Energy Function) ---
# 我们设计一个不对称的双势阱：
# 一个深坑（全局最优），一个浅坑（局部最优）
def energy_function(x):
    # (x^2 - 1)^2 是标准的W形双井
    # + 0.3*x 用来倾斜它，使左边的坑比右边的深
    return (x**2 - 1)**2 + 0.3 * x

# --- 2. 准备数据 ---
x = np.linspace(-2.5, 2.5, 500)
E = energy_function(x)

# 找到真正的全局最小值，用于绘图标记
min_idx = np.argmin(E)
global_min_x = x[min_idx]
global_min_y = E[min_idx]

# --- 3. 设置绘图布局 ---
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 10), constrained_layout=True)

# 上图：能量景观 E(x) - 永远不变
ax1.plot(x, E, 'k-', linewidth=3, label='Energy $E(x)$')
ax1.scatter(global_min_x, global_min_y, color='gold', s=150, zorder=5, edgecolors='k', label='Global Min')
ax1.set_title("The Problem: Energy Landscape $E(x)$", fontsize=14)
ax1.set_ylabel("Energy")
ax1.grid(True, alpha=0.3)
ax1.legend()
ax1.set_xlim(-2.5, 2.5)

# 下图：概率分布 f(x) - 随温度变化
line, = ax2.plot([], [], 'r-', linewidth=3, alpha=0.8)
fill_poly = ax2.fill_between(x, np.zeros_like(x), np.zeros_like(x), color='red', alpha=0.3)
ax2.set_xlim(-2.5, 2.5)
ax2.set_ylabel("Probability Density $f(x)$")
ax2.set_xlabel("x")
ax2.grid(True, alpha=0.3)

# 动态文本：显示当前的 Lambda 和 Temperature
text_info = ax2.text(0.05, 0.9, '', transform=ax2.transAxes, fontsize=12, 
                     bbox=dict(facecolor='white', alpha=0.8))

# --- 4. 动画逻辑 ---
# Lambda 从 0.1 (高温) 变到 10.0 (低温)
# 我们用对数刻度，让高温阶段展示得慢一点，低温阶段快一点
lambdas = np.logspace(np.log10(0.1), np.log10(15.0), 100)

def init():
    line.set_data([], [])
    return line,

def update(frame_lambda):
    global fill_poly
    
    # === 核心物理计算 ===
    # 1. 计算玻尔兹曼因子 (未归一化概率)
    # e^(-lambda * E)
    unnormalized_prob = np.exp(-frame_lambda * E)
    
    # 2. 计算归一化常数 A (Partition Function Z)
    # 利用梯形法则进行数值积分
    integral_Z = np.trapezoid(unnormalized_prob, x)
    A = 1.0 / integral_Z
    
    # 3. 得到最终概率分布 f(x)
    # f(x) = A * e^(-lambda * E)
    pdf = A * unnormalized_prob
    # =================
    
    # 更新曲线
    line.set_data(x, pdf)
    
    # 更新填充区域 (需要移除旧的，画新的)
    fill_poly.remove()
    fill_poly = ax2.fill_between(x, 0, pdf, color='red', alpha=0.3)
    
    # 动态调整Y轴高度 (因为尖峰会越来越高)
    ax2.set_ylim(0, np.max(pdf) * 1.2)
    
    # 更新标题和文字
    T = 1 / frame_lambda
    ax2.set_title(f"The Stochastic Solution: Probability Distribution", fontsize=14)
    text_info.set_text(f"$\lambda$ = {frame_lambda:.2f} (Inverse Temp)\n$T$ = {T:.2f} (Temperature)")
    
    return line, fill_poly, text_info

# --- 5. 生成并保存动画 ---
print("正在生成动画，请稍候...")
ani = FuncAnimation(fig, update, frames=lambdas, init_func=init, blit=False)

# 保存为 GIF (需要安装 imagemagick 或使用 pillow)
ani.save('simulated_annealing.gif', writer=PillowWriter(fps=15))
print("✅ 动画已保存为 'simulated_annealing.gif'")

# 如果在 Jupyter Notebook 中，可以直接解开下面这行注释来显示
#plt.show()

import numpy as np
import matplotlib.pyplot as plt

# --- 1. 目标函数 (Energy Function) ---
def E(x):
    return x**2 - np.cos(np.pi * x)

# --- 2. 采样核心: Metropolis 准则 ---
# 这是一个"抽样机"，负责根据当前温度 T 生成样本
def sample_one_step(x_curr, T):
    # a. 提议 (Propose): 随机往左或往右迈一步
    x_next = x_curr + np.random.uniform(-0.5, 0.5)
    
    # b. 能量差 (Delta E)
    dE = E(x_next) - E(x_curr)
    
    # c. 接受/拒绝 (Accept/Reject)
    # 核心逻辑：如果新位置能量低，一定去；如果高，有一定概率去（取决于温度T）
    if dE < 0 or np.random.rand() < np.exp(-dE / T):
        return x_next # 接受移动
    else:
        return x_curr # 拒绝移动，待在原地

# --- 主流程: 对应你的笔记步骤 ---
def run_stochastic_optimization():
    # 初始化
    x = -2.5       # 起点 (故意选在一个较远的局部最优附近)
    T = 10.0       # T_max (高温)
    T_min = 0.01   # T_min (低温截止)
    alpha = 0.99   # 降温系数 (实际常用的降温方式)
    
    path = []      # 记录走过的路径
    temps = []     # 记录温度变化
    
    print(f"{'Step':<6} | {'Temp':<8} | {'Current x':<10} | {'Action'}")
    print("-" * 45)
    
    step = 0
    # 流程 4: 直到温度非常低 (Till a very low T)
    while T > T_min:
        
        # 流程 2 & 3: 采样 并 降温
        x = sample_one_step(x, T)
        
        # 记录数据
        path.append(x)
        temps.append(T)
        
        # 打印中间过程 (每隔200步)
        if step % 200 == 0:
            status = "Explore 🎲" if T > 1.0 else "Exploit 🎯"
            print(f"{step:<6} | {T:<8.4f} | {x:<10.4f} | {status}")
            
        # 降温 (Decrease T)
        T = T * alpha 
        step += 1
        
    # 流程 5: 统计均值 (Avg Samples)
    # 取最后 100 个低温样本的均值
    final_samples = path[-100:]
    estimated_min = np.mean(final_samples)
    
    print("-" * 45)
    print(f"✅ 最终估算结果: x = {estimated_min:.4f}")
    print(f"✅ 真实最小值: x = 0.0000 (大概率重合)")
    
    return path, temps

# --- 运行并可视化 ---
path, temps = run_stochastic_optimization()

# 画图展示粒子是如何"从乱跑"到"归位"的
plt.figure(figsize=(10, 6))
plt.plot(path, alpha=0.6, label='Particle Path')
plt.xlabel('Iterations (Time)')
plt.ylabel('Position x')
plt.title('Algorithm Flow: From Exploration (High T) to Exploitation (Low T)')
plt.axhline(0, color='r', linestyle='--', label='Global Min (x=0)')
plt.grid(True, alpha=0.3)
plt.legend()
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# --- 1. 定义 N 维目标函数 (Rastrigin Function) ---
def rastrigin(x):
    # x 是一个向量 (numpy array)
    # A * n + sum(x^2 - A * cos(2*pi*x))
    A = 10
    n = len(x)
    return A * n + np.sum(x**2 - A * np.cos(2 * np.pi * x))

# --- 2. N 维邻域生成 (Proposal) ---
def get_neighbor(x_curr, step_size=0.5):
    # 关键点：我们在 N 维空间中随机游走
    # size=len(x_curr) 保证了生成的扰动向量和 x 维度一致
    perturbation = np.random.uniform(-step_size, step_size, size=len(x_curr))
    return x_curr + perturbation

# --- 3. 模拟退火主程序 ---
def simulated_annealing_nd(n_dim=2, n_iter=2000):
    # 初始化：在一个范围内随机生成起点 (-5.12 到 5.12 是 Rastrigin 的标准定义域)
    current_x = np.random.uniform(-5.12, 5.12, size=n_dim)
    current_E = rastrigin(current_x)
    
    # 记录最佳解 (Best So Far)，防止跑丢了
    best_x = current_x.copy()
    best_E = current_E
    
    # 温度设置
    T = 100.0
    T_min = 1e-4
    alpha = 0.99  # 降温系数
    
    path = [current_x] # 记录路径用于画图
    energy_history = [current_E]

    print(f"开始 {n_dim} 维优化...")
    print(f"起点: {np.round(current_x, 2)}, Energy: {current_E:.2f}")

    iter_count = 0
    while T > T_min and iter_count < n_iter:
        # 1. 提议新位置 (向量加法)
        new_x = get_neighbor(current_x)
        
        # 防止跑出定义域 (Rastrigin 通常限制在 [-5.12, 5.12])
        new_x = np.clip(new_x, -5.12, 5.12)
        
        new_E = rastrigin(new_x)
        
        # 2. 计算能量差
        dE = new_E - current_E
        
        # 3. Metropolis 准则 (和 1 维一模一样)
        if dE < 0 or np.random.rand() < np.exp(-dE / T):
            current_x = new_x
            current_E = new_E
            
            # 更新历史最佳
            if current_E < best_E:
                best_x = current_x.copy()
                best_E = current_E
        
        path.append(current_x)
        energy_history.append(current_E)
        
        T *= alpha
        iter_count += 1
        
    print(f"结束. 最终位置: {np.round(best_x, 4)}")
    print(f"最终能量: {best_E:.6f} (理论最优是 0.0)")
    
    return np.array(path), energy_history, best_x

# --- 运行: 这里我们设为 2 维以便画图，但算法支持 N 维 ---
DIMENSION = 2
path, energies, final_sol = simulated_annealing_nd(n_dim=DIMENSION, n_iter=3000)

# --- 4. 可视化 (仅适用于 2D) ---
if DIMENSION == 2:
    plt.figure(figsize=(12, 5))
    
    # 子图 1: 地形图和路径
    plt.subplot(1, 2, 1)
    x_grid = np.linspace(-5.12, 5.12, 100)
    y_grid = np.linspace(-5.12, 5.12, 100)
    X, Y = np.meshgrid(x_grid, y_grid)
    # 计算网格上每个点的 Z 值
    Z = np.zeros_like(X)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Z[i,j] = rastrigin(np.array([X[i,j], Y[i,j]]))
            
    plt.contourf(X, Y, Z, levels=50, cmap='viridis', alpha=0.7)
    plt.colorbar(label='Energy')
    
    # 画路径：起点是白色，终点是红色
    plt.plot(path[:, 0], path[:, 1], 'w-', linewidth=0.5, alpha=0.6)
    plt.scatter(path[0, 0], path[0, 1], c='white', s=50, label='Start')
    plt.scatter(final_sol[0], final_sol[1], c='red', marker='*', s=200, label='End')
    plt.legend()
    plt.title(f"2D Rastrigin Optimization Path\n(Escaping many local minima)")
    
    # 子图 2: 能量下降曲线
    plt.subplot(1, 2, 2)
    plt.plot(energies)
    plt.yscale('log') # 用对数坐标看，因为后期下降很微小
    plt.xlabel('Iteration')
    plt.ylabel('Energy (Log Scale)')
    plt.title('Energy Minimization Process')
    plt.grid(True, which="both", ls="--")
    
    plt.tight_layout()
    plt.show()