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

from egogym.utils import get_pose, include_in_scene, add_xml_to_scene


class Object:
    def __init__(self, name, path=None, func=None, np_random=None):
        self.name = name
        self.english_name = self.name.rsplit("_", 1)[0].replace("_", " ")
        self.path = path
        self.last_set_pose = None
        self.object_xml = None
        self.handle_type = None
        self.object_geoms_ids = None
        self.np_random = np_random
        if func is not None:
            self.object_xml = func(
                np_random=np_random
            )
        self.body_name = f"{self.name}_object"

    def get_pose(self, data):
        return get_pose(data, self.body_name, "body")

    def get_geom_ids(self, model):
        if self.object_geoms_ids is not None:
            return self.object_geoms_ids
        object_geoms_ids = []
        for i in range(model.ngeom):
            geom_bodyid = model.geom_bodyid[i]
            body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, geom_bodyid)
            if body_name is not None and f"{self.name}_object" in body_name:
                object_geoms_ids.append(i)
        self.object_geoms_ids = object_geoms_ids
        return object_geoms_ids

    def get_bottom_pose(self, data):
        return get_pose(data, f"{self.name}_bottom_site", "site")

    def get_center_to_bottom_z_distance(self, data):
        bottom_z = self.get_bottom_pose(data)[2, 3]
        center_z = self.get_pose(data)[2, 3]
        return center_z - bottom_z

    def set_pose(self, model, data, pose):
        body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, self.body_name)
        model.body_pos[body_id] = pose[0:3, 3]
        model.body_quat[body_id] = R.from_matrix(pose[0:3, 0:3]).as_quat()
        self.last_set_pose = pose.copy()

    def add_to_scene_xml(self, scene_xml):
        if self.object_xml is not None:
            return add_xml_to_scene(scene_xml, self.object_xml)
        elif self.path is not None:
            return include_in_scene(scene_xml, self.path)
        else:
            raise ValueError("Either func or path must be provided")


class ArticulableObject(Object):
    def __init__(self, name, path=None, func=None, np_random=None):
        super().__init__(name, path, func, np_random)
        self.joint_range = None
        self.initial_joint = None
    
    def get_handle_pose(self, data):
        return get_pose(data, "handle", "geom")
    
    def get_jotinpos(self, data):
        return data.joint("task_joint").qpos[0]
    
    def get_perecenttage_opened(self, model, data):
        if self.joint_range is None:
            self.joint_range = model.jnt_range[mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "task_joint")]
            self.initial_joint = self.get_jotinpos(data)
        current_task_joint = self.get_jotinpos(data)
        normalized_task_joint = abs(current_task_joint - self.initial_joint) / (
            abs(self.joint_range[1] - self.joint_range[0]) + 1e-8
        )
        return normalized_task_joint
    
    def open(self, model, data):
        if self.joint_range is None:
            self.joint_range = model.jnt_range[mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "task_joint")]
            self.initial_joint = self.get_jotinpos(data)
        data.joint("task_joint").qpos = (
            self.joint_range[1] * self.np_random.uniform(0.5, 0.9)
            if self.joint_range[0] == 0
            else self.joint_range[0] * self.np_random.uniform(0.5, 0.9)
        )
        mujoco.mj_step(model, data, nstep=100)
        