import gymnasium as gym
import torch
import os
import pickle
from agents.qac import QACAgent
from agents.a2c import A2CAgent
from agents.off_pac import OffPACAgent
from utils import plot_comparison

RESULTS_DIR = 'training_results'

def get_result_path(agent_class, env_name, num_episodes):
    os.makedirs(RESULTS_DIR, exist_ok=True)
    return os.path.join(RESULTS_DIR, f'{agent_class.__name__}_{env_name}_{num_episodes}.pkl')

def load_or_train(env_name, agent_class, device, num_episodes):
    """检查是否有保存的结果，有则加载，没有则训练并保存。"""
    result_path = get_result_path(agent_class, env_name, num_episodes)
    if os.path.exists(result_path):
        with open(result_path, 'rb') as f:
            rewards = pickle.load(f)
        print(f"[{agent_class.__name__}] 已找到训练结果文件，直接加载: {result_path}")
        return rewards

    # 没有结果文件，开始训练
    print(f"[{agent_class.__name__}] 未找到结果文件，开始训练...")
    rewards = _train_agent(env_name, agent_class, device, num_episodes)

    with open(result_path, 'wb') as f:
        pickle.dump(rewards, f)
    print(f"[{agent_class.__name__}] 训练完成，结果已保存至: {result_path}")
    return rewards

def _train_agent(env_name, agent_class, device, num_episodes):
    """实际的训练循环逻辑"""
    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

        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 = load_or_train(ENV_NAME, QACAgent, device, EPISODES)

    print("\n--- 开始训练 A2C ---")
    a2c_rewards = load_or_train(ENV_NAME, A2CAgent, device, EPISODES)

    print("\n--- 开始训练 Off-Policy A2C ---")
    off_pac_rewards = load_or_train(ENV_NAME, OffPACAgent, device, EPISODES)
    
    # 将其加入字典一起画图对比
    results = {
        'QAC': qac_rewards,
        'A2C': a2c_rewards,
        'Off-Policy A2C': off_pac_rewards
    }

    
    plot_comparison(results, window=50, save_path='qac_vs_a2c_gpu.png')