from __future__ import annotations

import torch
from typing import TYPE_CHECKING

from isaaclab.managers import SceneEntityCfg
from isaaclab.utils.math import euler_xyz_from_quat, quat_apply_inverse, quat_from_euler_xyz

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedEnv


def _single_body_id(asset_cfg: SceneEntityCfg, term_name: str) -> int:
    body_ids = asset_cfg.body_ids
    if body_ids is None:
        raise ValueError(f"{term_name} requires asset_cfg with exactly one body name.")
    if isinstance(body_ids, int):
        return body_ids
    if isinstance(body_ids, slice) or len(body_ids) != 1:
        raise ValueError(f"{term_name} requires exactly one body id, got {body_ids}.")
    return body_ids[0]


def _heading_quat_with_offset(body_quat_w: torch.Tensor, heading_yaw_offset: float = 0.0) -> torch.Tensor:
    roll, pitch, yaw = euler_xyz_from_quat(body_quat_w)
    zeros = torch.zeros_like(yaw)
    return quat_from_euler_xyz(zeros, zeros, yaw - heading_yaw_offset)


def robot_base_lin_vel_b(env: ManagerBasedEnv) -> torch.Tensor:
    """Base linear velocity expressed in the base frame."""
    asset = env.scene["robot"]
    # prefer direct base-frame velocity if available
    if getattr(asset.data, "root_lin_vel_b", None) is not None:
        return asset.data.root_lin_vel_b
    # fallback: rotate world velocity into base frame
    return quat_apply_inverse(asset.data.root_quat_w, asset.data.root_lin_vel_w)


def body_ang_vel_yaw_frame(
    env: ManagerBasedEnv,
    asset_cfg: SceneEntityCfg,
    heading_yaw_offset: float = 0.0,
) -> torch.Tensor:
    """Body angular velocity expressed in a yaw-aligned frame for that body."""
    body_id = _single_body_id(asset_cfg, "body_ang_vel_yaw_frame")
    asset = env.scene[asset_cfg.name]
    body_quat_w = asset.data.body_quat_w[:, body_id, :]
    body_ang_vel_w = asset.data.body_ang_vel_w[:, body_id, :]
    return quat_apply_inverse(_heading_quat_with_offset(body_quat_w, heading_yaw_offset), body_ang_vel_w)


def body_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """Gravity projection in a configured body's local frame."""
    body_id = _single_body_id(asset_cfg, "body_projected_gravity")
    asset = env.scene[asset_cfg.name]
    return quat_apply_inverse(asset.data.body_quat_w[:, body_id, :], asset.data.GRAVITY_VEC_W)