import random
import numpy
import numpy as np
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 standing_with_feet_reward(
        env: ManagerBasedRLEnv,
        min_head_height: float,
        min_pelvis_height: float,
        sensor_cfg: SceneEntityCfg,
        force_threshold: float = 20.0,
        max_v_z: float = 0.5
) -> torch.Tensor:
    # 增加防护：从场景中安全获取 body 索引
    head_idx, _ = env.scene["robot"].find_bodies("H2")
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")

    # 1. 高度奖励：使用更稳定的归一化，限制范围在 [0, 1]
    curr_head_h = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
    curr_pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]

    # 使用 sigmoid 或简单的 min-max 映射，避免除以极小值
    head_score = torch.clamp(curr_head_h / min_head_height, 0.0, 1.2)
    pelvis_score = torch.clamp(curr_pelvis_h / min_pelvis_height, 0.0, 1.2)
    height_reward = (head_score + pelvis_score) / 2.0

    # 2. 足部受力：增加对 NaN 的防御
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    # 某些步数传感器可能未初始化，加个判空
    if contact_sensor is None: return torch.zeros(env.num_envs, device=env.device)

    foot_forces_z = torch.sum(contact_sensor.data.net_forces_w[:, :, 2], dim=-1)
    # 对巨大的冲击力做剪裁，防止 sigmoid 输入过大
    foot_forces_z = torch.clamp(foot_forces_z, 0.0, 500.0)
    force_weight = torch.sigmoid((foot_forces_z - force_threshold) / 5.0)

    # 3. 垂直速度惩罚：使用更平滑的惩罚
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]
    vel_penalty = torch.exp(-torch.abs(root_vel_z) / max_v_z)

    # 逻辑组合：高度 * 稳定性
    return height_reward * (0.5 + 0.5 * force_weight * vel_penalty)


def universal_arm_support_reward(
        env: ManagerBasedRLEnv,
        sensor_cfg: SceneEntityCfg,
        height_threshold: float = 0.60,
        min_force: float = 15.0
) -> torch.Tensor:
    """
    通用手臂支撑奖励：同时支持仰卧起坐支撑和俯卧撑起。
    逻辑：只要手臂有向上的推力，且身体正在向上移动，就给奖。
    """
    # 1. 获取传感器数据
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    if contact_sensor is None:
        return torch.zeros(env.num_envs, device=env.device)

    # 获取所有定义的手臂/手部 link 的垂直总受力 (World Z)
    # net_forces_w 形状: (num_envs, num_bodies, 3)
    arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    # 取所有受力点的最大值或平均值，代表支撑强度
    max_arm_force = torch.max(arm_forces_z, dim=-1)[0]

    # 2. 获取状态数据
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    pelvis_pos_z = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]

    # 3. 计算奖励项
    # A. 受力奖励：鼓励手部与地面产生大于 min_force 的推力
    # 使用 tanh 归一化，防止力矩过大导致奖励爆炸 (NaN 风险)
    force_reward = torch.tanh(torch.clamp(max_arm_force - min_force, min=0.0) / 50.0)

    # B. 速度引导：只有当机器人正在“向上起”时，支撑奖励才翻倍
    # 这样可以防止它趴在地上乱按手骗分
    velocity_factor = torch.clamp(root_vel_z, min=0.0, max=2.0)

    # C. 姿态惩罚回避：
    # 不再检查手是否在盆骨下方，而是检查手是否“在干活”
    # 只要受力足够大，就认为是在支撑
    is_supporting = (max_arm_force > min_force).float()

    # 4. 阶段性退出机制 (Curriculum)
    # 当盆骨高度超过 height_threshold (0.6m) 时，奖励线性消失
    # 强迫机器人最终依靠腿部力量平衡，而不是一直扶着地
    height_fade = torch.clamp((height_threshold - pelvis_pos_z) / 0.15, min=0.0, max=1.0)

    # 最终组合
    # 逻辑：受力 * (1 + 垂直速度) * 高度衰减
    total_reward = force_reward * (1.0 + 2.0 * velocity_factor) * is_supporting * height_fade

    return total_reward

def is_standing_still(
        env: ManagerBasedRLEnv,
        min_head_height: float,
        min_pelvis_height: float,
        max_angle_error: float,
        standing_time: float,
        velocity_threshold: float = 0.15
) -> torch.Tensor:
    """判定逻辑：双高度达标 + 躯干垂直 + 全身静止"""
    head_idx, _ = env.scene["robot"].find_bodies("H2")
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")

    current_head_h = env.scene["robot"].data.body_state_w[:, head_idx[0], 2]
    current_pelvis_h = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]

    gravity_error = torch.norm(env.scene["robot"].data.projected_gravity_b[:, :2], dim=-1)
    root_vel_norm = torch.norm(env.scene["robot"].data.root_lin_vel_w, dim=-1)

    # 判定条件：头够高 且 盆骨够高 且 垂直误差小 且 速度低
    is_stable_now = (
            (current_head_h > min_head_height) &
            (current_pelvis_h > min_pelvis_height) &
            (gravity_error < max_angle_error) &
            (root_vel_norm < velocity_threshold)
    )

    if "stable_timer" not in env.extras:
        env.extras["stable_timer"] = torch.zeros(env.num_envs, device=env.device)

    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

# --- 2. 配置类 ---

T1_JOINT_NAMES = [

    'AAHead_yaw', 'Head_pitch',
                
    'Left_Shoulder_Pitch', 'Left_Shoulder_Roll', 'Left_Elbow_Pitch', 'Left_Elbow_Yaw',
    'Right_Shoulder_Pitch', 'Right_Shoulder_Roll', 'Right_Elbow_Pitch', 'Right_Elbow_Yaw',

    'Waist',
    
    '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:
    @configclass
    class PolicyCfg(ObsGroup):
        concatenate_terms = True
        base_lin_vel = ObsTerm(func=mdp.base_lin_vel)
        base_ang_vel = ObsTerm(func=mdp.base_ang_vel)
        projected_gravity = ObsTerm(func=mdp.projected_gravity)
        root_pos = ObsTerm(func=mdp.root_pos_w)
        joint_pos = ObsTerm(func=mdp.joint_pos_rel,
                            params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)})
        joint_vel = ObsTerm(func=mdp.joint_vel_rel,
                            params={"asset_cfg": SceneEntityCfg("robot", joint_names=T1_JOINT_NAMES)})
        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": (-1.57, 1.57),   # 左右侧卧
                "pitch": tuple(numpy.array([1.4, 1.6], dtype=np.float32) * random.choice([-1 , 1])),  # 仰卧/俯卧
                "yaw": (-3.14, 3.14), # 全向旋转
                "x": (0.0, 0.0),
                "y": (0.0, 0.0),
                "z": (0.3, 0.4),
            },
            "velocity_range": {},
        },
        mode="reset",
    )


@configclass
class T1ActionCfg:
    """关键修改：降低 scale 让动作变丝滑，增大阻尼效果"""
    joint_pos = JointPositionActionCfg(
        asset_name="robot",
        joint_names=T1_JOINT_NAMES,
        scale=0.5,
        use_default_offset=True
    )


@configclass
class T1GetUpRewardCfg:
    # 1. 姿态基础奖 (引导身体变正)
    upright = RewTerm(func=mdp.flat_orientation_l2, weight=2.0)

    # 2. 【条件高度奖】：双高度判定（头+盆骨），且必须脚踩地
    height_with_feet = RewTerm(
        func=standing_with_feet_reward,
        weight=20.0,  # 作为核心引导，增加权重
        params={
            "min_head_height": 1.10,
            "min_pelvis_height": 0.7,
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"]),
            "force_threshold": 20.0,
            "max_v_z": 0.3
        }
    )

    # 3. 手臂撑地奖：辅助脱离地面阶段
    arm_push_support = RewTerm(
        func=universal_arm_support_reward,
        weight=15.0,  # 显著增加权重（从 3.0 提到 15.0），让它成为起步的关键
        params={
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link", "AL3", "AR3"]),
            "height_threshold": 0.65,  # 躯干升到 0.6m 前都鼓励手臂用力
            "min_force": 8.0  # 只要有 15N 的力就触发
        }
    )

    # 4. 关节限位惩罚 (新增：防止关节撞死导致数值问题)
    joint_limits = RewTerm(
        func=mdp.joint_pos_limits,
        weight=-1.0,
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    # 5. 时间惩罚 (强制效率)
    time_penalty = RewTerm(
        func=mdp.is_alive,
        weight=-1.2
    )

    # 6. 成功终极大奖
    is_success = RewTerm(
        func=is_standing_still,
        weight=800.0,
        params={
            "min_head_height": 1.05,
            "min_pelvis_height": 0.75,
            "max_angle_error": 0.3,
            "standing_time": 0.2,
            "velocity_threshold": 0.5
        }
    )

@configclass
class T1GetUpTerminationsCfg:
    time_out = DoneTerm(func=mdp.time_out)

    # 失败判定：躯干倾斜超过 45 度重置
    #base_crash = DoneTerm(func=mdp.bad_orientation, params={"limit_angle": 0.785})

    # 成功判定：双高度 + 稳定
    standing_success = DoneTerm(
        func=is_standing_still,
        params={
            "min_head_height": 1.05,
            "min_pelvis_height": 0.75,
            "max_angle_error": 0.3,
            "standing_time": 0.2,
            "velocity_threshold": 0.5
        }
    )


@configclass
class T1EnvCfg(ManagerBasedRLEnvCfg):
    scene = T1SceneCfg(num_envs=8192, env_spacing=2.5)  # 5090 性能全开

    def __post_init__(self):
        super().__post_init__()
        self.scene.robot.init_state.pos = (0.0, 0.0, 0.2)

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

    episode_length_s = 6.0
    decimation = 4