# -*- coding: utf-8 -*-
"""
Based on Matteo Dell'Amico's solution
https://gist.github.com/matteodellamico/4451520
"""

from heapq import heapify, heappush, heappop

from six import iteritems


class PriorityDict(dict):
    """
    Dictionary that can be used as a priority queue.

    Keys of the dictionary are items to be put into the queue, and values
    are their respective priorities. All dictionary methods work as expected.
    The advantage over a standard heapq-based priority queue is
    that priorities of items can be efficiently updated (amortized O(1))
    using code as 'the_dict[item] = new_priority.'

    The 'smallest' method can be used to return the object with lowest
    priority, and 'pop_smallest' also removes it.

    The 'sorted_iter' method provides a destructive sorted iterator.
    """

    def __init__(self, *args, **kwargs):
        super(PriorityDict, self).__init__(*args, **kwargs)
        self._rebuild_heap()

    def _rebuild_heap(self):
        self._heap = [(val, key) for key, val in iteritems(self)]
        heapify(self._heap)

    def smallest(self):
        """
        Return the item with the lowest priority.
        Raises IndexError if the object is empty.
        """

        heap = self._heap
        val, key = heap[0]
        while key not in self or self[key] != val:
            heappop(heap)
            val, key = heap[0]
        return key, val

    def pop_smallest(self):
        """
        Return the item with the lowest priority and remove it.

        Raises IndexError if the object is empty.
        """

        heap = self._heap
        val, key = heappop(heap)
        while key not in self or self[key] != val:
            val, key = heappop(heap)
        del self[key]
        return key, val

    def __setitem__(self, key, val):
        """
        We are not going to remove the previous value from the heap,
        since this would have a cost O(n).
        """

        super(PriorityDict, self).__setitem__(key, val)

        if len(self._heap) < 2 * len(self):
            heappush(self._heap, (val, key))
        else:
            # When the heap grows larger than 2 * len(self), we rebuild it
            # from scratch to avoid wasting too much memory.
            self._rebuild_heap()

    def setdefault(self, key, d=None):
        if key not in self:
            self[key] = d
            return d
        return self[key]

    def update(self, *args, **kwargs):
        """
        Reimplementing dict.update is tricky -- see e.g.
        http://mail.python.org/pipermail/python-ideas/2007-May/000744.html
        We just rebuild the heap from scratch after passing to super.
        """

        super(PriorityDict, self).update(*args, **kwargs)
        self._rebuild_heap()

    def sorted_iter(self):
        """
        Sorted iterator of the priority dictionary items.
        Beware: this will destroy elements as they are returned.
        """

        while self:
            yield self.pop_smallest()
