"""
训练入口 — 支持 SAC / PPO / 对比(compare) 三种模式

用法:
    python train.py --algo sac      # 仅训练 SAC（行为与原始版本一致）
    python train.py --algo ppo      # 仅训练 PPO
    python train.py --algo compare  # 依次训练 SAC 和 PPO，结束后输出对比曲线

可选参数:
    --env   Gymnasium 环境 ID（默认 Pendulum-v1）
    --seed  随机种子（默认 0）
"""
import argparse
import os
import yaml
import numpy as np
import torch
import gymnasium as gym
from typing import cast
from gymnasium.spaces import Box

from algorithms.sac import SAC
from algorithms.ppo import PPO
from utils.replay_buffer import ReplayBuffer
from utils.logger import Logger


# ===========================================================================
#  配置加载
# ===========================================================================

def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


ROOT = os.path.dirname(os.path.abspath(__file__))
SAC_CFG_PATH = os.path.join(ROOT, "configs", "sac_pendulum_config.yaml")
PPO_CFG_PATH = os.path.join(ROOT, "configs", "ppo_pendulum_config.yaml")


# ===========================================================================
#  工具函数
# ===========================================================================

def set_seed(seed: int, env: gym.Env):
    np.random.seed(seed)
    torch.manual_seed(seed)
    env.reset(seed=seed)


def make_env(env_id: str):
    env = gym.make(env_id)
    obs_space = cast(Box, env.observation_space)
    act_space = cast(Box, env.action_space)
    if obs_space.shape is None or act_space.shape is None:
        raise ValueError("环境的 observation/action space 不支持 shape 维度读取。")
    state_dim  = obs_space.shape[0]
    action_dim = act_space.shape[0]
    max_action = float(act_space.high[0])
    return env, state_dim, action_dim, max_action


# ===========================================================================
#  SAC 训练循环 (off-policy)
# ===========================================================================

def train_sac(env_id: str, seed: int) -> list:
    """
    训练 SAC 并返回每个 episode 的总奖励列表。
    """
    config = load_config(SAC_CFG_PATH)
    env, state_dim, action_dim, max_action = make_env(env_id)
    set_seed(seed, env)

    print(f"\n{'='*60}")
    print(f"  Training SAC on {env_id}")
    print(f"  state_dim={state_dim}, action_dim={action_dim}, max_action={max_action}")
    print(f"{'='*60}\n")

    agent         = SAC(state_dim, action_dim, max_action, config)
    replay_buffer = ReplayBuffer(state_dim, action_dim, max_size=config['buffer_size'])
    logger        = Logger()

    total_steps = 0

    for episode in range(config['max_episodes']):
        state, _ = env.reset()
        episode_reward: float = 0.0

        for _ in range(config['max_steps']):
            # 动作选择：预热期使用纯随机，之后用策略
            if total_steps < config['start_steps']:
                action = env.action_space.sample()
            else:
                action = agent.select_action(state, evaluate=False)

            next_state, reward, terminated, truncated, _ = env.step(action)
            done = float(terminated)

            replay_buffer.add(state, action, reward, next_state, done)
            state = next_state
            episode_reward += float(reward)
            total_steps += 1

            if replay_buffer.size > config['batch_size']:
                agent.update(replay_buffer, config['batch_size'])

            if terminated or truncated:
                break

        logger.record(episode_reward)
        print(f"[SAC] Episode: {episode+1:03d} | Steps: {total_steps:06d} | Reward: {episode_reward:.2f}")

        if (episode + 1) % 50 == 0:
            path = os.path.join(ROOT, f"sac_actor_pendulum_ep{episode+1}.pth")
            torch.save(agent.actor.state_dict(), path)
            print(f"  [*] 模型已保存: {path}")

    env.close()
    logger.plot_learning_curve(save_dir=ROOT)
    print("\n[SAC] 训练完成！\n")
    return logger.episode_rewards


# ===========================================================================
#  PPO 训练循环 (on-policy)
# ===========================================================================

def train_ppo(env_id: str, seed: int) -> list:
    """
    训练 PPO 并返回每个 episode 的总奖励列表。

    PPO 是 on-policy 的：先收集固定 T 步数据（steps_per_update），
    然后用这批数据做 K 轮 epoch 更新，再继续收集。
    """
    config = load_config(PPO_CFG_PATH)
    env, state_dim, action_dim, max_action = make_env(env_id)
    set_seed(seed, env)

    print(f"\n{'='*60}")
    print(f"  Training PPO on {env_id}")
    print(f"  state_dim={state_dim}, action_dim={action_dim}, max_action={max_action}")
    print(f"{'='*60}\n")

    agent  = PPO(state_dim, action_dim, max_action, config)
    logger = Logger()

    steps_per_update = config.get('steps_per_update', 2048)
    max_episodes     = config['max_episodes']
    max_steps        = config['max_steps']

    total_steps    = 0
    buffer_steps   = 0
    episode        = 0
    episode_reward: float = 0.0
    state, _ = env.reset()

    while episode < max_episodes:
        action, action_raw, log_prob, value = agent.select_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = float(terminated)

        agent.rollout.add(state, action_raw, log_prob, reward, done, value)
        state = next_state
        episode_reward += float(reward)
        total_steps   += 1
        buffer_steps  += 1

        # ---- episode 结束 ----
        if terminated or truncated:
            logger.record(episode_reward)
            print(f"[PPO] Episode: {episode+1:03d} | Steps: {total_steps:06d} | Reward: {episode_reward:.2f}")

            # 阶段性保存（避免重复保存）
            if (episode + 1) % 50 == 0:
                path = os.path.join(ROOT, f"ppo_actor_pendulum_ep{episode+1}.pth")
                torch.save(agent.actor.state_dict(), path)
                print(f"  [*] 模型已保存: {path}")

            episode       += 1
            episode_reward = 0.0
            state, _       = env.reset()

            if episode >= max_episodes:
                break

        # ---- 收集够 T 步 → 触发 PPO 更新 ----
        if buffer_steps >= steps_per_update:
            last_value = 0.0 if done else agent.get_value(state)
            agent.rollout.compute_returns_and_advantages(last_value)
            info = agent.update()
            buffer_steps = 0
            print(f"  [PPO update] actor_loss={info['actor_loss']:.4f} | "
                  f"value_loss={info['value_loss']:.4f} | entropy={info['entropy']:.4f}")

    env.close()

    # 若缓冲区中还有剩余数据，做最后一次更新
    if buffer_steps > 0:
        agent.rollout.compute_returns_and_advantages(0.0)
        agent.update()

    logger.plot_learning_curve(save_dir=ROOT)
    print("\n[PPO] 训练完成！\n")
    return logger.episode_rewards


# ===========================================================================
#  对比模式：依次训练两者，输出对比曲线
# ===========================================================================

def train_compare(env_id: str, seed: int):
    print("\n" + "="*60)
    print("  Compare Mode: SAC vs PPO")
    print("="*60)

    rewards_sac = train_sac(env_id, seed)
    rewards_ppo = train_ppo(env_id, seed)

    Logger.plot_comparison(
        rewards_dict={'SAC': rewards_sac, 'PPO': rewards_ppo},
        window=10,
        save_dir=ROOT,
    )
    print("\n[Compare] 对比训练完成！已生成 3 张图：")
    print("  - comparison_curve.png    (同 episode 范围对比)")
    print("  - sac_learning_curve.png  (SAC 独立完整曲线)")
    print("  - ppo_learning_curve.png  (PPO 独立完整曲线)")


# ===========================================================================
#  命令行入口
# ===========================================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description="SAC / PPO 强化学习训练脚本 (Pendulum-v1)"
    )
    parser.add_argument(
        "--algo", type=str, default="sac",
        choices=["sac", "ppo", "compare"],
        help="选择训练的算法: sac | ppo | compare (默认: sac)"
    )
    parser.add_argument(
        "--env", type=str, default="Pendulum-v1",
        help="Gymnasium 环境 ID (默认: Pendulum-v1)"
    )
    parser.add_argument(
        "--seed", type=int, default=0,
        help="随机种子 (默认: 0)"
    )
    return parser.parse_args()


def main():
    args = parse_args()

    if args.algo == "sac":
        train_sac(args.env, args.seed)
    elif args.algo == "ppo":
        train_ppo(args.env, args.seed)
    elif args.algo == "compare":
        train_compare(args.env, args.seed)


if __name__ == "__main__":
    main()