from egogym.embodiments.arms.arm import Arm
from egogym.embodiments.robot import requires_mj
from egogym.utils import get_pose
import numpy as np


class DroidArm(Arm):    
    def __init__(self, control_steps=600, grasping_steps=100, action_space="delta", np_random=None):
        super().__init__(control_steps=control_steps, grasping_steps=grasping_steps, action_space=action_space, np_random=np_random)
        self.name = "droid"
        self.initial_q = np.array([0, -1/5 * np.pi, 0, -4/5 * np.pi, 0, 3/5 * np.pi, 0.0])
        self.joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"]
        self.gripper_actuator_name = "fingers_actuator"
    
    @requires_mj
    def prepare(self, desired_pose, step_fn=None):

        desired_pose[:3,3] += np.array([0.0, 0.6, 0.075])
        if self.add_noise_init:
            position_noise = self.np_random.uniform(-0.03, 0.03, size=3)
            desired_pose[:3, 3] += position_noise

        self.data.ctrl[self.act_ids] = self.initial_q
        step_fn(500)
        gripper_position = self.get_tcp_pose()[:3, 3]
        link_position = get_pose(self.data, "link0", "body")[:3, 3]
        delta_position = gripper_position - link_position
        desired_gripper_position = desired_pose[:3, 3]
        desired_link_position = desired_gripper_position - delta_position
        self.model.body_pos[self.model.body("link0").id] = desired_link_position
        step_fn(100)

    @requires_mj
    def get_grasp(self):
        val = self.data.qpos[self.model.joint("right_driver_joint").qposadr[0]] / 0.8102
        if val < 0.01:
            val = 0.0
        val = np.array([val], dtype=np.float32).reshape(1)
        current = np.array([(np.clip(val, 0.0, 1.0))], dtype=np.float32).reshape(1)[0]
        return current
    
    @requires_mj
    def set_grasp(self, grip_action):
        desired_grasp = 255.0 * grip_action
        self.data.ctrl[self.model.actuator("fingers_actuator").id] = desired_grasp.item()
        if desired_grasp != self.last_grasp:
            self.last_grasp = desired_grasp
            return 1
        return 0