import os
import copy
import numpy as np

import gym
from gym import error, spaces
from gym.utils import seeding
import cv2
import time

try:
    import mujoco_py
except ImportError as e:
    raise error.DependencyNotInstalled("{}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)".format(e))


class RobotEnv(gym.GoalEnv):
    def __init__(self, model_path, initial_qpos, n_actions, n_substeps):
        if model_path.startswith('/'):
            fullpath = model_path
        else:
            fullpath = os.path.join(os.path.dirname(__file__), 'assets', model_path)
        if not os.path.exists(fullpath):
            raise IOError('File {} does not exist'.format(fullpath))

        model = mujoco_py.load_model_from_path(fullpath)
        self.sim = mujoco_py.MjSim(model, nsubsteps=n_substeps)
        self.viewer = None

        self.max_u = [0.25, 0.27, 0.145]
        self.action_offset = [1.3, 0.75, 0.555]

        self.last_action = None

        self.metadata = {
            'render.modes': ['human', 'rgb_array'],
            'video.frames_per_second': int(np.round(1.0 / self.dt))
        }

        self.seed()
        self._env_setup(initial_qpos=initial_qpos)
        self.initial_state = copy.deepcopy(self.sim.get_state())

        # self._reset_sim()
        self.goal = self._sample_goal()
        self.flag = 1
        obs = self._get_obs()
        self.action_space = spaces.Box(-1., 1., shape=(n_actions,), dtype='float32')
        self.observation_space = spaces.Dict(dict(
            desired_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            achieved_goal=spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype='float32'),
            observation=spaces.Box(-np.inf, np.inf, shape=obs['observation'].shape, dtype='float32'),
        ))

        self.render_permit = 0

    @property
    def dt(self):
        return self.sim.model.opt.timestep * self.sim.nsubsteps

    # Env methods
    # ----------------------------

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step_maze(self, action):
        action = np.clip(action, self.action_space.low, self.action_space.high)
        self._set_action(action)
        self.sim.step()
        self._step_callback()
        obs = self._get_obs()

        done = False
        info = {
            'is_success': self._is_success(obs['achieved_goal'], self.goal),
        }
        reward = self.compute_reward(obs['achieved_goal'], self.goal, info)
        # print(reward)
        return obs, reward, done, info

    def step_pick(self, action):
        action = np.clip(action, self.action_space.low, self.action_space.high)
        self._set_action(action)
        self.sim.step()
        self._step_callback()
        obs = self._get_obs()

        done = False
        info = {
            'is_success': self._is_success(obs['achieved_goal'], self.goal),
        }
        reward = self.compute_reward(obs['achieved_goal'], self.goal, info)
        return obs, reward, done, info

    def step_rope(self, action):
        x_min = 1.05
        x_max = 1.55
        y_min = 0.48
        y_max = 1.02
        self.collision_with_initial_rope = 0
        action = np.clip(action, self.action_space.low, self.action_space.high)
        drag_len = 0.08
        action = action.copy()  # ensure that we don't change the action outside of this scope
        ac_start = np.array([action[0], action[1], 0.5])
        # ac_end = np.array([action[2], action[3], 0.5])
        drag_theta = action[2]*180+180.0
        # Bring in 0 to 1 range
        action[3] = (action[3]+1)/2.0
        # drag_len = action[3]*drag_len
        bins = 18
        # drag_theta = drag_theta - drag_theta % 20.0
        action_start = self.action_offset + (self.max_u * ac_start)
        # action_end = self.action_offset + (self.max_u * ac_end)

        action_x = action_start[0] + drag_len*np.cos(np.radians(drag_theta))*1.0
        action_y = action_start[1] + drag_len*np.sin(np.radians(drag_theta))*1.0
        action_x = max(x_min, action_x)
        action_x = min(x_max, action_x)
        action_y = max(y_min, action_y)
        action_y = min(y_max, action_y)
        action_end = np.array([action_x, action_y, 0.5])
        action_start[2] = 0.5
        action_end[2] = 0.5
        action_start_temp = action_start.copy()
        action_end_temp = action_end.copy()
        action_start_temp[2] = 0.42
        action_end_temp[2] = 0.42

        grip_pos = self.sim.data.get_site_xpos('robot0:grip').copy()
        object_oriented_goal = action_start - grip_pos
        initial_grip_pos = self.sim.data.get_site_xpos('robot0:grip').copy()
        initial_action_start = initial_grip_pos.copy()
        initial_action_start[2] = 0.6
        initial_object_oriented_goal = initial_action_start - initial_grip_pos
        
        image_count = 0
        count = 0

        while np.linalg.norm(initial_object_oriented_goal) >= 0.03:
            if count > 20:
                break
            else:
                count += 1
            
            image_count += 1
            action = np.array([0., 0., 0., 0.])
            action[0] = initial_object_oriented_goal[0]*6
            action[1] = initial_object_oriented_goal[1]*6
            action[2] = initial_object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            initial_grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            initial_object_oriented_goal = initial_action_start - initial_grip_pos

        grip_pos = self.sim.data.get_site_xpos('robot0:grip').copy()
        object_oriented_goal = action_start - grip_pos

        while np.linalg.norm(object_oriented_goal) >= 0.03:
            if count > 20:
                break
            else:
                count += 1
            if self.render_permit:
                self.render()
            image_count += 1
            
            action = np.array([0., 0., 0., 0.])
            action[0] = object_oriented_goal[0]*6
            action[1] = object_oriented_goal[1]*6
            action[2] = object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            object_oriented_goal = action_start - grip_pos

        rope_size = 15
        grip_pos = self.sim.data.get_site_xpos('robot0:grip')
        object_oriented_goal = action_start - grip_pos
        
        image_count = 0
        count = 0
        while np.linalg.norm(object_oriented_goal) >= 0.03:
            if count > 20:
                break
            else:
                count += 1
            if self.render_permit:
                self.render()
            image_count += 1
            action = np.array([0., 0., 0., 0.])
            action[0] = object_oriented_goal[0]*6
            action[1] = object_oriented_goal[1]*6
            action[2] = object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            object_oriented_goal = action_start - grip_pos


        action_start[2] = 0.25
        grip_pos = self.sim.data.get_site_xpos('robot0:grip')
        object_oriented_goal = action_start - grip_pos
        count = 0
        while np.linalg.norm(object_oriented_goal) >= 0.03:
            if count > 20:
                break
            else:
                count += 1
            if self.render_permit:
                self.render()
            if self.if_collision():
                self.collision_with_initial_rope = 1
            image_count += 1
            action = np.array([0., 0., 0., 0.])
            action[0] = object_oriented_goal[0]*6
            action[1] = object_oriented_goal[1]*6
            action[2] = object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            object_oriented_goal = action_start - grip_pos

        action_start[:2] = action_end[:2].copy()
        grip_pos = self.sim.data.get_site_xpos('robot0:grip')
        object_oriented_goal = action_start - grip_pos

        count = 0
        while np.linalg.norm(object_oriented_goal) >= 0.04:
            if count > 20:
                break
            else:
                count += 1
            if self.render_permit:
                self.render()
            if self.if_collision():
                self.collision_with_initial_rope = 1
            image_count += 1
            action = np.array([0., 0., 0., 0.])
            action[0] = object_oriented_goal[0]*6
            action[1] = object_oriented_goal[1]*6
            action[2] = object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            object_oriented_goal = action_start - grip_pos

        action_start[2] = 0.5
        grip_pos = self.sim.data.get_site_xpos('robot0:grip')
        object_oriented_goal = action_start - grip_pos

        count = 0
        while np.linalg.norm(object_oriented_goal) >= 0.01:
            if count > 20:
                break
            else:
                count += 1
            if self.render_permit:
                self.render()
            if self.if_collision():
                self.collision_with_initial_rope = 1
            image_count += 1
            action = np.array([0., 0., 0., 0.])
            action[0] = object_oriented_goal[0]*6
            action[1] = object_oriented_goal[1]*6
            action[2] = object_oriented_goal[2]*6
            action[3] = 0

            action_temp = action.copy()
            self._set_action(action_temp)
            self.sim.step()
            self._step_callback()

            grip_pos = self.sim.data.get_site_xpos('robot0:grip')
            object_oriented_goal = action_start - grip_pos

        obs = self._get_obs()
        done = False
        info = {
            'is_success': self._is_success(obs['achieved_goal'], self.goal),
        }
        reward = self.compute_reward(obs['achieved_goal'], self.goal, info)
        return obs, reward, done, info

    def reset(self):
        # Attempt to reset the simulator. Since we randomize initial conditions, it
        # is possible to get into a state with numerical issues (e.g. due to penetration or
        # Gimbel lock) or we may not achieve an initial condition (e.g. an object is within the hand).
        # In this case, we just keep randomizing until we eventually achieve a valid initial
        # configuration.
        did_reset_sim = False
        while not did_reset_sim:
            did_reset_sim = self._reset_sim()
        self.goal = self._sample_goal().copy()
        obs = self._get_obs()
        return obs

    def close(self):
        if self.viewer is not None:
            # self.viewer.finish()
            self.viewer = None

    def render(self, mode='human'):
        self._render_callback()
        if mode == 'rgb_array':
            self._get_viewer().render()
            # window size used for old mujoco-py:
            width, height = 500, 500
            data = self._get_viewer().read_pixels(width, height, depth=False)
            # original image is upside-down, so flip it
            return data[::-1, :, :]
        elif mode == 'human':
            self._get_viewer().render()

    def _get_viewer(self):
        if self.viewer is None:
            self.viewer = mujoco_py.MjViewer(self.sim)
            self._viewer_setup()
        return self.viewer


    def set_subgoal(self, name, action):
        # site_id = self.sim.model.site_name2id(name)
        # self.sim.model.site_pos[site_id] = action
        raise NotImplementedError()


    # Extension methods
    # ----------------------------

    def _reset_sim(self):
        """Resets a simulation and indicates whether or not it was successful.
        If a reset was unsuccessful (e.g. if a randomized state caused an error in the
        simulation), this method should indicate such a failure by returning False.
        In such a case, this method will be called again to attempt a the reset again.
        """
        self.sim.set_state(self.initial_state)
        self.sim.forward()
        return True

    def _get_obs(self):
        """Returns the observation.
        """
        raise NotImplementedError()

    def _set_action(self, action):
        """Applies the given action to the simulation.
        """
        raise NotImplementedError()

    def _is_success(self, achieved_goal, desired_goal):
        """Indicates whether or not the achieved goal successfully achieved the desired goal.
        """
        raise NotImplementedError()

    def _sample_goal(self):
        """Samples a new goal and returns it.
        """
        raise NotImplementedError()

    def if_collision(self):
        """Collision detection.
        """
        pass

    def _env_setup(self, initial_qpos):
        """Initial configuration of the environment. Can be used to configure initial state
        and extract information from the simulation.
        """
        pass

    def _viewer_setup(self):
        """Initial configuration of the viewer. Can be used to set the camera position,
        for example.
        """
        pass

    def _render_callback(self):
        """A custom callback that is called before rendering. Can be used
        to implement custom visualizations.
        """
        pass

    def _step_callback(self):
        """A custom callback that is called after stepping the simulation. Can be used
        to enforce additional constraints on the simulation state.
        """
        pass
