import mujoco
from functools import wraps
import numpy as np
from egogym.utils import get_pose


def requires_mj(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if not hasattr(self, 'model') or not hasattr(self, 'data') or self.model is None or self.data is None:
            raise RuntimeError(f"Must call setup_mj before using the {self.__class__.__name__}.")
        return func(self, *args, **kwargs)
    return wrapper


class Robot():    
    def __init__(self, control_steps=50, grasping_steps=1, action_space="delta", np_random=None):
        self.name = None
        self.model = None
        self.data = None
        self.control_steps = control_steps
        self.grasping_steps = grasping_steps
        self.action_space = action_space
        self.np_random = np_random
        self.camera_names = ["egocentric", "exocentric"]
        self.rgb_renderers = {}
        self.add_noise_init = True

    def additional_obs_spaces(self):
        return {}
    
    def get_additional_obs(self):
        return {}
    
    def action_space_shape(self):
        return (0,)
    
    def setup_mj(self, model, data):
        self.model = model
        self.data = data

    @requires_mj
    def get_camera_pose(self):
        return get_pose(self.data, "egocentric", "camera")
    
    
    @requires_mj
    def setup_renderers(self, render_width=960, render_height=720):
        if len(self.rgb_renderers) > 0:
            del self.rgb_renderers
        self.rgb_renderers = {}
        for camera_name in self.camera_names:
            self.rgb_renderers[camera_name] = mujoco.Renderer(self.model, height=render_height, width=render_width)
    
    @requires_mj
    def get_camera_view(self, camera_name):
        scene_opt = mujoco.MjvOption()
        self.rgb_renderers[camera_name].update_scene(self.data, camera_name, scene_opt)
        rgb = self.rgb_renderers[camera_name].render()
        return rgb
    
    @requires_mj
    def get_tcp_pose(self):
        return get_pose(self.data, "grasping_center", "site")
    
    @requires_mj
    def get_grasp(self):
        return np.array(self.data.qpos[self.model.actuator("fingers_actuator").id])
    
    @requires_mj
    def get_grasped_bodies(self):
        contact_points = self.data.contact
        body_names = set()
        for contact in contact_points:
            body1_id = self.model.geom_bodyid[contact.geom1]
            body1 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body1_id)
            body2_id = self.model.geom_bodyid[contact.geom2]
            body2 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body2_id)
            if body1 is None or body2 is None:
                continue
            if "left" in body1 or "right" in body1:
                body_names.add(body2)
            if "left" in body2 or "right" in body2:
                body_names.add(body1)
        return body_names
    

    @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