import numpy as np
from rand_param_envs.base import RandomEnv
from rand_param_envs.gym import utils

class CheetahVelRandParamsEnv(RandomEnv, utils.EzPickle):
    def __init__(self, log_scale_limit=3.0):
        self._goal_vel = 0.
        self.goal_radius = 0.3
        RandomEnv.__init__(self, log_scale_limit, 'half_cheetah.xml', 5)
        utils.EzPickle.__init__(self)

    def _step(self, action):
        xposbefore = self.data.qpos[0][0]
        self.do_simulation(action, self.frame_skip)
        xposafter = self.data.qpos[0][0]

        forward_vel = (xposafter - xposbefore) / self.dt
        forward_reward = -1.0 * abs(forward_vel - self._goal_vel)
        sparse_reward = self.sparsify_rewards(forward_reward)
        ctrl_cost = 0.5 * 1e-1 * np.sum(np.square(action))

        observation = self._get_obs()
        reward = forward_reward - ctrl_cost
        sparse_reward = sparse_reward - ctrl_cost
        reward = sparse_reward
        done = False
        infos = dict(reward_forward=forward_reward,
            reward_ctrl=-ctrl_cost, velocity=forward_vel,sparse_reward=sparse_reward)
        return (observation, reward, done, infos)

    def _get_obs(self):
        return np.concatenate([
            self.data.qpos.flat[1:],
            self.data.qvel.flat,
            self.get_body_com("torso").flat,
        ]).astype(np.float32).flatten()

    def viewer_setup(self):
        camera_id = self.model.camera_name2id('track')
        self.viewer.cam.type = 2
        self.viewer.cam.fixedcamid = camera_id
        self.viewer.cam.distance = self.model.stat.extent * 0.35
        # Hide the overlay
        self.viewer._hide_overlay = True

    def sparsify_rewards(self, r):
        ''' zero out rewards when outside the goal radius '''
        #mask = (r >= -self.goal_radius).astype(np.float32)
        #r = r * mask
        if r < - self.goal_radius:
            r = -2.
        r = r + 2.
        return r

'''
if __name__ == "__main__":

    env = Walker2DRandParamsEnv()
    tasks = env.sample_tasks(40)
    while True:
        env.reset()
        env.set_task(np.random.choice(tasks))
        print(env.model.body_mass)
        for _ in range(100):
            env.render()
            env.step(env.action_space.sample())  # take a random action
'''

