import pickle

import numpy as np
from agent_code.big_bertha_v1.parameters import (BATCH_SIZE,
                                                 EXPERIENCE_BUFFER_SIZE_MAX,
                                                 STATE_SHAPE)


class ExperienceBuffer(object):
    def __init__(self):
        self.size = EXPERIENCE_BUFFER_SIZE_MAX
        self.memory_index = 0
        self.n_entries = 0

        self.states = np.zeros((self.size, STATE_SHAPE))
        self.actions = np.zeros(self.size, dtype=np.int8)
        self.rewards = np.zeros(self.size)
        self.next_states = np.zeros((self.size, STATE_SHAPE))

    def remember(self, state, action, reward, next_state):
        self.states[self.memory_index] = state
        self.actions[self.memory_index] = action
        self.rewards[self.memory_index] = reward
        self.next_states[self.memory_index] = next_state
        self.memory_index = (self.memory_index + 1) % self.size
        if self.n_entries < self.size:
            self.n_entries += 1

    def sample(self, batch_size=BATCH_SIZE):
        if self.n_entries == self.size:
            indices = np.random.choice(self.size, batch_size)
        else:
            indices = np.random.choice(self.memory_index + 1, batch_size)

        return self.states[indices], self.actions[indices], self.rewards[indices], self.next_states[indices]

    def save_samples(self, file_name='expert_transitions'):
        with open('{}.p'.format(file_name), 'wb') as handle:
            pickle.dump([self.states, self.actions, self.rewards, self.next_states, self.memory_index, self.n_entries], handle,
                        protocol=pickle.HIGHEST_PROTOCOL)

    def load_samples(self, file_name='expert_transitions'):
        temp_trans = pickle.load(open("{}.p".format(file_name), "rb"))
        self.states = temp_trans[0]
        self.actions = temp_trans[1]
        self.rewards = temp_trans[2]
        self.next_states = temp_trans[3]
        self.memory_index = temp_trans[4]
        self.n_entries = temp_trans[5]
