from roboimi.assets.robots.arm_base import *
import os


class VX300s_insert_ee(ArmBase):
    def __init__(self):
        super().__init__(
            name="vx300s",
            urdf_path=None,
            xml_path="./assets/models/manipulators/vx300s/bi_vx300s_insertion_ee.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.joint_index = ['vx300s_left/waist', 'vx300s_left/shoulder', 'vx300s_left/elbow', 'vx300s_left/forearm_roll', 'vx300s_left/wrist_angle','vx300s_left/wrist_rotate']
        self.left_arm.gripper_index = ['vx300s_left/left_finger','vx300s_left/right_finger']
        self.left_arm.actuator_index = ['vx300s_left/left_finger_t', 'vx300s_left/right_finger_t']
        self.left_arm.setArmInitPose(np.array([0.5,0.0,0.0,0.0,0.0,0.0]))
        self.left_arm.ee_mocap = 0
        self.arms.append(self.left_arm)
        self.right_arm = self.Arm(self,'single', self.urdf_path)
        self.right_arm.joint_index = ['vx300s_right/waist', 'vx300s_right/shoulder', 'vx300s_right/elbow', 'vx300s_right/forearm_roll', 'vx300s_right/wrist_angle','vx300s_right/wrist_rotate']
        self.right_arm.gripper_index = ['vx300s_right/left_finger','vx300s_right/right_finger']
        self.right_arm.actuator_index = ['vx300s_right/left_finger_t', 'vx300s_right/right_finger_t']
        self.right_arm.setArmInitPose(np.array([-0.5,0.0,0.0,0.0,0.0,0.0]))
        self.right_arm.ee_mocap = 1
        self.arms.append(self.right_arm)
        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num

        # joint PD
        # self.kp = 500 * np.ones(self.jnt_num)
        # self.kd = 44.57 * np.ones(self.jnt_num)

        # print("Dual VX300s init !")

    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        return np.array([-0.5, 0.0, 0.0, 0.0, 0.0, 0.0])




class VX300s_transfer_ee(ArmBase):
    def __init__(self):
        super().__init__(
            name="vx300s",
            urdf_path=None,
            xml_path="./assets/models/manipulators/vx300s/bi_vx300s_transfer_ee.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.joint_index = ['vx300s_left/waist', 'vx300s_left/shoulder', 'vx300s_left/elbow', 'vx300s_left/forearm_roll', 'vx300s_left/wrist_angle','vx300s_left/wrist_rotate']
        self.left_arm.actuator_index = ['vx300s_left/left_finger_t', 'vx300s_left/right_finger_t']
        self.left_arm.setArmInitPose(self.init_qpos)
        self.left_arm.ee_mocap = 0
        self.arms.append(self.left_arm)
        self.right_arm = self.Arm(self,'single', self.urdf_path)
        self.right_arm.joint_index = ['vx300s_right/waist', 'vx300s_right/shoulder', 'vx300s_right/elbow', 'vx300s_right/forearm_roll', 'vx300s_right/wrist_angle','vx300s_right/wrist_rotate']
        self.right_arm.actuator_index = ['vx300s_right/left_finger_t', 'vx300s_right/right_finger_t']
        self.right_arm.setArmInitPose(self.init_qpos)
        self.right_arm.ee_mocap = 1
        self.arms.append(self.right_arm)
        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num

        print("VX300s left init !")

    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        return np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])



class VX300s_insert_jnt(ArmBase):
    def __init__(self):
        super().__init__(
            name="vx300s",
            urdf_path=None,
            xml_path="./assets/models/manipulators/vx300s/bi_vx300s_insert_jnt.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.joint_index = ['vx300s_left/waist', 'vx300s_left/shoulder', 'vx300s_left/elbow', 'vx300s_left/forearm_roll', 'vx300s_left/wrist_angle','vx300s_left/wrist_rotate']
        self.left_arm.gripper_index = ['vx300s_left/left_finger','vx300s_left/right_finger']
        self.left_arm.actuator_index = ['vx300s_left/waist_t',
                                        'vx300s_left/shoulder_t',
                                        'vx300s_left/elbow_t',
                                        'vx300s_left/forearm_roll_t',
                                        'vx300s_left/wrist_angle_t',
                                        'vx300s_left/wrist_rotate_t',
                                        'vx300s_left/left_finger_t', 
                                        'vx300s_left/right_finger_t']
        self.left_arm.setArmInitPose(self.init_qpos)
        self.arms.append(self.left_arm)
        self.right_arm = self.Arm(self,'single', self.urdf_path)
        self.right_arm.joint_index = ['vx300s_right/waist', 'vx300s_right/shoulder', 'vx300s_right/elbow', 'vx300s_right/forearm_roll', 'vx300s_right/wrist_angle','vx300s_right/wrist_rotate']
        self.right_arm.gripper_index = ['vx300s_right/left_finger','vx300s_right/right_finger']
        self.right_arm.actuator_index = ['vx300s_right/waist_t',
                                         'vx300s_right/shoulder_t',
                                         'vx300s_right/elbow_t',
                                         'vx300s_right/forearm_roll_t',
                                         'vx300s_right/wrist_angle_t',
                                         'vx300s_right/wrist_rotate_t',
                                         'vx300s_right/left_finger_t',
                                         'vx300s_right/right_finger_t']
        self.right_arm.setArmInitPose(self.init_qpos)
        self.arms.append(self.right_arm)
        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num

        # joint PD
        # self.kp = 500 * np.ones(self.jnt_num)
        # self.kd = 44.57 * np.ones(self.jnt_num)

        # print("Dual VX300s init !")

    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        return np.array([-0.5, 0.0, 0.0, 0.0, 0.0, 0.0])
    


class VX300s_transfer_jnt(ArmBase):
    def __init__(self):
        super().__init__(
            name="vx300s",
            urdf_path=None,
            xml_path="./assets/models/manipulators/vx300s/bi_vx300s_transfer_jnt.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.joint_index = ['vx300s_left/waist', 'vx300s_left/shoulder', 'vx300s_left/elbow', 'vx300s_left/forearm_roll', 'vx300s_left/wrist_angle','vx300s_left/wrist_rotate']
        self.left_arm.actuator_index = ['vx300s_left/waist_t',
                                        'vx300s_left/shoulder_t',
                                        'vx300s_left/elbow_t',
                                        'vx300s_left/forearm_roll_t',
                                        'vx300s_left/wrist_angle_t',
                                        'vx300s_left/wrist_rotate_t',
                                        'vx300s_left/left_finger_t', 
                                        'vx300s_left/right_finger_t']
        self.left_arm.setArmInitPose(self.init_qpos)
        self.left_arm.ee_mocap = 0
        self.arms.append(self.left_arm)
        self.right_arm = self.Arm(self,'single', self.urdf_path)
        self.right_arm.joint_index = ['vx300s_right/waist', 'vx300s_right/shoulder', 'vx300s_right/elbow', 'vx300s_right/forearm_roll', 'vx300s_right/wrist_angle','vx300s_right/wrist_rotate']
        self.right_arm.actuator_index = ['vx300s_right/waist_t',
                                         'vx300s_right/shoulder_t',
                                         'vx300s_right/elbow_t',
                                         'vx300s_right/forearm_roll_t',
                                         'vx300s_right/wrist_angle_t',
                                         'vx300s_right/wrist_rotate_t',
                                         'vx300s_right/left_finger_t',
                                         'vx300s_right/right_finger_t']
        self.right_arm.setArmInitPose(self.init_qpos)
        self.right_arm.ee_mocap = 1
        self.arms.append(self.right_arm)
        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num

        print("VX300s left init !")

    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        return np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])