from egogym.embodiments.robot import Robot, requires_mj
from gymnasium.spaces import Box
import numpy as np


class Arm(Robot):
    
    def __init__(self, control_steps=2000, grasping_steps=1000, 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.joint_names = []
        self.qpos_ids = None
        self.act_ids = None
        self.gripper_act_id = None
    
    def additional_obs_spaces(self):
        num_joints = len(self.joint_names)
        return {
            "joint_positions": Box(low=-np.inf, high=np.inf, shape=(num_joints,), dtype=np.float32)
        }
    
    def action_space_shape(self):
        return (len(self.joint_names) + 1,)
    
    @requires_mj
    def get_additional_obs(self):
        return {
            "joint_positions": self.get_joints().astype(np.float32)
        }
    
    def setup_mj(self, model, data):
        super().setup_mj(model, data)
        if self.joint_names:
            self.qpos_ids = np.array([self.model.joint(name).qposadr[0] for name in self.joint_names])
            self.act_ids = np.array([self.model.actuator(name).id for name in self.joint_names])
        if hasattr(self, 'gripper_actuator_name'):
            self.gripper_act_id = self.model.actuator(self.gripper_actuator_name).id
    
    @requires_mj
    def get_joints(self):
        return self.data.qpos[self.qpos_ids].copy()
    
    @requires_mj
    def set_joints(self, joint_positions):
        self.data.ctrl[self.act_ids] = joint_positions
    
    @requires_mj
    def step(self, action, step_fn):

        if len(action.shape) == 2:
            action = action[0]
        
        num_joints = len(self.joint_names)
        desired_joints = action[:num_joints]
        grip_action = action[num_joints] if len(action) > num_joints else 0.0
        
        if self.action_space == "delta":
            target_joints = self.get_joints() + desired_joints
        else:
            target_joints = desired_joints
        

        self.set_joints(target_joints)
        self.data.ctrl[self.gripper_act_id] = 255.0 * np.clip(grip_action, 0.0, 1.0)
        
        step_fn(self.control_steps)
