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()

import numpy as np
import matplotlib.pyplot as plt

# ==========================================
# 1. 定义目标函数 (Target Function)
# ==========================================
# 这是一个构造出来的多峰函数，有多个坑
def target_func(x):
    # 参数 (对应 MATLAB 中的 a, b, c)
    a = np.array([-40, -20, -40])
    b = np.array([150, 100, 300])
    c = np.array([-10, -40, 40])
    
    y = 0
    for i in range(3):
        y += a[i] / (b[i] + (x + c[i])**2)
    return y

# 绘图范围
xmin, xmax = -100, 100
x_vec = np.linspace(xmin, xmax, 1000)
y_vec = target_func(x_vec)

plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.title('Target Function (Optimization Landscape)')
plt.grid(True)
plt.show()

# ==========================================
# 2. 定义冷却计划 (Temperature Schedule)
# ==========================================
# 我们希望：
# - 初始温度 Tin 时：即使是最大的能量差，也有 80% 的概率被接受 (高探索)
# - 结束温度 Tend 时：即使是很小的能量差(1e-3)，也只有 5% 的概率被接受 (高锁定)

y_max, y_min = np.max(y_vec), np.min(y_vec)
delta_E_max = y_max - y_min

# 根据公式反推温度
# P = exp(-dE / T)  =>  ln(P) = -dE / T  =>  T = -dE / ln(P)
T_in = -delta_E_max / np.log(0.8)
T_end = -1e-3 / np.log(0.05)

n_iter = 10000

# 几何冷却法则: T(k) = T_in * exp(-tau * k)
# 在 k = n_iter-1 时 T = T_end
tau = -np.log(T_end / T_in) / (n_iter - 1)

iterations = np.arange(n_iter)
T_schedule = T_in * np.exp(-tau * iterations)

plt.figure(figsize=(6, 4))
plt.plot(iterations, T_schedule)
plt.title('Temperature Law (Geometric Decay)')
plt.xlabel('Iteration')
plt.ylabel('Temperature')
plt.grid(True)
plt.show()

# ==========================================
# 3. 模拟退火主循环 (Simulated Annealing)
# ==========================================
print("开始模拟退火...")
sx = np.zeros(n_iter)
alpha_hist = np.zeros(n_iter)

# 初始点 (随机)
current_x = np.random.randn() * 10
sx[0] = current_x
proposal_std = 50 # 提议分布的标准差 (对应 MATLAB 中的 sd=100)

for i in range(1, n_iter):
    # 1. 提议 (Proposal): 随机游走
    # 并在边界内截断 (Truncated Normal)
    while True:
        candidate = current_x + np.random.randn() * proposal_std
        if xmin <= candidate <= xmax:
            break
            
    # 2. 计算能量差
    E_curr = target_func(current_x)
    E_cand = target_func(candidate)
    dE = E_cand - E_curr # 我们想最小化，所以能量越低越好
    
    # 3. 计算接受率 (Metropolis Criterion)
    # 当前温度
    T_curr = T_schedule[i]
    
    # 概率公式 alpha = exp(-dE / T)
    # 如果 dE < 0 (变好了)，指数为正，exp > 1，min 取 1 (必接受)
    # 如果 dE > 0 (变差了)，指数为负，概率 < 1
    alpha = min(1, np.exp(-dE / T_curr))
    
    alpha_hist[i] = alpha
    
    # 4. 接受/拒绝
    if np.random.rand() < alpha:
        current_x = candidate
    # else: current_x 保持不变
    
    sx[i] = current_x

# 随机解结果 (取后20%的均值)
t_burn = int(0.8 * n_iter)
x_sol_stochastic = np.mean(sx[t_burn:])
y_sol_stochastic = target_func(x_sol_stochastic)

print(f"随机解 (SA): x = {x_sol_stochastic:.4f}, y = {y_sol_stochastic:.4f}")
# Plot on the image
plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.plot(x_sol_stochastic, y_sol_stochastic, 'ro')
plt.title('Target Function With Solution')
plt.grid(True)
plt.show()

# ==========================================
# 4. 确定性解法对比 (Newton's Method)
# ==========================================
# 牛顿法: x_new = x - f'(x) / f''(x)
# 我们故意给它一个离全局最优有点距离，但在某个局部最优附近的起点
x_newton = -10.0 
newton_path = [x_newton]

# 计算导数 (用于牛顿法) - 使用中心差分法近似
def get_derivatives(f, x, h=1e-5):
    # 一阶导数 f'
    df = (f(x + h) - f(x - h)) / (2 * h)
    # 二阶导数 f''
    ddf = (f(x + h) - 2*f(x) + f(x - h)) / (h**2)
    return df, ddf

for _ in range(100):
    df, ddf = get_derivatives(target_func, x_newton)
    if abs(ddf) < 1e-6: break # 防止除零
    
    step = df / ddf
    x_newton = x_newton - step
    newton_path.append(x_newton)
    
    if abs(step) < 1e-6: break # 收敛

y_sol_deterministic = target_func(x_newton)
print(f"确定性解 (Newton): x = {x_newton:.4f}, y = {y_sol_deterministic:.4f}")

# Plot on the image
plt.figure(figsize=(10, 4))
plt.plot(x_vec, y_vec, linewidth=2)
plt.plot(x_newton, y_sol_deterministic, 'ro')
plt.title('Target Function With Solution)')
plt.grid(True)
plt.show()

# ==========================================
# 5. 最终结果可视化
# ==========================================
plt.figure(figsize=(12, 6))
plt.plot(x_vec, y_vec, label='Target Function', color='blue', alpha=0.5)

# 画出 SA 的采样点 (用颜色深浅表示时间)
plt.scatter(sx[::10], target_func(sx[::10]), c=np.arange(0, n_iter, 10), 
            cmap='Wistia', alpha=0.5, s=10, label='SA Trajectory')

# 标记最终解
plt.plot(x_sol_stochastic, y_sol_stochastic, 'r*', markersize=20, label='SA Solution (Global)')
plt.plot(x_newton, y_sol_deterministic, 'kx', markersize=15, markeredgewidth=3, label='Newton Solution (Local)')

plt.colorbar(label='Iteration Time')
plt.legend()
plt.title('Global Optimization: Simulated Annealing vs Newton Method')
plt.xlim(xmin, xmax)
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# 绘图准备
x_vec = np.linspace(-100, 100, 1000)
y_vec = target_func(x_vec)

# ==========================================
# Phase 1: 模拟退火 (粗调 / Global Search)
# ==========================================
print("🚀 Phase 1: Simulated Annealing (Global Search)...")

# 1. 自动计算温度 (根据你的公式)
y_max, y_min = np.max(y_vec), np.min(y_vec)
delta_E_max = y_max - y_min

# 初始: 80% 概率接受最大跳跃; 结束: 5% 概率接受微小跳跃(1e-3)
T_in = -delta_E_max / np.log(0.8)
T_end = -1e-3 / np.log(0.05)

n_iter = 5000 # 迭代次数
# 几何冷却参数 tau
tau = -np.log(T_end / T_in) / (n_iter - 1)

# 初始化
current_x = np.random.uniform(-90, 90) # 随机起点
sx = np.zeros(n_iter)
proposal_std = 20 # 探索步长

# SA 主循环
for i in range(n_iter):
    # 计算当前温度
    T_curr = T_in * np.exp(-tau * i)
    
    # 提议新点 (限制在范围内)
    while True:
        candidate = current_x + np.random.randn() * proposal_std
        if -100 <= candidate <= 100: break
    
    # 能量差
    dE = target_func(candidate) - target_func(current_x)
    
    # Metropolis 准则
    if dE < 0 or np.random.rand() < np.exp(-dE / T_curr):
        current_x = candidate
        
    sx[i] = current_x

# SA 结果: 取最后 10% 样本的平均值 (平稳分布的期望)
t_burn = int(0.9 * n_iter)
x_sa_approx = np.mean(sx[t_burn:])
y_sa_approx = target_func(x_sa_approx)

print(f"✅ SA 完成。粗略最优解: x ≈ {x_sa_approx:.4f}")


# ==========================================
# Phase 2: 牛顿法 (精调 / Local Refinement)
# ==========================================
print("\n🎯 Phase 2: Hybrid Newton Method (Fine Tuning)...")

# --- 关键点: 将 SA 的结果作为 Newton 的起点 ---
x_hybrid = x_sa_approx 
hybrid_path = [x_hybrid]

for _ in range(50): # 牛顿法收敛极快，50次足够
    df, ddf = get_derivatives(target_func, x_hybrid)
    
    if abs(ddf) < 1e-9: break # 防止除零
    
    # 牛顿更新: x = x - f'/f''
    step = df / ddf
    x_hybrid = x_hybrid - step
    hybrid_path.append(x_hybrid)
    
    if abs(step) < 1e-10: # 极高精度停止准则
        print("   -> 牛顿法已收敛 (Converged).")
        break

y_hybrid = target_func(x_hybrid)
print(f"✅ 混合策略最终解: x = {x_hybrid:.8f} (精度极高)")


# ==========================================
# 对比实验: 如果直接用牛顿法 (Bad Start)
# ==========================================
print("\n⚠️ 对比: 直接使用牛顿法 (无 SA 辅助)...")
# 故意选一个离局部最优很近，但离全局最优很远的起点 (比如 -10)
x_bad_newton = -10.0 
for _ in range(50):
    df, ddf = get_derivatives(target_func, x_bad_newton)
    if abs(ddf) < 1e-9: break
    x_bad_newton = x_bad_newton - df / ddf

print(f"❌ 失败。陷入局部最优: x = {x_bad_newton:.4f}")


# ==========================================
# 可视化结果
# ==========================================
plt.figure(figsize=(12, 6))
plt.plot(x_vec, y_vec, 'k-', alpha=0.3, label='Target Function')

# 1. 画出 SA 的探索轨迹 (散点)
plt.scatter(sx[::10], target_func(sx[::10]), c='orange', s=10, alpha=0.3, label='Phase 1: SA Exploration')

# 2. 画出 混合策略的结果 (绿星)
plt.plot(x_hybrid, y_hybrid, 'g*', markersize=25, label='Phase 2: Hybrid Result (Global Opt)')

# 3. 画出 失败的牛顿法结果 (红叉)
plt.plot(x_bad_newton, target_func(x_bad_newton), 'rx', markersize=20, markeredgewidth=3, label='Newton Only (Local Trap)')

plt.title(f"Hybrid Optimization Strategy\nSA (Approx: {x_sa_approx:.2f}) + Newton (Final: {x_hybrid:.5f})")
plt.legend()
plt.grid(True)
plt.show()