import os
import cv2
import mujoco
from gymnasium.spaces import Box
import numpy as np
from scipy.spatial.transform import Rotation as R

from egogym.egogym import Egogym
from egogym.utils import include_in_scene, position_sampler, make_objects_manager
import egogym.assets.constants as constants

class OpenTask(Egogym):

    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), num_objs=1, seed=None, objects_set=None):
        super().__init__(robot=robot, action_space=action_space, render_mode=render_mode, render_size=render_size, seed=seed)
        self.num_objs = num_objs
        if objects_set is not None:
            self.objects_manager = make_objects_manager(objects_set, self.np_random, shuffle=False)
        else:
            self.objects_manager = make_objects_manager(constants.all_open_objects_set, self.np_random)
        self.observation_space["handle_pose"] = Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32)

    def make_task_scene(self):
        with open(f"{os.path.dirname(__file__)}/../assets/scenes/open.xml", "r") as f:
            scene_xml = f.read()

        self.object = self.objects_manager.sample()
        scene_xml = self.object.add_to_scene_xml(scene_xml)
        scene_xml = include_in_scene(scene_xml, f"{os.path.dirname(__file__)}/../assets/embodiments/{self.robot.name}/model_open.xml")
        return scene_xml

    def get_obs(self):
        obs = self.get_base_obs()
        obs["handle_pose"] = self.object.get_handle_pose(self.data).reshape(16).astype(np.float32)
        return obs

    def compute_reward(self) -> float:
        perecentage_opened = self.object.get_perecenttage_opened(self.model, self.data)
        return perecentage_opened

    def reset(self, seed=None, options=None):
        super().reset(seed=seed, options=options)
        scene_xml_string = self.make_task_scene()
        self.grasped_bodies = set()
        self.grasping_object = False
        
        gripper_init_pose = np.eye(4)
        gripper_init_pose[:3, :3] = R.from_euler("xyz", np.array([1.5, 0.000, 0.000])).as_matrix()
        self.setup_mujoco(scene_xml_string)
        mujoco.mj_step(self.model, self.data)
        handle_pose = self.object.get_handle_pose(self.data)
        camera_pos = handle_pose[0:3, 3].copy()
        camera_pos += np.array([0.0, -1.0, 0.28])
        self.model.cam_pos[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "exocentric")] = camera_pos
        mujoco.mj_step(self.model, self.data)
        gripper_init_pose[:3, 3] = handle_pose[:3, 3] + np.array([0.0, -0.6, 0.05])
        self.robot.prepare(gripper_init_pose, self.env_step)
        self.env_step()
        self.initial_robot_pose = self.robot.get_camera_pose()

        self.settle_env()
        self.enable_sleeping_islands()

        observation = self.get_obs()
        info = {"object_name": self.object.name}

        return observation, info

    def step(self, action):
        obs, reward, terminated, truncated, info = super().step(action)
        self.grasped_bodies.update(self.robot.get_grasped_bodies())
        grasped_bodies_list = list(self.grasped_bodies)
        if f"{self.object.name}_object" in grasped_bodies_list:
            self.grasping_object = True
        info = {
            "grasped_bodies": grasped_bodies_list,
            "object_name": self.object.name,
            "initial_robot_pose": self.initial_robot_pose,
            "initial_object_pose": self.object.last_set_pose,
            "is_grasping": self.robot.get_grasp(),
            "gripper_current_position": self.robot.get_tcp_pose(),
            "grasping_object": self.grasping_object,
            **info
        }
        return obs, reward, terminated, truncated, info
