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


class DianaMed(ArmBase):
    def __init__(self):
        self.robot = RobotGenerator(
            name='singlediana',
            scene='default',
            chassis=None,
            manipulator='DianaMed',
            gripper=None,
            g2m_body=['0_link7']
        )
        self.load_aruco()
        super().__init__(
            name="diana_med",
            urdf_path=os.path.join(os.path.dirname(__file__), "../models/manipulators/DianaMed/DianaMed.urdf"),
            xml_path="./assets/singlediana.xml",
            gripper=None
        )
        self.single_arm = self.Arm(self, 'single',self.urdf_path)
        self.single_arm.joint_index = ['0_j1', '0_j2', '0_j3', '0_j4', '0_j5', '0_j6', '0_j7']
        self.single_arm.actuator_index = ['0_a1', '0_a2', '0_a3', '0_a4', '0_a5', '0_a6', '0_a7']
        self.single_arm.setArmInitPose(self.init_qpos)
        self.arms.append(self.single_arm)

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

        print("DianaMed init !")

    def load_aruco(self):
        self.robot.add_texture('aruco', type='2d',
                               file=os.path.join(os.path.dirname(__file__), '../textures/aruco.png'))
        self.robot.add_material('aruco', texture='aruco', texrepeat='1 1', texuniform='false')
        self.robot.add_body(node='worldbody',name='aruco')
        self.robot.add_geom(node='aruco', name='aruco_box', pos='0.919 0 1.27', mass='0.001',
                            euler="0 -1.57 0", size='0.05 0.05 0.001', type='box', material='aruco')
        self.robot.add_joint(node='aruco', name='aruco_x' ,type='slide' ,axis='1 0 0')
        self.robot.add_joint(node='aruco', name='aruco_y' ,type='slide' ,axis='0 1 0')
        self.robot.add_joint(node='aruco', name='aruco_z' ,type='slide' ,axis='0 0 1')




    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        # return np.array([0.0, -np.pi / 4.0, 0.0, np.pi / 2.0, 0.00, np.pi / 4.0, 0.0])
        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])


class BiDianaMed(ArmBase):
    def __init__(self):
        super().__init__(
            name="Bidiana",
            urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf",
            xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.set_Arm_base_link('left_base_link')
        self.left_arm.set_Arm_ee_link('left_link7')
        self.left_arm.InitKDL
        self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7']
        self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left']
        self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left']
        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.set_Arm_base_link('right_base_link')
        self.right_arm.set_Arm_ee_link('right_link7')
        self.right_arm.InitKDL
        self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7']
        self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right']
        self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right']
        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("Double Diana init !")

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