import matplotlib.pyplot as plt
import numpy as np
import argparse

# 添加输入参数
parser = argparse.ArgumentParser(description="Visualize reward curve with smoothing.")
parser.add_argument('--file_path', type=str, default="exps/runs/reward.txt", help='Path to the reward.txt file')
parser.add_argument('--window_size', type=int, default=20, help='Smoothing window size')
args = parser.parse_args()

file_path = args.file_path
window_size = args.window_size
# 文件路径

save_path = file_path.split('.')[0] + f"_vis_{window_size}.png"

# 读取奖励数据
rewards = []
with open(file_path, 'r') as f:
    for line in f:
        rewards.append(float(line.strip()))

# 创建步数数组
steps = np.arange(len(rewards))

# 平滑函数（使用滑动窗口平均）
def smooth_data(data, window_size=10):
    smoothed = []
    for i in range(len(data)):
        start = max(0, i - window_size // 2)
        end = min(len(data), i + window_size // 2 + 1)
        window = data[start:end]
        smoothed.append(sum(window) / len(window))
    return smoothed

# 平滑处理（窗口大小可调整）
smoothed_rewards = smooth_data(rewards, window_size=window_size)

# 创建图表
plt.figure(figsize=(12, 6))
plt.plot(steps, rewards, color='skyblue', alpha=0.4, label='Raw Reward')
plt.plot(steps, smoothed_rewards, color='tab:blue', linewidth=2, label='Smoothed Reward')

# 添加标签和标题
plt.xlabel('Training Step', fontsize=12)
plt.ylabel('Reward', fontsize=12)
plt.title('Reward per Training Step', fontsize=14, pad=20)
plt.legend(loc='best', frameon=True)
plt.grid(True, alpha=0.3)

# 优化布局
plt.tight_layout()
plt.savefig(save_path, dpi=600)