#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed 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.

"""This file is a python-implementation for a heap.

This file allows for both a min-heap and a max-heap by making
use of the :meth:`MIN_HEAP` and :meth:`MAX_HEAP` comparator
functions listed below.
"""

MIN_HEAP = lambda x, l: x < l
"""Comparator function, indicating to use a min-heap as the heap.

See Also:
    :class:`heap`
"""

MAX_HEAP = lambda x, l: x > l
"""Comparator function, indicating to use a max-heap as the heap.

See Also:
    :class:`heap`
"""

class heap:
    """The heap class is a simple heap implementation.

    Note:
        The comparator function takes two arguments (``x`` and ``l``, as
        described below) and should return ``True`` if ``x`` should be
        swapped with ``l``.

        - ``x``: The value of a node in the heap.
        - ``l``: The root to compare ``x`` to.

    Warning:
        You shouldn't need any other than :meth:`MIN_HEAP` or :meth:`MAX_HEAP`
        for the comparison function. If you do, you're probably doing strange
        things and you should consider using another datastructure.

    Args:
        comparator (def):   The comparison function. Defaults to :meth:`MAX_HEAP`.
    """
    def __init__(self, comparator = MAX_HEAP):
        self.data = []
        self.comparator = comparator

    def __repr__(self):
        return repr(self.data)

    def __len__(self):
        """Gets the heap size."""
        return len(self.data)

    def items(self):
        """Get a copy of the contents of the heap as a list."""
        return self.data[:]

    def empty(self):
        """Check if the heap is empty.

        Returns:
            ``True`` if the heap is empty, ``False`` otherwise.
        """
        return len(self.data) == 0

    def _heapify(self, i):
        """Fix the heap, making sure the heap condition remains satisfied.

        Warning:
            This function should not be called from an external source.

        Args:
            i (int):    Index of the root node to heapify.
        """
        largest = i
        l = 2 * i + 1   # left
        r = 2 * i + 2   # right
        n = len(self.data)

        if l < n and self.comparator(self.data[l], self.data[largest]):
            largest = l

        if r < n and self.comparator(self.data[r], self.data[largest]):
            largest = r

        if largest != i:
            # Swap
            self.data[i], self.data[largest] = self.data[largest], self.data[i]

            self._heapify(largest)
            d = (i - 1) // 2
            if d >= 0:
                self._heapify(d)

    def insert(self, node):
        """Insert a new node into the heap.

        Args:
            node (object):  A comparable object to insert.
        """
        self.data.append(node)
        self._heapify((len(self.data) - 2) // 2)

    def pop(self, i=0):
        """Remove the ith item of the heap.

        Args:
            i (int):    The index of the node to pop.
                        Defaults to 0 (the top).

        Returns:
            The item that was removed.
        """
        top = self.peek(i)
        self.data[i] = self.data[-1]
        del self.data[-1]
        self._heapify(i)
        return top

    def indexOf(self, node):
        """Gets the index of a node.

        Args:
            node (object):  The node to get the index for.

        Returns:
            The index of that node.
        """
        return self.data.index(node)

    def peek(self, i=0):
        """Look at the item at location i in the heap.

        Args:
            i (int):    The index of the node to peek.
                        Defaults to 0 (the top).

        Returns:
            ``None`` if the heap is empty, otherwise the item at index ``i``.

        Raises:
           IndexError: When the index is not found.
        """
        return None if self.empty() else self.data[i]
