# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import print_function
import numpy as np


class ReplayMem(object):

    def __init__(
        self,
        obs_dim,
        act_dim,
        memory_size=1000000):

        # allocate space for memory cells
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.obss = np.zeros((memory_size, obs_dim))
        self.acts = np.zeros((memory_size, act_dim))
        self.rwds = np.zeros((memory_size, ))
        self.ends = np.zeros(memory_size, dtype='uint8')
        self.memory_size = memory_size
        self.top = -1
        self.size = 0

    def add_sample(self, obs, act, rwd, end):

        self.top = (self.top + 1) % self.memory_size
        self.obss[self.top] = obs
        self.acts[self.top] = act
        self.rwds[self.top] = rwd
        self.ends[self.top] = end

        if self.size < self.memory_size:
            self.size += 1

    def get_batch(self, batch_size):

        assert self.size >= batch_size

        indices = np.zeros(batch_size, dtype="uint64")
        transit_indices = np.zeros(batch_size, dtype="uint64")

        counter = 0
        while counter < batch_size:
            idx = np.random.randint(0, self.size)
            # case where the last piece of memory is sampled
            # which does not have a successor state
            if idx == self.top:
                    continue
            transit_idx = (idx + 1) % self.memory_size
            indices[counter] = idx
            transit_indices[counter] = transit_idx
            counter += 1

        return (self.obss[indices],
                self.acts[indices],
                self.rwds[indices],
                self.ends[indices],
                self.obss[transit_indices])


if __name__ == "__main__":

    memory = ReplayMem(2, 1, memory_size=10)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([2, 2]), np.array([2]), 10, 0)
    memory.add_sample(np.array([1, 1]), np.array([1]), 100, 1)
    print(memory.obss)
    print(memory.acts)
    print(memory.rwds)
    print(memory.ends)
    print(memory.get_batch(5))

