import argparse
import os
import random
from pathlib import Path

os.environ.setdefault("MKL_THREADING_LAYER", "GNU")

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm.auto import tqdm

from agent.ppo import PPOAgent
from agent.trpo import TRPOAgent


def parse_args():
    cpu_count = os.cpu_count() or 1
    default_envs = max(1, min(4, cpu_count))
    default_threads = max(1, min(8, cpu_count))

    parser = argparse.ArgumentParser(description="使用 GPU / 多核并行训练 PPO 与 TRPO，并输出对比图。")
    parser.add_argument("--env-name", type=str, default="Pendulum-v1", help="Gymnasium 环境名称。")
    parser.add_argument("--num-episodes", type=int, default=500, help="每个算法训练的 episode 数。")
    parser.add_argument("--batch-size", type=int, default=2000, help="每次策略更新前收集的环境步数。")
    parser.add_argument("--num-envs", type=int, default=default_envs, help="并行采样环境数。大于 1 时可显著提高吞吐。")
    parser.add_argument("--vector-mode", choices=("sync", "async"), default="async", help="向量环境模式。async 会启动多进程，更适合多核 CPU。")
    parser.add_argument("--cpu-threads", type=int, default=default_threads, help="PyTorch CPU 线程数。GPU 训练时主要影响 CPU 侧数据准备。")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto", help="训练设备。auto 会优先选择 CUDA。")
    parser.add_argument("--seed", type=int, default=42, help="全局随机种子。")
    parser.add_argument("--hidden-dim", type=int, default=128, help="Actor/Critic 隐层宽度。")
    parser.add_argument("--output-dir", type=str, default="outputs", help="图像与训练数据的输出目录。")
    parser.add_argument("--no-show", action="store_true", help="仅保存图像，不弹出 matplotlib 窗口。")
    return parser.parse_args()


def resolve_device(device_name):
    if device_name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if device_name == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("请求使用 CUDA，但当前环境不可用。请检查 GPU 驱动和 PyTorch CUDA 版本。")

    return torch.device(device_name)


def configure_runtime(cpu_threads, seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    if cpu_threads is not None:
        torch.set_num_threads(max(1, cpu_threads))
        if hasattr(torch, "set_num_interop_threads"):
            torch.set_num_interop_threads(max(1, min(cpu_threads, 4)))

    if hasattr(torch, "set_float32_matmul_precision"):
        torch.set_float32_matmul_precision("high")


def make_env(env_name, seed, worker_idx):
    def _thunk():
        env = gym.make(env_name)
        env.reset(seed=seed + worker_idx)
        env.action_space.seed(seed + worker_idx)
        return env

    return _thunk


def create_vector_env(env_name, num_envs, seed, vector_mode):
    env_fns = [make_env(env_name, seed, idx) for idx in range(num_envs)]
    if num_envs == 1 or vector_mode == "sync":
        return gym.vector.SyncVectorEnv(env_fns)
    return gym.vector.AsyncVectorEnv(env_fns, context="spawn")


def format_postfix(episode_rewards, total_env_steps, updates, last_update_stats):
    postfix = {
        "avg20": f"{np.mean(episode_rewards[-20:]):.1f}" if episode_rewards else "n/a",
        "steps": total_env_steps,
        "updates": updates,
    }
    for key, value in last_update_stats.items():
        if isinstance(value, (int, float, np.floating)):
            postfix[key] = f"{float(value):.4f}"
    return postfix


def train_agent(agent, env_name, num_episodes=500, batch_size=2000, num_envs=1, seed=42, vector_mode="async", desc=None):
    """通用训练函数，适用于 PPO 和 TRPO，支持 GPU 和向量化并行采样。"""
    env = create_vector_env(env_name, num_envs, seed, vector_mode)
    episode_rewards = []
    memory = {
        "states": [],
        "actions": [],
        "rewards": [],
        "next_states": [],
        "masks": [],
    }
    total_env_steps = 0
    updates = 0
    last_update_stats = {}

    states, _ = env.reset(seed=seed)
    running_rewards = np.zeros(num_envs, dtype=np.float64)
    step_count = 0

    progress = tqdm(total=num_episodes, desc=desc or agent.__class__.__name__, dynamic_ncols=True)

    try:
        while len(episode_rewards) < num_episodes:
            actions = agent.get_action(states)
            next_states, rewards, terminated, truncated, _ = env.step(actions)
            done = np.logical_or(terminated, truncated)

            masks = (~done).astype(np.float32)
            reward_store = rewards.astype(np.float32).copy()
            truncated_only = np.logical_and(truncated, ~terminated)
            if np.any(truncated_only):
                bootstrap_values = np.asarray(agent.get_value(next_states[truncated_only]), dtype=np.float32)
                reward_store[truncated_only] += agent.gamma * bootstrap_values

            memory["states"].append(np.asarray(states, dtype=np.float32).copy())
            memory["actions"].append(np.asarray(actions, dtype=np.float32).copy())
            memory["rewards"].append(np.asarray(reward_store, dtype=np.float32).copy())
            memory["next_states"].append(np.asarray(next_states, dtype=np.float32).copy())
            memory["masks"].append(np.asarray(masks, dtype=np.float32).copy())

            running_rewards += rewards
            step_count += num_envs
            total_env_steps += num_envs

            if np.any(done):
                done_indices = np.flatnonzero(done)
                new_episode_rewards = running_rewards[done_indices].tolist()
                remaining = num_episodes - len(episode_rewards)
                accepted = new_episode_rewards[:remaining]
                if accepted:
                    episode_rewards.extend(accepted)
                    progress.update(len(accepted))

                running_rewards[done_indices] = 0.0
                progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))

                if len(episode_rewards) >= num_episodes:
                    break

                reset_mask = np.zeros(num_envs, dtype=bool)
                reset_mask[done_indices] = True
                reset_states, _ = env.reset(options={"reset_mask": reset_mask})
                next_states[done_indices] = reset_states[done_indices]

            states = next_states

            if step_count >= batch_size:
                last_update_stats = agent.update(memory)
                for key in memory:
                    memory[key].clear()
                step_count = 0
                updates += 1
                progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))

        if memory["states"]:
            last_update_stats = agent.update(memory)
            updates += 1
            progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
    finally:
        progress.close()
        env.close()

    summary = {
        "episodes": len(episode_rewards),
        "total_env_steps": total_env_steps,
        "updates": updates,
        "final_avg20": float(np.mean(episode_rewards[-20:])) if episode_rewards else float("nan"),
    }
    summary.update(last_update_stats)
    return episode_rewards, summary


def smooth(rewards, window=10):
    """滑动平均平滑曲线。"""
    rewards = np.asarray(rewards, dtype=np.float32)
    if rewards.size == 0:
        return rewards

    smoothed = []
    for i in range(rewards.size):
        start = max(0, i - window + 1)
        smoothed.append(float(np.mean(rewards[start:i + 1])))
    return np.asarray(smoothed, dtype=np.float32)


def save_artifacts(ppo_rewards, trpo_rewards, output_dir, env_name):
    output_dir.mkdir(parents=True, exist_ok=True)

    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(f"PPO vs TRPO on {env_name}")
    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()

    figure_path = output_dir / "ppo_vs_trpo_comparison.png"
    metrics_path = output_dir / "training_metrics.npz"
    fig.savefig(figure_path, dpi=150)
    np.savez(metrics_path, ppo_rewards=np.asarray(ppo_rewards), trpo_rewards=np.asarray(trpo_rewards))

    return fig, figure_path, metrics_path


def main():
    args = parse_args()
    configure_runtime(args.cpu_threads, args.seed)
    device = resolve_device(args.device)
    output_dir = Path(args.output_dir)

    env = gym.make(args.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()

    tqdm.write(
        f"配置: env={args.env_name}, device={device}, num_envs={args.num_envs}, "
        f"vector_mode={args.vector_mode}, cpu_threads={args.cpu_threads}, batch_size={args.batch_size}"
    )

    ppo_agent = PPOAgent(
        state_dim=state_dim,
        action_dim=action_dim,
        action_bound=action_bound,
        hidden_dim=args.hidden_dim,
        device=device,
    )
    ppo_rewards, ppo_summary = train_agent(
        ppo_agent,
        args.env_name,
        num_episodes=args.num_episodes,
        batch_size=args.batch_size,
        num_envs=args.num_envs,
        seed=args.seed,
        vector_mode=args.vector_mode,
        desc="PPO",
    )
    tqdm.write(f"PPO 完成: {ppo_summary}")

    trpo_agent = TRPOAgent(
        state_dim=state_dim,
        action_dim=action_dim,
        action_bound=action_bound,
        hidden_dim=args.hidden_dim,
        device=device,
    )
    trpo_rewards, trpo_summary = train_agent(
        trpo_agent,
        args.env_name,
        num_episodes=args.num_episodes,
        batch_size=args.batch_size,
        num_envs=args.num_envs,
        seed=args.seed + 10_000,
        vector_mode=args.vector_mode,
        desc="TRPO",
    )
    tqdm.write(f"TRPO 完成: {trpo_summary}")

    fig, figure_path, metrics_path = save_artifacts(ppo_rewards, trpo_rewards, output_dir, args.env_name)

    tqdm.write(f"对比图已保存至 {figure_path}")
    tqdm.write(f"训练曲线原始数据已保存至 {metrics_path}")

    if args.no_show:
        plt.close(fig)
    else:
        plt.show()


if __name__ == "__main__":
    main()