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 = 30.0
) -> torch.Tensor:
    """
    【双高度条件奖励】：只有脚踩地，且头和躯干同时达到高度，才给予高度奖励。
    """
    # 1. 获取脚部触地力判定
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    foot_forces_z = torch.sum(contact_sensor.data.net_forces_w[:, :, 2], dim=-1)
    is_feet_on_ground = foot_forces_z > force_threshold

    # 2. 获取头部和躯干索引并提取高度
    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]

    # 3. 计算高度达标度 (0.0 - 1.0)
    head_reward = torch.clamp(current_head_h / min_head_height, max=1.2)
    pelvis_reward = torch.clamp(current_pelvis_h / min_pelvis_height, max=1.0)

    # 综合高度奖励（取平均值）
    combined_height_reward = (head_reward + pelvis_reward) / 2.0

    # 4. 逻辑门：脚不着地，奖励为 0；脚着地后，根据高度给分
    return torch.where(is_feet_on_ground, combined_height_reward, torch.zeros_like(combined_height_reward))


def arm_push_up_reward(
        env: ManagerBasedRLEnv,
        sensor_cfg: SceneEntityCfg,
        height_threshold: float = 0.6
) -> torch.Tensor:
    """手臂撑地奖励：辅助机器人从趴/躺状态利用手臂反作用力起身"""
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    if contact_sensor is None: return torch.zeros(env.num_envs, device=env.device)

    arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    max_arm_force = torch.max(arm_forces_z, dim=-1)[0]

    # 当躯干还很低时，鼓励手撑地
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    current_height = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]

    pushing_reward = torch.clamp(max_arm_force, max=200.0) / 100.0
    return torch.where(current_height < height_threshold, pushing_reward, torch.zeros_like(pushing_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 = [
    '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)
        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": (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.0, 0.0),
            },
            "velocity_range": {},
        },
        mode="reset",
    )


@configclass
class T1ActionCfg:
    """关键修改：降低 scale 让动作变丝滑，增大阻尼效果"""
    joint_pos = JointPositionActionCfg(
        asset_name="robot",
        joint_names=T1_JOINT_NAMES,
        scale=0.2,  # 从 0.5 降到 0.2，防止电机暴力抽搐
        use_default_offset=True
    )


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

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

    # 3. 手臂撑地奖：辅助脱离地面阶段
    arm_push_support = RewTerm(
        func=arm_push_up_reward,
        weight=2.0,
        params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link"])}
    )

    # 4. 惩罚项
    undesired_contacts = RewTerm(
        func=mdp.undesired_contacts,
        weight=-5.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["H2", "Trunk", ".*_Left", ".*_Right"]),
            # 注意：这里我们排除了脚部，惩罚大腿/小腿/躯干着地
            "threshold": 1.0
        }
    )

    # 5. 成功终极大奖
    is_success = RewTerm(
        func=lambda env, keys: env.termination_manager.get_term(keys),
        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.15,
            "standing_time": 0.8,
            "velocity_threshold": 0.15
        }
    )


@configclass
class T1EnvCfg(ManagerBasedRLEnvCfg):
    scene = T1SceneCfg(num_envs=16384, 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