# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this
# software and associated documentation files (the "Software"), to deal in the Software
# without restriction, including without limitation the rights to use, copy, modify,
# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

class CachedScoreValue:
    def __init__(self, value, score, timestamp):
        self.value = value
        self.score = score
        self.timestamp = timestamp


class HeapWithInverseIndex():
    """
    Heap where elements can be retrieved using a "key". This allows e.g. for
    updating arbitrary elements. It also supports invalidating arbitrary
    elements in a (time) efficient way.

    The key can be customized through the key_function parameter.
    """
    def __init__(self,
                 ini_vector=None,
                 compare_function=lambda a, b: a > b,
                 value_score_function=lambda x: x,
                 key_function=lambda x: x,
                 use_score_caching=False):
        self.cmp = lambda a, b: compare_function(self._get_value_score(a), self._get_value_score(b))
        self.timestep = 0
        self.use_score_caching = use_score_caching
        self.value_score_function = value_score_function
        self.inverse_index = {}
        self.keyf = lambda v: key_function(v.value)
        self.invalid_pos = set()
        self.data = []
        if ini_vector:
            for (i, x) in enumerate(ini_vector):
                value = CachedScoreValue(x, None, None)
                self.data.append(value)
                self.inverse_index[self.keyf(value)] = i
            self.heapify()

    def _get_value_score(self, cached_score_data):
        if not self.use_score_caching:
            return self.value_score_function(cached_score_data.value)
        else:
            if cached_score_data.timestamp != self.timestep:
                cached_score_data.score = self.value_score_function(cached_score_data.value)
                cached_score_data.timestamp = self.timestep
            return cached_score_data.score

    def increase_timestep(self):
        self.timestep += 1

    def __len__(self):
        return len(self.data)

    def get_max(self):
        return self.data[0].value

    def _swap_pos(self, i, j):
        ori_i_is_invalid = i in self.invalid_pos
        ori_j_is_invalid = j in self.invalid_pos
        (self.data[i], self.data[j]) = (self.data[j], self.data[i])

        # Note: we have swapped positions now
        if ori_i_is_invalid:
            self.invalid_pos.add(j)
        else:
            self.invalid_pos.discard(j)
            self.inverse_index[self.keyf(self.data[j])] = j

        if ori_j_is_invalid:
            self.invalid_pos.add(i)
        else:
            self.invalid_pos.discard(i)
            self.inverse_index[self.keyf(self.data[i])] = i

    def insert(self, elem):
        if self.invalid_pos:
            # Take one of the invalidated positions
            target_pos = self.invalid_pos.pop()
            new_elem = CachedScoreValue(elem, None, None)
            self.data[target_pos] = new_elem
            self.inverse_index[self.keyf(new_elem)] = target_pos
            self.sift_down(target_pos)
            self.sift_up(target_pos)
        else:
            new_elem = CachedScoreValue(elem, None, None)
            self.inverse_index[self.keyf(new_elem)] = len(self.data)
            self.data.append(new_elem)
            self.sift_up(len(self.data) - 1)

    def get(self, key):
        pos = self.inverse_index.get(key)
        if pos is not None:
            return self.data[pos].value
        else:
            return None

    def pop_max_cached_value(self):
        if not self.data:
            raise IndexError("Heap is empty")
        if 0 in self.invalid_pos:
            self.push_invalid(0)
        if 0 in self.invalid_pos:
            raise IndexError("No valid elements in heap")

        maximum = self.data[0]
        del self.inverse_index[self.keyf(self.data[0])]
        if len(self.data) > 1:
            self.clean_invalid()
            last_element_pos = len(self.data) - 1
            if last_element_pos in self.invalid_pos:
                # TODO: This shouldn't happen...
                self.invalid_pos.add(0)
                self.invalid_pos.discard(last_element_pos)
            self.data[0] = self.data.pop()
            self.inverse_index[self.keyf(self.data[0])] = 0
            self.sift_down(0)
        else:
            self.data = []
            self.invalid_pos = set()
        #~ print("%.2f%%" % (100 * float(len(self.invalid_pos)) / len(self.data)))
        return maximum

    def pop_max(self):
        return self.pop_max_cached_value().value

    def parent(self, i):
        return int((i - 1) / 2)

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def has_child(self, i):
        return self.left_child(i) < len(self.data)

    def heapify(self):
        start = self.parent(len(self.data) - 1)
        while start >= 0:
            self.sift_down(start)
            start -= 1

    def substitute_down_key(self, key, new_value):
        pos = self.inverse_index[key]
        self.data[pos].value = new_value
        self.data[pos].timestamp = None
        self.sift_down(pos)

    def substitute_up_key(self, key, new_value):
        pos = self.inverse_index[key]
        self.data[pos].value = new_value
        self.data[pos].timestamp = None
        self.sift_up(pos)

    def remove_key(self, key):
        pos = self.inverse_index.get(key)
        if pos is not None:
            self.remove(pos)
            del self.inverse_index[key]

    def invalidate_key(self, key):
        pos = self.inverse_index.get(key)
        if pos is not None:
            self.invalid_pos.add(pos)
            del self.inverse_index[key]

    def push_invalid(self, pos):
        root = pos
        if pos not in self.invalid_pos:
            return
        if self.has_child(root):
            # Push an invalid node towards the bottom of the heap
            left_child = self.left_child(root)
            right_child = self.right_child(root)

            if left_child in self.invalid_pos:
                self.push_invalid(left_child)
            if right_child < len(self.data):
                if right_child in self.invalid_pos:
                    self.push_invalid(right_child)

                if right_child not in self.invalid_pos and left_child not in self.invalid_pos:
                    # Both are valid
                    exchange_with = left_child if self.cmp(self.data[left_child], self.data[right_child]) else right_child
                elif left_child not in self.invalid_pos:
                    # Left is the only valid
                    exchange_with = left_child
                else:
                    # Either both are invalid or right is the only valid
                    exchange_with = right_child
            else:
                exchange_with = left_child

            if exchange_with not in self.invalid_pos:
                self._swap_pos(root, exchange_with)
                self.push_invalid(exchange_with)
                self.clean_invalid()

    def clean_invalid(self):
        removed = 0
        while (len(self.data) - 1) in self.invalid_pos:
            self.invalid_pos.discard(len(self.data) - 1)
            self.data.pop()
            removed += 1

    def push_all_invalid(self):
        for pos in range(len(self.data)-1, -1, -1):
            if pos in self.invalid_pos:
                self.push_invalid(pos)

    def sift_down(self, pos):
        root = pos
        if root in self.invalid_pos:
            return
        while self.has_child(root):
            left_child = self.left_child(root)
            right_child = self.right_child(root)

            if left_child in self.invalid_pos:
                self.push_invalid(left_child)
            if right_child < len(self.data) and right_child in self.invalid_pos:
                self.push_invalid(right_child)

            if left_child in self.invalid_pos:
                if right_child >= len(self.data) or right_child in self.invalid_pos:
                    break
                else:
                    child = right_child
            else:
                if right_child >= len(self.data) or right_child in self.invalid_pos:
                    child = left_child
                else:
                    child = left_child if self.cmp(self.data[left_child], self.data[right_child]) else right_child

            if self.cmp(self.data[child], self.data[root]):
                self._swap_pos(root, child)
                root = child
            else:
                break
        self.clean_invalid()

    def sift_up(self, pos):
        root = pos
        if root in self.invalid_pos:
            return
        while root > 0:
            parent = self.parent(root)
            if parent in self.invalid_pos:
                self._swap_pos(root, parent)
                # Because the root was invalid, we have to consider if the
                # other subtree may have a bigger element
                self.sift_down(parent)
                root = parent
            elif self.cmp(self.data[root], self.data[parent]):
                self._swap_pos(root, parent)
                root = parent
            else:
                break
        self.clean_invalid()

    def remove(self, pos):
        if pos != len(self.data) - 1:
            self._swap_pos(pos, len(self.data) - 1)
            self.invalid_pos.discard(len(self.data) - 1)
            self.data.pop()
            if pos > 0 and self.cmp(self.data[pos], self.data[self.parent(pos)]):
                self.sift_up(pos)
            else:
                self.sift_down(pos)
        else:
            self.invalid_pos.discard(len(self.data) - 1)
            self.data.pop()

    def update_compare_function(self, new_cmp):
        self.cmp = new_cmp
        self.heapify()

    def dfs(self, pos):
        stack = [pos]
        while stack:
            root = stack.pop()
            yield root
            child = self.left_child(root)
            if child < len(self.data):
                stack.append(child)
                child += 1  # Right child
                if child < len(self.data):
                    stack.append(child)

    def check_heap_property(self, root):
        if root in self.invalid_pos:
            return True # Actually we should check the children, but we will call it in the whole heap
        holds = True
        for node in self.dfs(root):
            if node == root:
                continue
            holds = node in self.invalid_pos or self.cmp(self.data[root], self.data[node])
            if not holds:
                print("Does not hold  root(%d) = %d node(%d) = %d" % (root, self.data[root], node, self.data[node]))
                break
        return holds
