import gymnasium as gym
import torch
from agents.qac import QACAgent
from agents.a2c import A2CAgent
from utils import plot_comparison

def train_agent(env_name, agent_class, device, num_episodes=500):
    """
    通用的训练循环函数
    """
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0] # type: ignore
    action_dim = int(env.action_space.n)       # type: ignore
    
    # 实例化传入的算法代理
    agent = agent_class(state_dim, action_dim, device)
    
    rewards_history = []
    
    for episode in range(num_episodes):
        state, _ = env.reset()
        episode_reward = 0
        
        # QAC 属于 Sarsa 类，需要提前采样第一个动作
        action = agent.select_action(state)
        
        while True:
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            
            # 预采样下一个动作
            next_action = agent.select_action(next_state)
            
            # 统一的接口调用更新
            agent.update(state, action, reward, next_state, next_action, done)
            
            state = next_state
            action = next_action
            episode_reward += reward # type: ignore
            
            if done:
                break
                
        rewards_history.append(episode_reward)
        if (episode + 1) % 100 == 0:
            print(f"[{agent_class.__name__}] 回合 {episode+1}/{num_episodes}, 近100回合均分: {sum(rewards_history[-100:])/100:.2f}")
            
    env.close()
    return rewards_history

if __name__ == "__main__":
    # 检测 GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"正在使用的计算设备: {device}")
    
    ENV_NAME = 'CartPole-v1'
    EPISODES = 600
    
    print("\n--- 开始训练 QAC ---")
    qac_rewards = train_agent(ENV_NAME, QACAgent, device, EPISODES)
    
    print("\n--- 开始训练 A2C ---")
    a2c_rewards = train_agent(ENV_NAME, A2CAgent, device, EPISODES)
    
    # 收集结果并绘图对比
    results = {
        'QAC (High Variance)': qac_rewards,
        'A2C (Low Variance)': a2c_rewards
    }
    
    plot_comparison(results, window=50, save_path='qac_vs_a2c_gpu.png')