import gymnasium as gym
import torch
import os
import pickle
from agents.dpac import DPACAgent
from agents.ddpg import DDPGAgent
from utils import plot_comparison


RESULTS_DIR = 'training_results'


# ---------- 路径工具 ----------
def _path(agent_class, env_name, num_episodes, suffix):
    os.makedirs(RESULTS_DIR, exist_ok=True)
    return os.path.join(RESULTS_DIR, f'{agent_class.__name__}_{env_name}_{num_episodes}{suffix}')


def result_path(agent_class, env_name, num_episodes):
    return _path(agent_class, env_name, num_episodes, '.pkl')


def model_path(agent_class, env_name, num_episodes):
    return _path(agent_class, env_name, num_episodes, '_model.pt')


# ---------- 训练 & 保存 ----------
def train(env_name, agent_class, device, num_episodes):
    rp = result_path(agent_class, env_name, num_episodes)
    mp = model_path(agent_class, env_name, num_episodes)

    if os.path.exists(rp) and os.path.exists(mp):
        with open(rp, 'rb') as f:
            rewards = pickle.load(f)
        print(f"[{agent_class.__name__}] 已加载训练结果: {rp}")
        return rewards

    print(f"[{agent_class.__name__}] 开始训练 ...")
    if agent_class.__name__ == 'DDPGAgent':
        rewards, agent = _train_ddpg(env_name, agent_class, device, num_episodes)
    else:
        rewards, agent = _train_dpac(env_name, agent_class, device, num_episodes)

    with open(rp, 'wb') as f:
        pickle.dump(rewards, f)
    agent.save(mp)
    print(f"[{agent_class.__name__}] 训练完成 | 曲线: {rp} | 模型: {mp}")
    return rewards


def _train_dpac(env_name, agent_class, device, num_episodes):
    env = gym.make(env_name)
    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

    agent = agent_class(state_dim, action_dim, max_action, device)
    rewards = []

    for ep in range(num_episodes):
        state, _ = env.reset()
        episode_reward = 0
        while True:
            action = agent.select_action(state)
            ns, reward, terminated, truncated, _ = env.step(action)
            agent.update(state, action, reward, ns, None, terminated or truncated)
            state = ns
            episode_reward += reward                    # type: ignore
            if terminated or truncated:
                break
        rewards.append(episode_reward)
        if (ep + 1) % 50 == 0:
            avg = sum(rewards[-50:]) / 50
            print(f"  回合 {ep+1}/{num_episodes}, 近50回合均分: {avg:.2f}")

    env.close()
    return rewards, agent


def _train_ddpg(env_name, agent_class, device, num_episodes, batch_size=64, start_steps=2000):
    env = gym.make(env_name)
    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

    agent = agent_class(state_dim, action_dim, max_action, device)
    rewards = []
    total_steps = 0

    for ep in range(num_episodes):
        state, _ = env.reset()
        episode_reward = 0
        while True:
            action = agent.select_action(state)
            ns, reward, terminated, truncated, _ = env.step(action)
            agent.replay_buffer.add(state, action, reward, ns, terminated or truncated)
            total_steps += 1
            state = ns
            episode_reward += reward                    # type: ignore
            if total_steps > start_steps:
                agent.update(batch_size)
            if terminated or truncated:
                break
        rewards.append(episode_reward)
        if (ep + 1) % 50 == 0:
            avg = sum(rewards[-50:]) / 50
            print(f"  回合 {ep+1}/{num_episodes}, 近50回合均分: {avg:.2f}")

    env.close()
    return rewards, agent


# ---------- 视频演示 ----------
def demo(env_name, agent_class, device, num_demo=5, num_episodes=600):
    mp = model_path(agent_class, env_name, num_episodes)
    if not os.path.exists(mp):
        print(f"[{agent_class.__name__}] 未找到模型权重 {mp}，跳过演示。")
        return

    video_dir = os.path.join(RESULTS_DIR, 'videos')
    os.makedirs(video_dir, exist_ok=True)

    base_env = gym.make(env_name, render_mode='rgb_array')
    env = gym.wrappers.RecordVideo(base_env, video_dir, name_prefix=f'{agent_class.__name__}_demo')
    state_dim = env.env.observation_space.shape[0]              # type: ignore
    action_dim = env.env.action_space.shape[0]                  # type: ignore
    max_action = float(env.env.action_space.high[0])            # type: ignore

    agent = agent_class(state_dim, action_dim, max_action, device)
    agent.load(mp)
    print(f"[{agent_class.__name__}] 加载模型: {mp}，录制 {num_demo} 局演示 ...")

    for ep in range(num_demo):
        state, _ = env.reset()
        total_reward = 0.0
        while True:
            action = agent.select_action(state)
            state, reward, terminated, truncated, _ = env.step(action)
            total_reward += float(reward)
            if terminated or truncated:
                break
        print(f"  回合 {ep+1}/{num_demo}, 总奖励: {total_reward:.2f}")

    env.close()
    base_env.close()
    print(f"[{agent_class.__name__}] 视频已保存至: {video_dir}")


# ---------- 主程序 ----------
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"计算设备: {device}\n")

    ENV_NAME = 'Pendulum-v1'
    EPISODES = 600

    dpac_rewards = train(ENV_NAME, DPACAgent, device, EPISODES)
    ddpg_rewards = train(ENV_NAME, DDPGAgent, device, EPISODES)

    results = {
        'DPAC': dpac_rewards,
        'DDPG': ddpg_rewards,
    }
    plot_comparison(results, window=50, save_path='compare_dpac_vs_ddpg.png')

    demo(ENV_NAME, DPACAgent, device)
    demo(ENV_NAME, DDPGAgent, device)