import numpy as np
import kaiwu.optimizer as optimizer
import kaiwu.solver as solver
import matplotlib.pyplot as plt
# ====================== 步骤1：构建Qubo模型的Q矩阵（自动驾驶场景） ======================
# 变量数：200（100个路径节点+50个传感器+50个算力任务）
n_vars = 200
# 初始化Q矩阵（对称矩阵，Qubo模型核心）
Q = np.zeros((n_vars, n_vars))
# 1. 填充对角元（对应单个变量的基础贡献：时间+能耗+风险）
# 模拟路径节点时间贡献（a_i）、传感器能耗（c_j）、算力任务风险（e_l）
np.random.seed(42) # 固定随机种子保证可复现
a_i = np.random.uniform(1, 10, size=100) # 路径节点基础时间
c_j = np.random.uniform(0.5, 5, size=50) # 传感器基础能耗
e_l = np.random.uniform(2, 8, size=50) # 算力任务基础风险
# 权重（时间w1=0.4，能耗w2=0.3，风险w3=0.3）
w1, w2, w3 = 0.4, 0.3, 0.3
# 填充对角元：前100个为路径节点，100-150为传感器，150-200为算力任务
Q[:100, :100] = np.diag(w1 * a_i)
Q[100:150, 100:150] = np.diag(w2 * c_j)
Q[150:200, 150:200] = np.diag(w3 * e_l)
# 2. 填充非对角元（变量间协同贡献：节点转移时间、传感器协同能耗、任务风险叠加）
# 模拟节点转移时间b_ij（仅填充前100×100区域，随机生成对称矩阵）
b_ij = np.random.uniform(0.1, 2, size=(100, 100))
b_ij = (b_ij + b_ij.T) / 2 # 保证对称
Q[:100, :100] += w1 * b_ij / 2 # 非对角元需除以2避免重复计算
# 3. 添加约束惩罚项（算力上限约束）
lambda1 = 30 # 惩罚系数（实验验证的平衡值）
C_max = 50 # 算力上限
g_l = np.random.uniform(0.5, 2, size=50) # 算力任务占用量
# 惩罚项展开：λ*(Σg_l x_l - C_max)² = λ*(Σg_l²x_l² + 2Σg_l g_p x_l x_p - 2C_maxΣg_l x_l + C_max²)
# 因x_l是0-1变量，x_l²=x_l，故惩罚项融入Q矩阵：
penalty_diag = lambda1 * (g_l **2 - 2 * C_max * g_l) # 对角元惩罚
penalty_offdiag = lambda1 * 2 * np.outer(g_l, g_l) # 非对角元惩罚
Q[150:200, 150:200] += np.diag(penalty_diag)
Q[150:200, 150:200] += penalty_offdiag / 2 # 非对角元除以2
# ====================== 步骤2：初始化Qubo优化器 ======================
# 加载Q矩阵到优化器
qubo_opt = optimizer.QuboOptimizer()
qubo_opt.set_qubo_matrix(Q) # 传入200×200的Q矩阵
# ====================== 步骤3：配置模拟退火求解器 ======================
# 初始化求解器
sa_solver = solver.SimulatedAnnealingSolver()
# 设置核心参数（匹配竞赛要求）
sa_solver.set_params(
temperature_decay=0.95, # 退火温度衰减系数
max_iter=10000, # 迭代次数
init_solution_method="random" # 初始解生成方法：随机
)
# ====================== 步骤4：执行求解 ======================
# 传入优化器求解
result = sa_solver.solve(qubo_opt)
# ====================== 步骤5：结果解析与可视化 ======================
# 提取核心结果
optimal_solution = result["optimal_solution"] # 最优解（200维0-1数组）
optimal_value = result["optimal_value"] # 最优目标函数值
iteration_curve = result["iteration_curve"] # 迭代过程目标函数值曲线
# 输出关键信息
print("===== 求解结果 =====")
print(f"最优目标函数值：{optimal_value:.4f}")
print("最优解（前10个变量）：", optimal_solution[:10])
# 解析最优解对应的实际决策（示例）
selected_nodes = np.where(optimal_solution[:100]==1)[0] # 选中的路径节点
selected_sensors = np.where(optimal_solution[100:150]==1)[0] + 100 # 开启的传感器
print(f"选中的路径节点数量：{len(selected_nodes)}，节点索引：{selected_nodes[:5]}")
print(f"开启的传感器数量：{len(selected_sensors)}，传感器索引：{selected_sensors[:5]}")
# 绘制迭代曲线
plt.figure(figsize=(10, 6))
plt.plot(iteration_curve, label="目标函数值")
plt.xlabel("迭代次数")
plt.ylabel("目标函数值")
plt.title("模拟退火求解迭代过程")
plt.legend()
plt.grid(True)
plt.savefig("iteration_curve.png")
plt.show()
