#ifndef HEAP_PRIORITY_QUEUE_H
#define HEAP_PRIORITY_QUEUE_H

#include "IPriorityQueue.h"

#include <iostream>
#include <vector>
#include <functional>

template<class K, class V> class HeapPriorityQueue : public IPriorityQueue < K, V > {
private:
    std::vector<Entry<K, V>> entries;
    std::function<int(K, K)> compareFunction;

    // ------------------------------

    void upheap(unsigned int currentEntryIndex) {
        while (currentEntryIndex > 0) {
            int parentEntryIndex = indexOfParent(currentEntryIndex);

            Entry<K, V> parent = entries.at(parentEntryIndex);
            K parentKey = parent.getKey();

            Entry<K, V> current = entries.at(currentEntryIndex);
            K currentKey = current.getKey();

            if (compareFunction(parentKey, currentKey) > 0) {
                swap(parentEntryIndex, currentEntryIndex);
                currentEntryIndex = parentEntryIndex;
            }
            else {
                return;
            }
        }
    }

    void downheap(unsigned int currentEntryIndex) {
        while (currentEntryIndex < entries.size() - 1) {
            bool hasLeft = hasLeftEntry(currentEntryIndex);
            bool hasRight = hasRightEntry(currentEntryIndex);
            int leftIndex = indexOfLeft(currentEntryIndex);
            int rightIndex = indexOfRight(currentEntryIndex);

            unsigned int smallestBelowIndex = -1;

            if (hasLeft) {
                smallestBelowIndex = leftIndex;
            }
            else { break; }

            if (hasRight) {
                // If leftKey is bigger than rightKey
                if (compareFunction(entries.at(leftIndex).getKey(), entries.at(rightIndex).getKey()) > 0) {
                    smallestBelowIndex = rightIndex;
                }
            }

            // --------------------------------
            // if the smallest below is smaller than current -> swap + update i
            if (smallestBelowIndex != -1) {
                K currentKey = entries.at(currentEntryIndex).getKey();
                K smallestKey = entries.at(smallestBelowIndex).getKey();
                if (compareFunction(currentKey, smallestKey) > 0) {
                    swap(currentEntryIndex, smallestBelowIndex);
                    currentEntryIndex = smallestBelowIndex;
                }
            }
            else { break; }
        }
    }

    unsigned int indexOfParent(unsigned int i) {
        return (i - 1) / 2;
    }

    unsigned int indexOfLeft(unsigned int i) {
        return i * 2 + 1;
    }

    unsigned int indexOfRight(unsigned int i) {
        return i * 2 + 2;
    }

    bool hasLeftEntry(unsigned int i) {
        return this->indexOfLeft(i) < entries.size();
    }

    bool hasRightEntry(unsigned int i) {
        return this->indexOfRight(i) < entries.size();
    }

    void swap(unsigned int i, unsigned int j) {
        if (i >= entries.size() || j >= entries.size()) {
            std::cerr << "swap(i, j) :: Invalid arguments!";
        }

        iter_swap(entries.begin() + i, entries.begin() + j);
    }

public:
    HeapPriorityQueue<K, V>(std::function<int(K, K)> compareFunction) {
        entries = std::vector<Entry<K, V>>();
        this->compareFunction = compareFunction;
    }

    unsigned int size() override {
        return entries.size();
    }

    bool isEmpty() override {
        return this->size() == 0;
    };

    Entry<K, V> * insert(K key, V value) override {
        Entry<K, V> newEntry = Entry<K, V>(key, value);
        entries.push_back(newEntry);
        this->upheap(entries.size() - 1);

        return &newEntry;
    }

    Entry<K, V> * min() override {
        return (!isEmpty()) ? &entries.at(0) : NULL;
    }

    Entry<K, V> * removeMin() override {
        if (isEmpty()) {
            return NULL;
        }

        Entry<K, V> removed = *(this->min());
        this->swap(0, entries.size() - 1);
        entries.erase(entries.end() - 1);
        this->downheap(0);

        return &removed;
    }

    // --------------- Utility ------------------

    void clear() {
        entries = std::vector<Entry<K, V>>();
    }

    void print() {
        std::cout << "Priority queue content:" << std::endl;
        for (Entry<K, V> entry : entries) {
            std::cout << "[" << entry.getKey() << " -> " << entry.getValue() << "]" << std::endl;
        }
        std::cout << std::endl;
    }
};

#endif