import isaaclab.sim as sim_utils
from isaaclab.actuators import ImplicitActuatorCfg
from isaaclab.assets.articulation import ArticulationCfg

import numpy as np


####### ROBOT DEFINITIONS #######
ROBOTIQ_2F85_DEFAULT_JOINT_POS = {
    "finger_joint": 0.0,
    "right_outer.*": 0.0,
    "left_outer.*": 0.0,
    "left_inner_finger_knuckle_joint": 0.0,
    "right_inner_finger_knuckle_joint": 0.0,
    "left_inner_finger_joint": -0.785398,
    "right_inner_finger_joint": 0.785398,
}

DROID_DEFAULT_JOINT_POS = {
    "panda_joint1": 0.0,
    "panda_joint2": -1 / 5 * np.pi,
    "panda_joint3": 0.0,
    "panda_joint4": -4 / 5 * np.pi,
    "panda_joint5": 0.0,
    "panda_joint6": 3 / 5 * np.pi,
    "panda_joint7": 0,
    **ROBOTIQ_2F85_DEFAULT_JOINT_POS,
}

DROID_ARTICULATION = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        usd_path=f"./envs/assets/patrick_droid/franka_robotiq_gripper.usd",
        activate_contact_sensors=False,
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=True,
            max_depenetration_velocity=5.0,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=True, solver_position_iteration_count=36, solver_velocity_iteration_count=0
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0), rot=(1, 0, 0, 0), joint_pos=DROID_DEFAULT_JOINT_POS),
    soft_joint_pos_limit_factor=1,
)

ROBOTIQ_2F85 = ArticulationCfg(
    prim_path="{ENV_REGEX_NS}/RobotiqGripper",
    spawn=sim_utils.UsdFileCfg(
        # usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/UniversalRobots/2f85RobotiqGripper/robotiq_2f85_gripper.usd",
        usd_path = "we probably never need this",
        activate_contact_sensors=False,
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=True,
            max_depenetration_velocity=5.0,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False, solver_position_iteration_count=36, solver_velocity_iteration_count=0
        ),
        mass_props=sim_utils.MassPropertiesCfg(mass=0.5),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        pos=(0, 0, 0.1), rot=(1, 0, 0, 0), joint_pos=ROBOTIQ_2F85_DEFAULT_JOINT_POS
    ),
    actuators={
        "gripper": ImplicitActuatorCfg(
            joint_names_expr=["finger_joint"],
            stiffness=17,
            damping=5,
            effort_limit_sim=165,
        ),
        "inner_finger": ImplicitActuatorCfg(
            joint_names_expr=[".*_inner_finger_joint"],
            stiffness=0.2,
            damping=0.02,
            effort_limit_sim=0.5,
        ),
    },
    soft_joint_pos_limit_factor=1,
)

IMPLICIT_DROID = DROID_ARTICULATION.copy()  # type: ignore
IMPLICIT_DROID.actuators = {
    "arm": ImplicitActuatorCfg(
        joint_names_expr=["panda_joint.*"],
        stiffness=400.0,
        damping=80.0,
        effort_limit=87.0,
        velocity_limit=2.175,
    ),
    "gripper": ROBOTIQ_2F85.actuators["gripper"],
    "inner_finger": ROBOTIQ_2F85.actuators["inner_finger"],
}

EXPLICIT_DROID = DROID_ARTICULATION.copy()  # type: ignore
EXPLICIT_DROID.actuators = {
    "arm": ImplicitActuatorCfg(
        joint_names_expr=["panda_joint.*"],
        stiffness=0.0,
        damping=0.0,
        velocity_limit_sim=3.14,
        effort_limit_sim={
            "panda_joint1": 150.0,
            "panda_joint2": 150.0,
            "panda_joint3": 150.0,
            "panda_joint4": 150.0,
            "panda_joint5": 150.0,
            "panda_joint6": 150.0,
            "panda_joint7": 150.0,
        },
    ),
    "gripper": ROBOTIQ_2F85.actuators["gripper"],
    "inner_finger": ROBOTIQ_2F85.actuators["inner_finger"],
}
####### ROBOT DEFINITIONS #######


####### ACTION DEFINITIONS #######
from isaaclab.envs.mdp.actions.actions_cfg import (
    BinaryJointPositionActionCfg,
    JointPositionActionCfg,
    RelativeJointPositionActionCfg,
)
from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg
from isaaclab.utils import configclass
from sim_improvement.environments.mdp import DefaultJointPositionStaticActionCfg, TransformedOneShotDifferentialIKActionCfg

DROID_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    scale=1.0,
    use_default_offset=False,
)

DROID_RELATIVE_JOINT_POSITION: RelativeJointPositionActionCfg = RelativeJointPositionActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    scale=0.02,
    use_zero_offset=True,
)

DROID_IK_RELATIVE_ARM = TransformedOneShotDifferentialIKActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    body_name="robotiq_base_link",
    controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"),
    # scale=0.002,
    scale=1.0,
)

ROBOTIQ_GRIPPER_BINARY_ACTIONS = BinaryJointPositionActionCfg(
    asset_name="robot",
    joint_names=["finger_joint"],
    open_command_expr={"finger_joint": 0.0},
    close_command_expr={"finger_joint": 0.785398},
)

ROBOTIQ_COMPLIANT_JOINTS = DefaultJointPositionStaticActionCfg(
    asset_name="robot", joint_names=["left_inner_finger_joint", "right_inner_finger_joint"]
)


@configclass
class DROIDIkRelativeAction:
    arm = DROID_IK_RELATIVE_ARM
    gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS
    compliant_joints = ROBOTIQ_COMPLIANT_JOINTS

@configclass
class DROIDJointPositionAction:
    arm = DROID_JOINT_POSITION
    gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS
    compliant_joints = ROBOTIQ_COMPLIANT_JOINTS

@configclass
class DROIDRelativeJointPositionAction:
    arm = DROID_RELATIVE_JOINT_POSITION
    gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS
    compliant_joints = ROBOTIQ_COMPLIANT_JOINTS

####### ACTION DEFINITIONS #######



__all__ = ["IMPLICIT_DROID", "EXPLICIT_DROID"]