import mujoco
import numpy as np
import src.mujoco_viewer as mujoco_viewer
import src.pinocchio_kinematic as pinocchio_kinematic
import src.utils as utils

class RobotController(mujoco_viewer.CustomViewer):    
    def __init__(self, scene_path, arm_path):
        super().__init__(scene_path, distance=3.5, azimuth=180, elevation=-90)
        self.scene_path = scene_path
        self.arm_path = arm_path
        
        self.arm = pinocchio_kinematic.Kinematics("ee_center_body")
        self.arm.buildFromMJCF(arm_path)
        self.end_effector_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "ee_center_body")
        self.last_dof = None
        
    def runBefore(self):
        self.model.opt.timestep = 0.005
        # 设定初始位置
        self.initial_pos = self.model.key_qpos[0]  
        print("Initial position: ", self.initial_pos)
        for i in range(self.model.nq):
            self.data.qpos[i] = self.initial_pos[i]
        ee_pos = self.data.body(self.end_effector_id).xpos.copy()
        # self.x = ee_pos[0]
        # self.y = ee_pos[1]
        # self.z = ee_pos[2]
        self.x = 0.1
        self.y = 0.0
        self.z = 0.4
        self.last_dof = self.data.qpos[:9].copy()
    
    def runFunc(self):
        self.x += 0.001
        if self.x > 0.5:
            self.x = 0.5
        tf = utils.transform2mat(self.x, self.y, self.z, np.pi, 0, 0)
        self.dof, info = self.arm.ik(tf, current_arm_motor_q=self.last_dof)
        print("ik result", info["success"])
        self.last_dof = self.dof
        self.data.ctrl[:7] = self.dof[:7]
        print("ee pos", self.getBodyPoseByName("link7"))


if __name__ == "__main__":
    SCENE_XML_PATH = 'model/franka_emika_panda/scene_pos.xml'
    ARM_XML_PATH = 'model/franka_emika_panda/panda_pos.xml'
    robot = RobotController(SCENE_XML_PATH, ARM_XML_PATH)
    robot.run_loop()
