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:

    head_idx, _ = env.scene["robot"].find_bodies("H2")
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")

    curr_head_h = torch.clamp(env.scene["robot"].data.body_state_w[:, head_idx[0], 2], 0.0, 2.0)
    curr_pelvis_h = torch.clamp(env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2], 0.0, 2.0)

    head_score = torch.tanh(curr_head_h / (min_head_height + 1e-6) * 2.0)
    pelvis_score = torch.tanh(curr_pelvis_h / (min_pelvis_height + 1e-6) * 2.0)
    height_reward = (head_score + pelvis_score) / 2.0

    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    foot_forces_z = torch.sum(contact_sensor.data.net_forces_w[:, :, 2], dim=-1)
    force_weight = torch.sigmoid((foot_forces_z - force_threshold) / 5.0)
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]
    vel_penalty = torch.exp(-2.0 * torch.clamp(torch.abs(root_vel_z) - max_v_z, min=0.0))

    influence_weight = torch.clamp((curr_pelvis_h - 0.2) / 0.4, min=0.0, max=1.0)
    combined_reward = height_reward * ((1.0 - influence_weight) + influence_weight * force_weight * vel_penalty)

    return combined_reward


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

def joint_deviation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """计算关节相对于默认姿态(default_joint_pos)的偏差平方和"""
    # 获取当前关节位置相对于默认位置的差值
    # mdp.joint_pos_rel 返回的是 (current_pos - default_pos)
    diff = mdp.joint_pos_rel(env, asset_cfg)
    return torch.sum(torch.square(diff), dim=-1)

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

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',
    # 手臂
    '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'
]


@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.1, 0.2),
            },
            "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=lambda env, keys: env.termination_manager.get_term(keys).float(),
        weight=1000.0,
        params={"keys": "standing_success"}
    )

@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