import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv  # 使用DummyVecEnv避免多进程问题
from stable_baselines3.common.env_util import make_vec_env
from typing import Callable, Any
from typing import Any, Literal
import numpy as np
import pybullet as p
from gymnasium import spaces
from PyFlyt.core.aviary import Aviary
from PyFlyt.core.utils.compile_helpers import check_numpy
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv

import gymnasium
import PyFlyt.gym_envs
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy


# --------------------------
# 核心：综合Wrapper（解决不动+调高度+自定义奖励）
# --------------------------
class QuadXPoleFullWrapper(gymnasium.Wrapper):
    def __init__(
        self,
        env,
        hover_bias=0.2,       # 基础悬停PWM（解决不动：必须>0.5才够升力）
        action_scale=0.2,       # 动作微调范围（控制电机微调幅度，避免过大/过小）
        target_height=2.5,      # 目标悬停高度（调高默认高度，可改3.0/4.0）
        reward_scaling=0.1      # 奖励缩放（避免奖励值过大导致训练不稳定）
    ):
        super().__init__(env)
        self.hover_bias = hover_bias    # 基础悬停推力（确保无人机能起飞）
        self.action_scale = action_scale# 动作微调范围（[-scale, +scale]）
        self.target_height = target_height  # 目标高度
        self.reward_scaling = reward_scaling# 奖励缩放系数

    def reset(self, **kwargs):
        """重置时将无人机初始高度设为目标高度"""
        obs, info = self.env.reset(** kwargs)
        # 修改无人机初始z轴位置（PyFlyt无人机状态的第3个元素是高度）
        if hasattr(self.env.unwrapped, "drone"):
            self.env.unwrapped.drone.state[2] = self.target_height  # z轴=目标高度
        return obs, info

    def step(self, action):
        """1. 处理动作（确保有足够升力）；2. 自定义奖励；3. 返回新状态"""
        # 1. 动作映射：模型输出[-1,1] → 实际PWM[hover_bias-scale, hover_bias+scale]
        # 保证电机有基础悬停推力，解决“不动”问题
        action = action * self.action_scale + self.hover_bias
        # 限制动作在[0,1]（避免PWM超出物理范围导致报错）
        action = np.clip(action, 0.0, 1.0)

        # 2. 执行动作，获取原始环境反馈
        obs, _, term, trunc, info = self.env.step(action)

        # 3. 解析观测值（按PyFlyt QuadX-Pole-Balance-v3观测空间定义）
        pos = obs[:3]          # 无人机位置 (x, y, z)
        orn = obs[3:7]         # 无人机姿态（四元数 x, y, z, w）
        pole_angle = obs[10]   # 杆倾斜角度（核心平衡指标，索引10为主要倾斜角）
        # （可选）如果需要更精准，可查看官方文档：观测空间包含杆的多个角度，取影响最大的一个

        # 4. 自定义奖励计算（多维度鼓励稳定）
        # ① 高度奖励：越接近目标高度，奖励越高（惩罚高度误差）
        height_error = pos[2] - self.target_height
        height_reward = -1.5 * (height_error ** 2)  # 权重1.5，误差越小奖励越高

        # ② 姿态奖励：无人机越水平，奖励越高（惩罚姿态偏移）
        # 四元数x/y/z越小，姿态越接近水平（w为实部，代表水平状态）
        orientation_reward = -0.8 * np.sum(orn[:3] ** 2)  # 权重0.8

        # ③ 杆平衡奖励：杆越竖直，奖励越高（惩罚杆倾斜）
        pole_reward = -2.0 * (pole_angle ** 2)  # 权重2.0，杆平衡是核心任务，权重更高

        # ④ 动作平滑奖励：避免电机大幅调整（惩罚过大动作）
        action_penalty = -0.1 * np.sum(action ** 2)  # 权重0.1，抑制动作波动

        # ⑤ 存活奖励：每步给固定奖励，鼓励持续存活（核心目标是“尽可能久”）
        alive_bonus = 1.2  # 每步+1.2，存活越久总奖励越高

        # 总奖励：加权求和 + 缩放
        total_reward = (
            height_reward + orientation_reward + pole_reward + action_penalty + alive_bonus
        ) * self.reward_scaling

        # 5. 返回处理后的结果
        return obs, total_reward, term, trunc, info


# --------------------------
# 1. 创建并包装环境
# --------------------------
# 原始环境配置（按官方文档，render_mode="human"实时显示）

# --------------------------
env_id = "PyFlyt/QuadX-Pole-Balance-v4"

# 用 make_vec_env 创建多个环境（n_envs 是并行环境数量）
env = make_vec_env(
    env_id,
    n_envs=4,  # 4个环境同时运行（可根据CPU核心数调整）
    wrapper_class=QuadXPoleFullWrapper,  # 我们的自定义包装器
    env_kwargs={
        "render_mode": None,  # 多环境训练时先不渲染，加快速度
        "max_duration_seconds": 30.0,
        "flight_dome_size": 5.0,
        "angle_representation": "quaternion"
    },
    wrapper_kwargs={
        "hover_bias": 0.2,
        "action_scale": 0.3,
        "target_height": 2.5,
        "reward_scaling": 0.1
    }
)
# 查看环境空间（确认配置正确）
print("动作空间（4个电机PWM）：", env.action_space)
print("观测空间（无人机+杆状态）：", env.observation_space)


# --------------------------
# 2. 定义PPO模型（适合连续动作，收敛快）
# --------------------------
model = PPO(
    policy="MlpPolicy",               # 多层感知器（处理连续动作）
    env=env,
    verbose=1,                        # 训练时打印详细信息（loss、reward等）
    tensorboard_log="./quadx_log/",   # 日志保存路径（可在TensorBoard查看训练曲线）
    learning_rate=3e-4,               # 学习率（连续动作任务常用3e-4）
    n_steps=2048,                     # PPO每批收集2048步数据
    batch_size=64,                    # 每批数据分64个batch训练（2048÷64=32，整除）
    n_epochs=10,                      # 每批数据训练10轮
    gamma=0.99,                       # 折扣因子（重视长期奖励）
    gae_lambda=0.95,                  # GAE参数（平衡偏差和方差）
    clip_range=0.2,                   # PPO裁剪范围（经典值0.2）
    ent_coef=0.01,                    # 熵系数（鼓励探索，避免过早收敛到局部最优）
    device="auto"                     # 自动使用GPU/CPU（有GPU会自动调用）
)


# --------------------------
# 3. 训练模型
# --------------------------
print("\n=== 开始训练 ===")
model.learn(
    total_timesteps=300000,  # 总训练步数（30万步，该任务较复杂，需足够步数）
    log_interval=10,         # 每10个批次打印一次训练信息
    progress_bar=True        # 显示训练进度条
)

# 保存训练好的模型（后续可直接加载，不用重新训练）
model.save("quadx_pole_balance_trained_model")
print("\n=== 模型已保存为：quadx_pole_balance_trained_model ===")


# --------------------------
# 4. 评估训练效果
# --------------------------
print("\n=== 开始评估（5局平均奖励） ===")
mean_reward, std_reward = evaluate_policy(
    model=model,
    env=env,
    n_eval_episodes=5,  # 评估5局
    render=True,        # 评估时实时显示
    deterministic=True  # 用确定性策略（避免随机动作，体现真实训练效果）
)
print(f"评估结果：平均奖励 = {mean_reward:.2f} ± {std_reward:.2f}")
# （说明：平均奖励越高、标准差越小，模型越稳定；若平均存活时间接近30秒，说明训练成功）


# --------------------------
# 5. 手动测试（可视化训练成果）
# --------------------------
print("\n=== 开始手动测试（持续1000步） ===")
obs, _ = env.reset()  # 重置环境
for step in range(1000):
    # 模型预测动作（确定性策略）
    action, _ = model.predict(obs, deterministic=True)
    # 执行动作
    obs, reward, term, trunc, info = env.step(action)
    # 若终止（坠毁/杆落地/超时），重置环境继续测试
    if term or trunc:
        print(f"第{step+1}步终止，重置环境...")
        obs, _ = env.reset()

# 关闭环境（释放资源）
env.close()
print("\n=== 测试结束 ===")