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,
        max_v_z: float = 0.25
) -> torch.Tensor:
    """
    平滑切换的高度奖励：
    低高度 -> 纯高度引导
    高高度 -> 高度 + 足底力 + 速度约束
    """
    # 1. 获取基本状态
    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]

    # 2. 计算基础高度得分 (0.0 - 1.0)
    head_score = torch.clamp(current_head_h / min_head_height, max=1.0)
    pelvis_score = torch.clamp(current_pelvis_h / min_pelvis_height, max=1.0)
    combined_height_score = (head_score + pelvis_score) / 2.0

    # 3. 计算足底力判定
    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).float()

    # 4. 计算速度惩罚 (抑制乱跳)
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]
    vel_penalty_factor = torch.exp(-4.0 * torch.clamp(torch.abs(root_vel_z) - max_v_z, min=0.0))

    # --- 核心逻辑切换 ---
    # 定义一个“过渡高度” (例如盆骨达到 0.4m)
    transition_h = 0.4

    # 如果高度很低：给纯高度奖，诱导它向上动
    low_height_reward = combined_height_score

    # 如果高度较高：给 综合奖 (高度 * 速度限制 * 必须踩地)
    high_height_reward = combined_height_score * vel_penalty_factor * is_feet_on_ground

    return torch.where(current_pelvis_h < transition_h, low_height_reward, high_height_reward)

def arm_push_up_reward(
        env: ManagerBasedRLEnv,
        sensor_cfg: SceneEntityCfg,
        height_threshold: float = 0.5
) -> 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)

    # 获取躯干(Pelvis/Trunk)当前高度
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    current_height = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]

    # 1. 手臂受力奖：只有在躯干较低时，鼓励手臂产生向上的反作用力
    arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    max_arm_force = torch.max(arm_forces_z, dim=-1)[0]
    push_reward = torch.clamp(max_arm_force / 50.0, max=1.0)  # 归一化受力奖励

    # 2. 躯干高度线性引导：只要在 height_threshold 以下，越高分越高
    # 这解决了“动不动奖励都一样”的问题
    lifting_reward = torch.clamp(current_height / height_threshold, max=1.0)

    # 只有在还没站起来（低于阈值）的时候，才给这两个组合奖励
    return torch.where(current_height < height_threshold,
                       push_reward + lifting_reward,
                       torch.zeros_like(push_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 feet_airtime_penalty_local(
        env: ManagerBasedRLEnv,
        sensor_cfg: SceneEntityCfg,
        threshold: float = 1.0
) -> torch.Tensor:
    """
    自定义滞空惩罚逻辑：
    如果脚部的垂直合力小于阈值，说明脚离地了。
    返回一个 Tensor，离地时为 1.0，着地时为 0.0。
    """
    # 1. 获取传感器对象
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)

    if contact_sensor is None:
        # 如果没搜到传感器，返回全 0，防止程序崩溃
        return torch.zeros(env.num_envs, device=env.device)

    # 2. 获取触地力 (num_envs, num_bodies_in_sensor, 3)
    # 我们取所有被监测 Body (左右脚) 的 Z 轴推力
    # 如果所有脚的力都小于 threshold，判定为“完全腾空”
    foot_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    is_in_air = torch.all(foot_forces_z < threshold, dim=-1)

    return is_in_air.float()

def root_vel_z_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    # 严厉惩罚 Z 轴正向速度（向上窜）
    vel_z = env.scene[asset_cfg.name].data.root_lin_vel_w[:, 2]
    return torch.square(torch.clamp(vel_z, min=0.0))

def joint_pos_rel_l2_local(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    # 获取相对默认位置的偏差 (num_envs, num_joints)
    rel_pos = mdp.joint_pos_rel(env, asset_cfg)
    # 计算平方和 (L2)
    return torch.sum(torch.square(rel_pos), dim=-1)

def strict_feet_contact_reward(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """如果脚不着地，直接给一个很大的负分，强制它必须寻找支点"""
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    # 只要有一只脚没力，就判定为不稳
    foot_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    all_feet_cond = torch.min(foot_forces_z, dim=-1)[0] > 5.0  # 左右脚都要有至少5N的力

    return (~all_feet_cond).float()  # 返回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'
]


@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.2, 0.2),#(-1.57, 1.57),
                "pitch": (-1.6, -1.4),#(-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=5.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=3.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. 抑制跳跃：严厉惩罚向上窜的速度
    root_vel_z_penalty = RewTerm(
        func=root_vel_z_l2_local,
        weight=-50.0,  # 增大负权重
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    # 6. 抑制滞空 (Airtime Penalty)
    feet_airtime = RewTerm(
        func=strict_feet_contact_reward,
        weight=-20.0,  # 加大权重，跳一下扣的分比站起来得的分还多
        params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_foot_link"])}
    )

    joint_vel_penalty = RewTerm(
        func=mdp.joint_vel_l2,
        weight=-0.5,  # 惩罚过快的关节运动
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    action_rate = RewTerm(
        func=mdp.action_rate_l2,
        weight=-0.5,  # 惩罚动作的突变，让动作更丝滑，减少爆发力
    )

    # 惩罚躯干的翻转和俯仰角速度
    base_ang_vel_penalty = RewTerm(
        func=lambda env, asset_cfg: torch.norm(mdp.base_ang_vel(env, asset_cfg), dim=-1),
        weight=-0.1,
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    joint_deviation = RewTerm(
        func=joint_pos_rel_l2_local,
        weight=0.1,  # 权重不要太高，只是为了让它动起来
        params={"asset_cfg": SceneEntityCfg("robot")}
    )

    # 7. 成功终极大奖
    is_success = RewTerm(
        func=lambda env, keys: env.termination_manager.get_term(keys),
        weight=500.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