from gymnasium.utils import seeding
import gymnasium as gym
from gymnasium.spaces import Box, Dict
import mujoco
import numpy as np
import cv2

from egogym.embodiments import ROBOT_REGISTRY

class Egogym(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30}

    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), seed=None):
        super().__init__()
        self.model = None
        self.data = None
        self.step_idx = 0
        self.render_mode = render_mode
        self.np_random, _ = seeding.np_random(seed)
    
        self.render_width = render_size[0]
        self.render_height = render_size[1]
        
        self.robot_type = robot
        self.action_space_type = action_space
        
        if self.robot_type not in ROBOT_REGISTRY:
            raise ValueError(f"Unknown robot type '{self.robot_type}'. Available robots: {list(ROBOT_REGISTRY.keys())}")
        
        robot_class = ROBOT_REGISTRY[self.robot_type]
        self.robot = robot_class(action_space=self.action_space_type, np_random=self.np_random)
        
        obs_spaces = {
            "rgb_ego": Box(low=0, high=255, shape=(self.render_height, self.render_width, 3), dtype=np.uint8),
            "rgb_exo": Box(low=0, high=255, shape=(self.render_height, self.render_width, 3), dtype=np.uint8),
            "tcp_pose": Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32),
            "camera_pose": Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32),
            "grasp": Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32),
        }
        obs_spaces.update(self.robot.additional_obs_spaces())
        self.observation_space = Dict(obs_spaces)
        self.action_space = Box(low=-np.inf, high=np.inf, shape=self.robot.action_space_shape(), dtype=np.float32)
        self._mujoco_step_counter = 0
        self._render_freq = 0
        self._render_callback = None
        self._should_render = self.render_mode == 'human'

    def setup_mujoco(self, scene_xml: str):
        if hasattr(self, "model") and self.model is not None:
            del self.model
            del self.data
        self.model = mujoco.MjModel.from_xml_string(scene_xml)
        self.data = mujoco.MjData(self.model)
        self.robot.setup_mj(self.model, self.data)
        self.robot.setup_renderers(self.render_width, self.render_height)

    def env_step(self, nsteps=1):
        for _ in range(nsteps):
            mujoco.mj_step(self.model, self.data)
            self._mujoco_step_counter += 1
            if self._render_freq > 0 and self._render_callback is not None and self._mujoco_step_counter >= self._render_freq:
                self._render_callback()
                self._mujoco_step_counter = 0

    def get_base_obs(self):
        obs = {
            "rgb_ego": self.robot.get_camera_view(self.robot.camera_names[0]),
            "rgb_exo": self.robot.get_camera_view(self.robot.camera_names[-1]),
            "tcp_pose": self.robot.get_tcp_pose().reshape(16).astype(np.float32),
            "camera_pose": self.robot.get_camera_pose().reshape(16).astype(np.float32),
            "grasp": np.array([self.robot.get_grasp()], dtype=np.float32),
        }
        obs.update(self.robot.get_additional_obs())
        return obs

    def render(self):
        rgb_ego = self.robot.get_camera_view(self.robot.camera_names[0])
        rgb_exo = self.robot.get_camera_view(self.robot.camera_names[1])
        rgb_stack = np.concatenate((rgb_ego, rgb_exo), axis=1)
        rgb_stack = cv2.cvtColor(rgb_stack, cv2.COLOR_RGB2BGR)
        cv2.imshow("Egogym", rgb_stack)
        cv2.waitKey(1)

    def close(self):
        if self.render_mode == "human":
            cv2.destroyAllWindows()

    def step(self, action):
        if action is None or np.isnan(action).any():
            return self.get_obs(), 0.0, False, False, {}
        self.robot.step(action, step_fn=self.env_step)
        obs = self.get_obs()
        reward = self.compute_reward()
        terminated = False
        truncated = False
        return obs, reward, terminated, truncated, {}
    
    def enable_sleeping_islands(self):
        try:
            self.model.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_SLEEP
        except AttributeError:
            pass
    
    def settle_env(self, nsteps=200):
        if hasattr(self.robot, 'free_joint_id') and self.robot.free_joint_id is not None: #TODO: check if its floating_gripper instead
            self.data.qpos[:] = np.where((np.arange(len(self.data.qpos)) >= self.robot.free_joint_id) & (np.arange(len(self.data.qpos)) < self.robot.free_joint_id + 6), self.data.qpos, 0.0)