from isaaclab.utils import configclass
from isaaclab.managers.action_manager import ActionTerm
from isaaclab.envs.mdp import JointActionCfg, DifferentialInverseKinematicsActionCfg

from .actions import DefaultJointPositionStaticAction, EnvFrameDiffIKAction, TransformedOneShotDifferentialIKAction
from .gravity_compensated_joint_action_cfg import (  # noqa: F401
    GravityCompensatedJointPositionActionCfg,
    GravityCompensatedRelativeJointPositionActionCfg,
    GravityCompensatedDifferentialIKActionCfg,
)

@configclass
class DefaultJointPositionStaticActionCfg(JointActionCfg):
    """Configuration for the joint position action term.

    See :class:`JointPositionAction` for more details.
    """

    class_type: type[ActionTerm] = DefaultJointPositionStaticAction

    use_default_offset: bool = True


@configclass
class TransformedOneShotDifferentialIKActionCfg(DifferentialInverseKinematicsActionCfg):
    """Configuration for One-Shot Differential IK action term.

    Unlike standard IK which recomputes at every physics step (500 Hz), this computes IK
    ONCE per policy step (10 Hz) and holds the joint target fixed. This makes it equivalent
    to JointPos in terms of control structure, eliminating distillation gap.

    Supports coordinate frame transformations via action_root_offset for policies trained
    in a transformed frame (e.g., 180° flipped base for real-robot alignment).
    """

    class_type: type[ActionTerm] = TransformedOneShotDifferentialIKAction

    action_root_offset: DifferentialInverseKinematicsActionCfg.OffsetCfg | None = None
    """Offset for the action root frame (e.g., 180° flip for real-robot alignment)."""


@configclass
class EnvFrameDiffIKActionCfg(DifferentialInverseKinematicsActionCfg):
    """DiffIK that operates in env frame (world-aligned axes) instead of root body frame.

    Use this when the robot's base link has a baked-in rotation in the USD and you
    want IK commands to align with the environment/world axes.

    Includes singularity protection inspired by anzu's QP-based DiffIK:
    - ``k_vx`` scales down the pose error before solving (like anzu's K_VX gain)
    - ``max_linear_vel`` / ``max_angular_vel`` clamp the desired Cartesian velocity
    - ``max_joint_vel`` clamps the IK output to respect joint velocity limits
    """

    class_type: type[ActionTerm] = EnvFrameDiffIKAction

    k_vx: float = 1.0
    """Proportional gain applied to pose error before IK solve.
    Values < 1.0 (e.g. 0.1) prevent singularity blowup by only closing a fraction
    of the error per sim step. Anzu uses 0.1 for Panda. Default 1.0 preserves
    backward-compatible behavior."""

    max_linear_vel: float | None = None
    """Max linear velocity (m/s) for clamping desired Cartesian translation rate.
    Anzu uses 0.4 m/s for Panda. None disables clamping."""

    max_angular_vel: float | None = None
    """Max angular velocity (rad/s) for clamping desired Cartesian rotation rate.
    Anzu uses ~1.57 rad/s (90 deg/s) for Panda. None disables clamping."""

    max_joint_vel: list[float] | float | None = None
    """Max joint velocity (rad/s) for clamping the IK output delta-q.
    Scalar applies the same limit to all joints; list gives per-joint limits.
    None disables clamping."""

    open_loop: bool = False
    """When True, compute pose error relative to the previous *commanded* pose
    instead of the current sim EE pose. This makes the IK open-loop: it
    dead-reckons from its own commands rather than using sim feedback."""