from gym.envs.mujoco.humanoid import HumanoidEnv, mujoco_env, utils
import numpy as np
from gym import register
import os
from pathlib import Path


class WallHumanoidEnv(HumanoidEnv):
    def __init__(self):
        mujoco_env.MujocoEnv.__init__(self, str(Path('./safe/envs/assets/humanoid.xml').resolve()), 5)
        utils.EzPickle.__init__(self)

    def _get_obs(self):
        data = self.sim.data
        return np.concatenate([data.qpos.flat[2:],
                               data.qvel.flat,
                               data.cinert[:-1].flat,  # remove the wall
                               data.cvel[:-1].flat,  # remove the wall
                               data.qfrc_actuator.flat,
                               data.cfrc_ext[:-1].flat])  # remove the wall

    def viewer_setup(self):
        # self.viewer.cam.type = mujoco_py.const.CAMERA_FIXED
        # self.viewer.cam.fixedcamid = 0
        self.viewer.cam.lookat[:] = [-0.5, 0, 1]
        self.viewer.cam.elevation = -1.4
        self.viewer.cam.distance = 10
        self.viewer.cam.azimuth = 70


register('WallHumanoid-v2', entry_point=WallHumanoidEnv, max_episode_steps=1000)
