import torch
from isaaclab.assets import ArticulationCfg
from isaaclab.envs import ManagerBasedRLEnvCfg, ManagerBasedRLEnv
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.envs.mdp import JointPositionActionCfg
from isaaclab.managers import SceneEntityCfg
from isaaclab.utils import configclass
from rl_game.get_up.env.t1_env import T1SceneCfg

import isaaclab.envs.mdp as mdp


# --- 1. 自定义 MDP 逻辑函数 (放在配置类之前) ---

def is_standing_still(
        env: ManagerBasedRLEnv,
        minimum_height: float,
        max_angle_error: float,
        standing_time: float
) -> torch.Tensor:
    """
    判定机器人是否稳定站立了一段时间。
    逻辑：高度达标且姿态垂直，持续时间超过 standing_time 则返回 True。
    """
    # 获取当前状态：高度 (Z轴) 和 投影重力 (前两个分量越小越垂直)

    # 1. 首先获取 H2 (头部) 的索引 (建议在环境初始化时获取一次，或者如下所示获取)
    # find_bodies 返回 (indices, names)
    head_idx, _ = env.scene["robot"].find_bodies("H2")

    # 2. 修改后的位置获取逻辑
    # data.body_link_pos_w 的维度是 (num_envs, num_bodies, 3)
    # 我们取所有环境 (:)，对应的头部索引 (head_idx[0])，以及 Z 轴坐标 (2)
    current_head_height = env.scene["robot"].data.body_link_pos_w[:, head_idx[0], 2]

    # 3. 姿态判定保持不变（通常依然以躯干 Trunk 的垂直度为准，因为头部可能会摆动）
    gravity_error = torch.norm(env.scene["robot"].data.projected_gravity_b[:, :2], dim=-1)

    # 4. 更新判断逻辑
    is_stable_now = (current_head_height > minimum_height) & (gravity_error < max_angle_error)

    # 在 env.extras 中维护一个计时器
    if "stable_timer" not in env.extras:
        env.extras["stable_timer"] = torch.zeros(env.num_envs, device=env.device)

    # 核心计时逻辑：达标累加 dt，不达标清零
    # dt = physics_dt * decimation (即控制频率的步长)
    dt = env.physics_dt * env.cfg.decimation
    env.extras["stable_timer"] = torch.where(
        is_stable_now,
        env.extras["stable_timer"] + dt,
        torch.zeros_like(env.extras["stable_timer"])
    )

    # 判定是否达到目标时长
    return env.extras["stable_timer"] > standing_time

def get_success_reward(env: ManagerBasedRLEnv, term_keys: str) -> torch.Tensor:
    """检查是否触发了特定的成功终止条件"""
    return env.termination_manager.get_term(term_keys)

def root_vel_z_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """专门惩罚 Z 轴方向的速度"""
    # 获取根部速度 (num_envs, 3) -> [vx, vy, vz]
    vel = env.scene[asset_cfg.name].data.root_lin_vel_w
    # 只取 Z 轴：vel[:, 2]
    return torch.square(vel[:, 2])

def joint_torques_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """计算机器人所有关节施加扭矩的平方和"""
    # 从 data.applied_torques 获取数据，通常形状为 (num_envs, num_joints)
    torques = env.scene[asset_cfg.name].data.applied_torque
    return torch.sum(torch.square(torques), dim=-1)

def joint_vel_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """计算机器人所有关节速度的平方和"""
    # 从 data.joint_vel 获取数据
    vel = env.scene[asset_cfg.name].data.joint_vel
    return torch.sum(torch.square(vel), dim=-1)

def joint_pos_limits_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """惩罚关节接近或超过限位"""
    # 获取关节位置 (num_envs, num_joints)
    joint_pos = env.scene[asset_cfg.name].data.joint_pos
    # 获取限位 (num_joints, 2) -> [lower, upper]
    limits = env.scene[asset_cfg.name].data.soft_joint_pos_limits
    lower_limits = limits[..., 0]
    upper_limits = limits[..., 1]

    # 计算超出限位的部分
    out_of_lower = torch.clamp(lower_limits - joint_pos, min=0.0)
    out_of_upper = torch.clamp(joint_pos - upper_limits, min=0.0)

    # 返回超出量的平方和
    return torch.sum(torch.square(out_of_lower + out_of_upper), dim=-1)

# --- 2. 配置类定义 ---

## 1. 定义与你的类一致的关节列表 (按照 ROBOT_MOTORS 的顺序)
T1_JOINT_NAMES = [
    'Left_Hip_Pitch', 'Right_Hip_Pitch',
    'Left_Hip_Roll', 'Right_Hip_Roll',
    'Left_Hip_Yaw', 'Right_Hip_Yaw',
    'Left_Knee_Pitch', 'Right_Knee_Pitch',
    'Left_Ankle_Pitch', 'Right_Ankle_Pitch',
    'Left_Ankle_Roll', 'Right_Ankle_Roll'
]

@configclass
class T1ObservationCfg:
    """观察值空间配置：严格对应你的 Robot 基类数据结构"""

    @configclass
    class PolicyCfg(ObsGroup):
        concatenate_terms = True
        enable_corruption = False

        # --- 状态量 (对应你的 Robot 类属性) ---

        # 1. 基体线速度 (accelerometer 相关的速度项)
        base_lin_vel = ObsTerm(func=mdp.base_lin_vel)

        # 2. 角速度 (对应你的 gyroscope 属性: degrees/s -> IsaacLab 默认为 rad/s)
        base_ang_vel = ObsTerm(func=mdp.base_ang_vel)

        # 3. 重力投影 (对应 global_orientation_euler/quat 相关的姿态感知)
        projected_gravity = ObsTerm(func=mdp.projected_gravity)

        # 4. 关节位置 (对应 motor_positions)
        # 使用 joint_pos_rel 获取相对于默认姿态的偏差，显式指定关节顺序
        joint_pos = ObsTerm(
            func=mdp.joint_pos_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)}
        )

        # 5. 关节速度 (对应 motor_speeds)
        joint_vel = ObsTerm(
            func=mdp.joint_vel_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)}
        )

        # 6. 上一次的动作 (对应 motor_targets)
        actions = ObsTerm(func=mdp.last_action)

    policy = PolicyCfg()


@configclass
class T1EventCfg:
    """重置时的随机姿态：背躺、趴着、侧躺"""
    reset_robot_rotation = EventTerm(
        func=mdp.reset_root_state_uniform,
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "pose_range": {
                "roll": (0, 0),#(-1.57, 1.57),
                "pitch": (1.57, 1.57),#(-1.57, 1.57),
                "yaw": (0, 0),#(-3.14, 3.14),
                "x": (0.0, 0.0),
                "y": (0.0, 0.0),
                "z": (0.15, 0.15),
            },
            "velocity_range": {},
        },
        mode="reset",
    )


@configclass
class T1ActionCfg:
    """动作空间"""
    joint_pos = JointPositionActionCfg(
        asset_name="robot",
        joint_names=T1_JOINT_NAMES,
        scale=0.5,
        use_default_offset=True
    )


@configclass
class T1GetUpRewardCfg:
    """优化后的奖励函数：抑制跳跃，引导稳健起身"""

    # 1. 高度引导 (改为平滑的指数奖励)
    # 相比 root_height_below_minimum，这个函数会让机器人越接近目标高度得分越高，且曲线平稳
    height_tracking = RewTerm(
        func=mdp.root_height_below_minimum,  # 如果没有自定义函数，保留这个但调低权重
        weight=2.0,  # 降低权重，防止“弹射”
        params={
            "minimum_height": 1.05,
            "asset_cfg": SceneEntityCfg("robot", body_names="H2"),
        }
    )

    # 2. 姿态奖 (保持不变，这是核心)
    upright = RewTerm(func=mdp.flat_orientation_l2, weight=5.0)

    # 3. 稳定性引导 (增加对速度的惩罚，抑制跳跃)
    # 惩罚过大的垂直速度，防止“跳起”
    root_vel_z_penalty = RewTerm(
        func=root_vel_z_l2_local,  # 使用本地函数
        weight=-5,
        params={"asset_cfg": SceneEntityCfg("robot")}  # 传入资产配置
    )

    feet_contact = RewTerm(
        func=mdp.contact_forces,
        weight=0.5,
        params={
            "sensor_cfg": SceneEntityCfg("feet_contact_sensor"),
            "threshold": 1.0
        }
    )

    # 4. 关节与能量约束 (防止 NaN 和乱跳的关键)
    joint_vel = RewTerm(
        func=joint_vel_l2_local,
        weight=-1,
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    applied_torque = RewTerm(
        func=joint_torques_l2_local,
        weight=-1.0e-2,
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    # 5. 动作平滑 (非常重要)
    action_rate = RewTerm(
        func=mdp.action_rate_l2,
        weight=-1.0 # 增大权重，强制动作连贯
    )

    # 6. 软限位惩罚：防止关节撞击
    joint_limits = RewTerm(
        func=joint_pos_limits_l2_local,
        weight=-10.0,
        params = {"asset_cfg": SceneEntityCfg("robot")}
    )

    # 7. 时间惩罚 (强制效率)
    # 每一帧都扣除固定分数，迫使机器人尽快达成 is_success 以停止扣分
    time_penalty = RewTerm(
        func=mdp.is_alive,
        weight=-1.2
    )

    # 8. 核心终点奖励
    is_success = RewTerm(
        func=get_success_reward,
        weight=500.0,  # 成功奖励可以给高点，但前提是动作要平稳
        params={"term_keys": "standing_success"}
    )

@configclass
class T1GetUpTerminationsCfg:
    """终止条件：站稳即算任务完成，且包含强制超时重置"""

    # --- 关键：必须显式添加这一行，episode_length_s 才会生效 ---
    time_out = DoneTerm(func=mdp.time_out)

    # 失败：跌倒 (Trunk 倾斜过大)
    # limit_angle 是弧度，1.0 约等于 57度，如果想严格点可以调小
    base_crash = DoneTerm(func=mdp.bad_orientation, params={"limit_angle": 1.0})

    # 成功：满足“头部高度”和“姿态要求”，且维持 0.8 秒
    standing_success = DoneTerm(
        func=is_standing_still, # 确保你已经把这个函数里的 current_height 改成了 H2 的 Z 轴
        params={
            # T1 头部 (H2) 站直高度约 1.15-1.2m，设为 1.1m 比较稳健
            "minimum_height": 1.05,
            # 姿态误差 (投影重力分量)，0.15 约等于 8.6 度，要求很直
            "max_angle_error": 1.0,
            # 维持时间
            "standing_time": 0.8
        }
    )


@configclass
class T1EnvCfg(ManagerBasedRLEnvCfg):
    """主环境配置"""
    scene = T1SceneCfg(num_envs=32768, env_spacing=2.5)

    def __post_init__(self):
        super().__post_init__()
        # 初始高度设低，配合随机旋转事件实现“从地上爬起来”
        self.scene.robot.init_state.pos = (0.0, 0.0, 0.4)

    observations = T1ObservationCfg()
    rewards = T1GetUpRewardCfg()
    terminations = T1GetUpTerminationsCfg()
    events = T1EventCfg()
    actions = T1ActionCfg()

    episode_length_s = 10.0  # 3秒强制重置
    decimation = 4  # 控制频率