from scipy.spatial.transform import Rotation as R
import numpy as np

from egogym.utils import get_pose
from egogym.embodiments.robot import Robot, requires_mj


class FloatingGripper(Robot):    
    def __init__(self, control_steps=50, grasping_steps=1, 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.last_grasp = 0.0
    
    def action_space_shape(self):
        return (17,)

    def setup_mj(self, model, data):
        super().setup_mj(model, data)
        self.last_grasp = 0.0
        self.free_joint_id = model.joint("free_joint").qposadr[0]
    
    @requires_mj
    def prepare(self, desired_pose, step_fn=None):
        if self.add_noise_init:
            position_noise = self.np_random.uniform(-0.03, 0.03, size=3)
            desired_pose[:3, 3] += position_noise
            euler_angles = R.from_matrix(desired_pose[:3, :3]).as_euler("xyz")
            euler_angles[0] += self.np_random.uniform(-0.1, 0.1)
            euler_angles[1] += self.np_random.uniform(-0.05, 0.05)
            euler_angles[2] += self.np_random.uniform(-0.1, 0.1)
            desired_pose[:3, :3] = R.from_euler("xyz", euler_angles).as_matrix()
        self.last_teleop_pose = desired_pose
        self.data.mocap_pos[0] = desired_pose[:3, 3]
        self.data.mocap_quat[0] = np.roll(R.from_matrix(desired_pose[:3, :3]).as_quat(), 1)
        step_fn(self.control_steps)

    @requires_mj
    def move(self, desired_pose):
        self.data.mocap_pos[0] = desired_pose[:3, 3]
        self.data.mocap_quat[0] = np.roll(R.from_matrix(desired_pose[:3, :3]).as_quat(), 1)

    @requires_mj
    def step(self, action, step_fn):
        if len(action.shape) == 2:
            action = action[0]

        pose = action[:16].reshape(4,4)
        grip_action = action[16]
        
        if self.action_space == "delta":
            goal_pose = self.get_camera_pose() @ pose
        else:
            goal_pose = pose
        
        self.move(goal_pose)
        updated = self.set_grasp(grip_action)
        if updated:
            step_fn(self.grasping_steps)
        step_fn(self.control_steps)