"""
 heap.py
 
 a miminal implementation of a priority queue
 as a binary heap including heapsort
 
 See https://runestone.academy/runestone/books/published/pythonds/
     Trees/BinaryHeapImplementation.html
 
 ---------------------------------------------------

    $ python heap.py
    === original ===
     5
        9
           14
              33
              17
           18
              27
        11
           19
           21
    === adding 8 ===
     -- bubble_up(10) -- 
     5
        9
           14
              33
              17
           18
              27
              8
        11
           19
           21
     -- bubble_up(4) -- 
     5
        9
           14
              33
              17
           8
              27
              18
        11
           19
           21
     -- bubble_up(1) -- 
     5
        8
           14
              33
              17
           9
              27
              18
        11
           19
           21
    === removing smallest ===
     -- bubble_down(0) -- 
     18
        8
           14
              33
              17
           9
              27
        11
           19
           21
     -- bubble_down(1) -- 
     8
        18
           14
              33
              17
           9
              27
        11
           19
           21
     -- bubble_down(4) -- 
     8
        9
           14
              33
              17
           18
              27
        11
           19
           21
    === popping until empty ===
    8 9 11 14 17 18 19 21 27 33 




 -------------------------------------------------

Jim Mahoney | cs.bennington.college | March 2021 | MIT License
"""

#
# The heap stores these indeces
#
#    0 1 2 3 4 5 6 
#
# as these places in a binary tree
#
#                0
#           1        2
#         3   4     5   6          
#
# and we can find the index of parents and children
# with these formulas.
#
#      left_child        2*i + 1     (for parent at i)
#      right_child       2*i + 2
#
#      parent           (n-1)//2   (for child at n)
#
#      check : 
#      (5 - 1) // 2       =>   2
#      (6 - 1) // 2  = 5 //2 =  round_down(2.5) => 2
#
#

infinity = float('inf')  # part of the IEEE floating point standard

def parent(i):
    """ return the index of the parent of node with index i 
        >>> parent(5)
        2
        >>> parent(6)
        2
    """
    # Remember that // is the "integer division" operator in python.
    return (i-1) // 2

def left(i):
    """ return index of left child
        >>> left(1)
        3
        >>> left(2)
        5
    """
    return 2*i+1

def right(i):
    """ return index of right child
        >>> right(1)
        4
        >>> right(2)
        6
    """
    return 2*i+2

class Heap:
    
    def __init__(self):
        """ initialize with runestone example as test case """
        self.verbose = True
        self.array = [5, 9, 11, 14, 18, 19, 21, 33, 17, 27]

    def __str__(self):
        """ define what str(self) does """
        return self.tree(0)[:-1]

    def __len__(self):
        """ define what len(self) does """
        return len(self.array)

    def has(self, i):
        """ True if index i exists """
        return 0 <= i < len(self.array)

    def is_empty(self):
        """ True if there are no values stored """
        return len(self) == 0
    
    def value(self, i):
        """ return value of i'th element """
        return self.array[i]
    
    def tree(self, i, indent=1):
        """ return string view of binary tree from index i """
        # recursive ... indenting each child further to right
        if not self.has(i):
            return ''
        else:
            result = ' ' * indent + str(self.array[i]) + '\n'
            for child in (left(i), right(i)):
                if self.has(child):
                    result += self.tree(child, indent+3)
            return result

    def bubble_up(self, i):
        """ percolate a node upwards until it reaches the right place """
        if self.verbose:
            print(f" -- bubble_up({i}) -- ")
            print(self)
        if i > 0:
            if self.value(i) < self.value(parent(i)):
                self.swap(i, parent(i))
                self.bubble_up(parent(i))

    def bubble_down(self, i):
        """ percolate a node downwards until it reaches the right place """
        if self.verbose:
            print(f" -- bubble_down({i}) -- ")
            print(self)
        destination = self.smallest(i)
        if destination:
            self.swap(i, destination)
            self.bubble_down(destination)
                
    def smallest(self, i):
        """ return index of smaller of two children below i or None """
        # If no children, or current is smallest, return None.
        self_value = self.value(i)
        right_value = self.value(right(i)) if self.has(right(i)) else infinity
        left_value = self.value(left(i)) if self.has(left(i)) else infinity
        if self_value <= right_value and self_value <= left_value:
            return None
        elif left_value < self_value and left_value < right_value:
            return left(i)
        else:
            return right(i)
                
    def swap(self, i, j):
        """ swap the values at indices i,j """
        (self.array[i], self.array[j]) = (self.array[j], self.array[i])

    def last(self):
        """ return index of last element """
        return len(self.array) - 1

    def push(self, value):
        """ add a value to the heap """
        # First stick it on the end of the array ...
        self.array.append(value)
        # ... and then swap it upwards until the one above is smaller
        self.bubble_up(self.last())

    def pop(self):
        """ return smallest value """
        if len(self) == 1:
            return self.array.pop()   # only 1 element; remove & return it.
        # First remember the value which was at index 0, to return when done;
        result = self.array[0]
        # replace that with end value, to restore the proper shape;
        self.array[0] = self.array.pop()  # (list pop removes last element)
        # then swap it downwards until the ones below are bigger, 
        self.bubble_down(0)
        # and finally, return the value we saved.
        return result
        
def main():
    heap = Heap()

    print("=== original ===")
    print(heap)

    print("=== adding 8 ===")
    heap.push(8)

    print("=== removing smallest ===")
    small = heap.pop()

    print("=== popping until empty ===")
    heap.verbose = False
    while not heap.is_empty():
        print(heap.pop(), end=' ')
    print()
    
    
if __name__ == "__main__":
    import doctest
    doctest.testmod()
    main()