# ==============================================================================
# Author: Hongru Liu
# Affiliation: School of Power and Energy, Northwestern Polytechnical University
# Version: 1.0
# Contact: hongruliu@mail.nwpu.edu.cn
# ==============================================================================

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

class Logger(object):
    def __init__(self):
        """
        初始化日志记录器，用于暂存训练过程中的各项指标
        """
        self.episode_rewards = []
        
    def record(self, reward):
        """
        记录每个 Episode 的总奖励
        """
        self.episode_rewards.append(reward)
        
    def plot_learning_curve(self, save_dir="."):
        """
        绘制并保存学习曲线 (Learning Curve)
        """
        # 确保保存目录存在
        os.makedirs(save_dir, exist_ok=True)
        
        # 创建画布，严格设置白色背景
        fig, ax = plt.subplots(figsize=(10, 6), facecolor='white')
        ax.set_facecolor('white')
        
        # 绘制奖励曲线，使用加粗线条以满足论文发表的视觉要求
        ax.plot(self.episode_rewards, linewidth=2.5, color='#1f77b4', label='Episode Reward')
        
        # 计算并绘制 10 个 Episode 的滑动平均线，让趋势更清晰
        if len(self.episode_rewards) >= 10:
            moving_avg = np.convolve(self.episode_rewards, np.ones(10)/10, mode='valid')
            ax.plot(range(9, len(self.episode_rewards)), moving_avg, 
                    linewidth=2.5, color='#ff7f0e', label='10-Episode Moving Average')

        # 设置全英文的坐标轴标签和图例，调整字体大小
        ax.set_xlabel('Episodes', fontsize=14, fontweight='bold')
        ax.set_ylabel('Total Reward', fontsize=14, fontweight='bold')
        ax.set_title('Training Learning Curve (Pendulum-v1)', fontsize=16, fontweight='bold')
        
        # 设置刻度字体大小
        ax.tick_params(axis='both', which='major', labelsize=12)
        
        # 增加网格线并设置图例
        ax.grid(True, linestyle='--', alpha=0.7)
        ax.legend(fontsize=12, loc='lower right')
        
        # 紧凑布局并保存
        plt.tight_layout()
        save_path = os.path.join(save_dir, "learning_curve.png")
        plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none')
        plt.close()
        print(f"[*] Learning curve saved to: {save_path}")