import os
from gymnasium.spaces import Box, Text
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, make_textures_manager
import egogym.assets.constants as constants

class PickTask(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.spread = 0.22
        self.num_objs = num_objs
        if objects_set is not None:
            self.objects_manager = make_objects_manager(objects_set, self.np_random, shuffle=True)
        else:
            self.objects_manager = make_objects_manager(constants.lite_pick_objects_set, self.np_random)
        self.textures_manager = make_textures_manager([f"wood/{i}.png" for i in range(10)], self.np_random)
        self.observation_space["object_pose"] = Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32)
        self.observation_space["object_name"] = Text(max_length=256)


    def make_task_scene(self):
        assets_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets"))
        
        with open(os.path.join(assets_dir, "scenes", "pick.xml"), "r") as f:
            scene_xml = f.read()

        self.object = self.objects_manager.sample()
        self.secondary_objects = []
        used_objects = {self.object.name}
        scene_xml = self.object.add_to_scene_xml(scene_xml)
        while len(used_objects) < self.num_objs:
            object = self.objects_manager.sample(random=True)
            if object.name not in used_objects and object.name.split("_")[0] != self.object.name.split("_")[0]:
                self.secondary_objects.append(object)
                used_objects.add(object.name)
                scene_xml = object.add_to_scene_xml(scene_xml)
        self.texture = self.textures_manager.sample(random=True)
        scene_xml = include_in_scene(scene_xml, os.path.join(assets_dir, "embodiments", self.robot.name, "model_pick.xml"))
        scene_xml = scene_xml.replace("{ASSETS_DIR}", assets_dir)
        scene_xml = scene_xml.replace("{TEXTURE_PATH}", self.texture)
        scene_xml = scene_xml.replace("{TABLE_WIDTH}", str(self.spread+0.24))
        scene_xml = scene_xml.replace("{TABLE_HEIGHT}", str(self.spread+0.24))
        return scene_xml

    def get_obs(self):
        obs = self.get_base_obs()
        obs["object_pose"] = self.object.get_pose(self.data).reshape(16).astype(np.float32)
        obs["object_name"] = self.object.name.split("_")[0]
        return obs

    def compute_reward(self) -> float:
        object_pose = self.object.get_pose(self.data)
        lift_distance = object_pose[2][3] - self.object.last_set_pose[2][3]
        return lift_distance

    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] = np.array([0.00, -0.68, 1.1])
        gripper_init_pose[:3, :3] = R.from_euler("xyz", np.array([1.3, 0.000, 0.000])).as_matrix()
        self.setup_mujoco(scene_xml_string)
        initial_positions = position_sampler(self.np_random, len(self.secondary_objects) + 1, [self.spread,self.spread], self.spread/2, 100, start_position=gripper_init_pose[:2, 3], thickness=self.spread/2.5)
        pose = np.eye(4)
        self.initial_robot_pose = self.robot.get_camera_pose()
        self.env_step(10)

        for i, obj in enumerate([self.object] + self.secondary_objects):
            z = obj.get_center_to_bottom_z_distance(self.data) + 0.78
            pose[:2, 3] = initial_positions[i][:2]
            pose[2, 3] = z
            pose[:3, :3] = R.from_euler(
                "x", self.np_random.uniform(0, 2 * np.pi)
            ).as_matrix()
            obj.set_pose(self.model, self.data, pose)
        
        self.env_step(100)
        self.settle_env()
        self.env_step(100)
        for obj in [self.object] + self.secondary_objects:
            obj.last_set_pose = obj.get_pose(self.data).copy()
        
        self.robot.prepare(gripper_init_pose, self.env_step)
        self.env_step(100)        
        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,
            "texture_name": self.texture,
            "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
        }

        if np.linalg.norm(self.robot.get_camera_pose()[:3,3] - self.object.get_pose(self.data)[:3,3]) > np.linalg.norm(self.initial_robot_pose[:3,3] - self.object.last_set_pose[:3,3])*1.5:
            truncated = True
            terminated = True

        return obs, reward, terminated, truncated, info
