import gym
import numpy as np


class NoisyObservationWrapper(gym.Wrapper):
    def __init__(self, env, sigma=1.0):
        super().__init__(env)
        self.env = env
        self.sigma = sigma

    def step(self, action):

        obs, reward, done, info = self.env.step(action)
        noise = np.random.randn(*self.observation_space.shape) * self.sigma
        noisy_obs = obs + noise
        len_qpos = self.env.unwrapped.sim.data.qpos.shape[0]
        len_qvel = self.env.unwrapped.sim.data.qvel.shape[0]
        # import pdb; pdb.set_trace()
        self.env.unwrapped.sim.data.qpos[1:] = self.env.unwrapped.sim.data.qpos[1:] + noise[:len_qpos-1]
        self.env.unwrapped.sim.data.qvel[:] = self.env.unwrapped.sim.data.qvel[:] + noise[len_qpos - 1:]



        return noisy_obs, reward, done, info
