import gymnasium as gym
import torch
import numpy as np
import matplotlib.pyplot as plt  # 新增：用于画图
from agent import TRPOAgent
from utils import RolloutBuffer

# 新增：用于计算滑动平均，让曲线更平滑
def moving_average(data, window_size=5):
    """计算简单滑动平均"""
    if len(data) < window_size:
        return data
    return np.convolve(data, np.ones(window_size)/window_size, mode='valid')

def plot_learning_curve(rewards, smoothed_rewards):
    """绘制学习曲线"""
    plt.figure(figsize=(10, 6))
    
    # 绘制原始奖励（浅色）
    plt.plot(rewards, color='blue', alpha=0.3, label='Raw Average Reward')
    
    # 绘制平滑后的奖励（深色，粗线）
    # 注意：滑动平均后数据点会变少，需要调整 X 轴起始位置
    if len(smoothed_rewards) > 0:
        x_ticks = np.arange(len(rewards) - len(smoothed_rewards), len(rewards))
        plt.plot(x_ticks, smoothed_rewards, color='red', linewidth=2, label='Smoothed Reward (MA-5)')
    
    plt.title('TRPO Training Performance on Pendulum-v1')
    plt.xlabel('Iteration')
    plt.ylabel('Average Reward')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.legend()
    
    # 保存图片
    plt.savefig('trpo_training_curve.png')
    print("\n训练曲线图已保存为 'trpo_training_curve.png'")
    
    # 如果有 GUI 界面则显示
    try:
        plt.show()
    except Exception:
        print("无法显示图形界面（可能是无头服务器），已跳过 plt.show()。")

def main():
    # 1. 初始化环境
    #env_name = 'Pendulum-v1'
    # env = gym.make('Pendulum-v1')
    # 兼容 Gymnasium
    try:
        env = gym.make('Pendulum-v1', render_mode=None)
    except Exception:
        env = gym.make('Pendulum-v1')
        
    state_dim = env.observation_space.shape[0]          # type: ignore
    action_dim = env.action_space.shape[0]              # type: ignore
    max_action = float(env.action_space.high[0])        # type: ignore
    min_action = float(env.action_space.low[0])         # type: ignore

    print(f"环境加载成功! 状态维度: {state_dim}, 动作维度: {action_dim}")
    print(f"动作范围: [{min_action}, {max_action}]")

    # 2. 初始化智能体和经验池
    agent = TRPOAgent(state_dim, action_dim, max_kl=0.01)
    buffer = RolloutBuffer()

    # 3. 设置训练超参数
    max_iterations = 300       # 训练迭代总轮数
    batch_size = 2000          # 每次更新收集的步数
    
    # 新增：用于记录绘图数据
    history_rewards = []

    # 4. 主训练循环
    for iteration in range(max_iterations):
        state = env.reset()
        if isinstance(state, tuple): state = state[0]
            
        episode_rewards = []
        ep_reward = 0
        steps = 0
        done = False

        while steps < batch_size:
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            
            with torch.no_grad():
                action, log_prob = agent.actor.get_action(state_tensor)
                value = agent.critic(state_tensor)
            
            action_np = action.squeeze(0).numpy()
            clipped_action = np.clip(action_np, min_action, max_action)
            
            # env.step 兼容性
            step_result = env.step(clipped_action)
            if len(step_result) == 5:
                next_state, reward, terminated, truncated, _ = step_result
                done = terminated or truncated
            else:
                next_state, reward, done, _ = step_result
                
            ep_reward += reward
            buffer.add(state=state, action=action_np, reward=reward, 
                       next_state=next_state, done=done, 
                       log_prob=log_prob.item(), value=value.item())
            
            state = next_state
            steps += 1
            
            if done:
                episode_rewards.append(ep_reward)
                state = env.reset()
                if isinstance(state, tuple): state = state[0]
                ep_reward = 0

        print(f"正在更新参数 (Iteration {iteration + 1}/{max_iterations})...")
        agent.update(buffer, state, done)
        
        # 5. 记录和打印日志
        if episode_rewards:
            avg_reward = np.mean(episode_rewards)
            history_rewards.append(avg_reward) # 新增：记录数据
            print(f"Iteration: {iteration + 1} | Average Reward: {avg_reward:.2f} | Max Reward: {np.max(episode_rewards):.2f}")
            print("-" * 50)
        else:
            # 如果 batch_size 刚好结束时没有完成任何 episode，
            # 为了画图不中断，我们沿用上一次的奖励（或者简单处理）
            if history_rewards:
                history_rewards.append(history_rewards[-1])
            else:
                history_rewards.append(-2000) # 初始默认低分

    print("训练结束！")
    env.close()

    # 6. 新增：绘图
    print("正在生成训练曲线图...")
    smoothed = moving_average(history_rewards, window_size=5)
    plot_learning_curve(history_rewards, smoothed)

if __name__ == '__main__':
    # 确保安装了 matplotlib: pip install matplotlib
    main()