import matplotlib
# 关键设置：针对 Linux 服务器无 GUI 环境，强制使用 Agg 后端进行纯文件渲染
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
import numpy as np

def plot_comparison(results_dict, window=50, save_path='comparison_result.png'):
    """
    绘制并保存算法对比曲线
    :param results_dict: 字典格式 {'QAC': [奖励列表], 'A2C': [奖励列表]}
    :param window: 移动平均的窗口大小
    :param save_path: 图片保存路径
    """
    plt.figure(figsize=(10, 6))
    
    for algo_name, rewards in results_dict.items():
        # 绘制原始透明度较低的曲线
        plt.plot(rewards, alpha=0.3, label=f'{algo_name} (Raw)')
        
        # 计算并绘制移动平均曲线，使趋势更平滑
        if len(rewards) >= window:
            moving_avg = np.convolve(rewards, np.ones(window)/window, mode='valid')
            plt.plot(np.arange(window-1, len(rewards)), moving_avg, linewidth=2, label=f'{algo_name} (Avg {window})')
            
    plt.xlabel('Episode')
    plt.ylabel('Total Reward')
    plt.title('Algorithm Comparison: QAC vs A2C')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 将图像保存到服务器硬盘
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"对比图像已成功保存至: {save_path}")