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']
        )
        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 !")



    @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])
