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.utils.math as math_utils

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,
        min_force: float = 20.0
) -> torch.Tensor:
    """
    强化版手臂撑地奖励：
    1. 鼓励手臂产生超过阈值的垂直反作用力。
    2. 当手臂用力且躯干有向上速度时，给予额外加成。
    """
    # 获取手臂传感器数据
    contact_sensor = env.scene.sensors.get(sensor_cfg.name)
    if contact_sensor is None:
        return torch.zeros(env.num_envs, device=env.device)

    # 1. 获取手臂 Z 轴受力 (取所有手臂 Body 的合力或最大力)
    arm_forces_z = contact_sensor.data.net_forces_w[:, :, 2]
    max_arm_force = torch.max(arm_forces_z, dim=-1)[0]

    # 归一化受力奖励：在 20N 到 100N 之间线性增长
    force_reward = torch.clamp((max_arm_force - min_force) / 80.0, min=0.0, max=1.0)

    # 2. 获取躯干高度和垂直速度
    pelvis_idx, _ = env.scene["robot"].find_bodies("Trunk")
    current_height = env.scene["robot"].data.body_state_w[:, pelvis_idx[0], 2]
    root_vel_z = env.scene["robot"].data.root_lin_vel_w[:, 2]

    # 3. 协同奖励：当手臂在用力推，且躯干正在上升时，给高分
    # 只有在高度低于阈值（还在撑起阶段）时生效
    pushing_up_bonus = torch.where(
        (max_arm_force > min_force) & (root_vel_z > 0.05),
        force_reward * (1.0 + root_vel_z * 2.0),  # 速度越快奖励越高
        force_reward
    )

    # 只有在躯干较低时才发放此奖励
    return torch.where(current_height < height_threshold,
                       pushing_up_bonus,
                       torch.zeros_like(pushing_up_bonus))

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表示违规


def reset_root_state_symmetric(
        env: ManagerBasedRLEnv,
        env_ids: torch.Tensor,
        asset_cfg: SceneEntityCfg,
        pose_range: dict,
        velocity_range: dict
):
    """随机对称采样：让机器人随机以 趴/躺/左倾/右倾 及其组合姿态重置"""
    robot = env.scene[asset_cfg.name]
    num_resets = len(env_ids)
    device = env.device

    # 1. 采样 Euler 角 (roll, pitch, yaw)
    # 我们先在 positive range 内采样
    def get_rand(key):
        low, high = pose_range[key]
        return (high - low) * torch.rand(num_resets, device=device) + low

    roll = get_rand("roll")
    pitch = get_rand("pitch")
    yaw = (pose_range["yaw"][1] - pose_range["yaw"][0]) * torch.rand(num_resets, device=device) + pose_range["yaw"][0]

    # 2. 核心：随机符号化 (Random Sign)
    # 50% 概率保持正，50% 概率变负。Roll 和 Pitch 独立随机。
    roll_sign = (torch.randint(0, 2, (num_resets,), device=device) * 2 - 1).float()
    pitch_sign = (torch.randint(0, 2, (num_resets,), device=device) * 2 - 1).float()

    roll *= roll_sign
    pitch *= pitch_sign

    # 3. 应用到状态
    quat = math_utils.quat_from_euler_xyz(roll, pitch, yaw)

    # 获取默认状态并叠加随机位置偏移 (x, y, z)
    root_states = robot.data.default_root_state[env_ids].clone()
    for i, key in enumerate(["x", "y", "z"]):
        low, high = pose_range[key]
        root_states[:, i] += (high - low) * torch.rand(num_resets, device=device) + low

    root_states[:, 3:7] = quat

    # 4. 写入仿真 (同时清除速度，确保是静态开始)
    robot.write_root_state_to_sim(root_states, env_ids)

# --- 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=reset_root_state_symmetric, # 使用上面那个通用的函数
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "pose_range": {
                "roll": (0, 1.57),   # 左右侧卧
                "pitch": (1.4, 1.6),  # 仰卧/俯卧
                "yaw": (-3.14, 3.14), # 全向旋转
                "x": (0.0, 0.0),
                "y": (0.0, 0.0),
                "z": (0.4, 0.5),
            },
            "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=30.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=15.0,  # 显著增加权重（从 3.0 提到 15.0），让它成为起步的关键
        params={
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*_hand_link", "AL3", "AR3"]),
            "height_threshold": 0.6,  # 躯干升到 0.6m 前都鼓励手臂用力
            "min_force": 15.0  # 只要有 15N 的力就触发
        }
    )

    # 4. 惩罚项
    undesired_contacts = RewTerm(
        func=mdp.undesired_contacts,
        weight=-10.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_sensor", body_names=["H2", "Trunk"]),
            # 注意：此处必须排除手臂相关 link，否则手臂用力时会同时被扣分
            "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