import numpy as np


class SumTree(object):

    def __init__(self, size):
        self.size = size
        self.n_entries = 0
        self.memory_index = 0
        self.tree = np.zeros(2 * size - 1)

    def total(self):
        return self.tree[0]

    def add(self, priority):
        index = self.memory_index + (self.size-1)
        self.update(index, priority)
        self.memory_index = (self.memory_index + 1) % self.size
        if self.n_entries < self.size:
            self.n_entries += 1

    def update(self, index, priority):
        change = priority - self.tree[index]
        self.tree[index] = priority

        parent_index = (index - 1) // 2
        while parent_index != 0:
            self.tree[parent_index] += change
            parent_index = (parent_index - 1) // 2
        self.tree[parent_index] += change

    def get_sample_priority(self, sum_until_value):
        index = 0
        left_node_index = 2 * index + 1
        right_node_index = left_node_index + 1

        while left_node_index < len(self.tree):
            if sum_until_value <= self.tree[left_node_index]:
                index = left_node_index
                left_node_index = 2 * index + 1
                right_node_index = left_node_index + 1
            else:
                sum_until_value = sum_until_value - self.tree[left_node_index]
                index = right_node_index
                left_node_index = 2 * index + 1
                right_node_index = left_node_index + 1

        return index - (self.size-1), self.tree[index]
