import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
from agent.ppo import PPOAgent
from agent.trpo import TRPOAgent


def train_agent(agent, env_name, num_episodes=500, batch_size=2000):
    """通用训练函数，适用于 PPO 和 TRPO"""
    env = gym.make(env_name)
    episode_rewards = []

    state, _ = env.reset()
    memory = []
    current_ep_reward = 0
    episodes_completed = 0
    step_count = 0

    while episodes_completed < num_episodes:
        action = agent.get_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated

        mask = 0.0 if done else 1.0
        reward_store = reward
        if truncated and not terminated:
            reward_store = reward + agent.gamma * agent.get_value(next_state)

        memory.append([state, action, reward_store, next_state, mask])
        state = next_state
        current_ep_reward += reward
        step_count += 1

        if done:
            episode_rewards.append(current_ep_reward)
            episodes_completed += 1
            state, _ = env.reset()
            current_ep_reward = 0

            if episodes_completed % 10 == 0:
                avg_reward = np.mean(episode_rewards[-10:])
                print(f"  Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")

        if step_count >= batch_size:
            agent.update(memory)
            memory.clear()
            step_count = 0

    env.close()
    return episode_rewards


def smooth(rewards, window=10):
    """滑动平均平滑曲线"""
    smoothed = []
    for i in range(len(rewards)):
        start = max(0, i - window + 1)
        smoothed.append(np.mean(rewards[start:i + 1]))
    return smoothed


def main():
    env_name = 'Pendulum-v1'
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    action_bound = float(env.action_space.high[0])
    env.close()

    num_episodes = 500

    # --- 训练 PPO ---
    print("=" * 50)
    print("开始训练 PPO 智能体...")
    print("=" * 50)
    ppo_agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
    ppo_rewards = train_agent(ppo_agent, env_name, num_episodes)

    # --- 训练 TRPO ---
    print("=" * 50)
    print("开始训练 TRPO 智能体...")
    print("=" * 50)
    trpo_agent = TRPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
    trpo_rewards = train_agent(trpo_agent, env_name, num_episodes)

    # --- 对比画图 ---
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    # 左图：原始奖励曲线
    axes[0].plot(ppo_rewards, alpha=0.3, color='blue', label='PPO (raw)')
    axes[0].plot(trpo_rewards, alpha=0.3, color='red', label='TRPO (raw)')
    axes[0].plot(smooth(ppo_rewards, 20), color='blue', linewidth=2, label='PPO (smooth)')
    axes[0].plot(smooth(trpo_rewards, 20), color='red', linewidth=2, label='TRPO (smooth)')
    axes[0].set_title('PPO vs TRPO on Pendulum-v1')
    axes[0].set_xlabel('Episode')
    axes[0].set_ylabel('Total Reward')
    axes[0].legend()
    axes[0].grid(True)

    # 右图：滑动平均对比（更清晰）
    axes[1].plot(smooth(ppo_rewards, 20), color='blue', linewidth=2, label='PPO')
    axes[1].plot(smooth(trpo_rewards, 20), color='red', linewidth=2, label='TRPO')
    axes[1].set_title('PPO vs TRPO (Smoothed, window=20)')
    axes[1].set_xlabel('Episode')
    axes[1].set_ylabel('Total Reward')
    axes[1].legend()
    axes[1].grid(True)

    plt.tight_layout()
    plt.savefig('ppo_vs_trpo_comparison.png', dpi=150)
    plt.show()
    print("对比图已保存至 ppo_vs_trpo_comparison.png")


if __name__ == '__main__':
    main()