import numpy as np

class ReplayBuffer:
    def __init__( self, obs_dim, act_dim, size):
        self.obs1_buf = np.zeros( [size, obs_dim], dtype=np.float32)
        self.obs2_buf = np.zeros( [size, obs_dim], dtype=np.float32)
        self.acts_buf = np.zeros( [size, act_dim], dtype=np.float32)
        self.rews_buf = np.zeros( size, dtype=np.float32)
        self.done_buf = np.zeros( size, dtype=np.float32)
        self.ptr, self.size, self.max_size = 0,0, size

    def store( self, obs1, act, rew, obs2, done):
        self.obs1_buf[self.ptr] = obs1
        self.obs2_buf[self.ptr] = obs2
        self.acts_buf[self.ptr] = act
        self.rews_buf[self.ptr] = rew
        self.done_buf[self.ptr] = done
        self.ptr = (self.ptr+1) % self.max_size
        self.size = min( self.size+1, self.max_size)

    def sample_batch( self, batch_size=32):
        idxs = np.random.randint( 0, self.size, size=batch_size)

        return dict( obs1 = self.obs1_buf[idxs],
                     obs2 = self.obs2_buf[idxs],
                     rews = self.rews_buf[idxs],
                     acts = self.acts_buf[idxs],
                     done = self.done_buf[idxs])
