#ifndef PERFECT_HASH_MAP_H
#define PERFECT_HASH_MAP_H

#include <algorithm>
#include <iostream>
#include <vector>

// Avoid a dependence on libHalide by defining a local variant we can use
struct PerfectHashMapAsserter {
    const bool c;

    PerfectHashMapAsserter(bool c)
        : c(c) {
    }

    template<typename T>
    PerfectHashMapAsserter &operator<<(T &&t) {
        if (!c) {
            std::cerr << t;
        }
        return *this;
    }
    ~PerfectHashMapAsserter() {
        if (!c) {
            exit(-1);
        }
    }
};

// A specialized hash map used in the autoscheduler. It can only grow,
// and it requires a perfect hash in the form of "id" and "max_id"
// fields on each key. If the keys don't all have a consistent max_id,
// or if you call make_large with the wrong max_id, you get UB. If you
// think that might be happening, uncomment the assertions below for
// some extra checking.

template<typename K, typename T, int max_small_size = 4, typename phm_assert = PerfectHashMapAsserter>
class PerfectHashMap {

    using storage_type = std::vector<std::pair<const K *, T>>;

    storage_type storage;

    int occupied = 0;

    // Equivalent to storage[i], but broken out into a separate method
    // to allow for bounds checks when debugging this.
    std::pair<const K *, T> &storage_bucket(int i) {
        /*
        phm_assert(i >= 0 && i < (int)storage.size())
            << "Out of bounds access: " << i << " " << storage.size() << "\n";
        */
        return storage[i];
    }

    const std::pair<const K *, T> &storage_bucket(int i) const {
        /*
        phm_assert(i >= 0 && i < (int)storage.size())
            << "Out of bounds access: " << i << " " << storage.size() << "\n";
        */
        return storage[i];
    }

    enum {
        Empty = 0,  // No storage allocated
        Small = 1,  // Storage is just an array of key/value pairs
        Large = 2   // Storage is an array with empty slots, indexed by the 'id' field of each key
    } state = Empty;

    void upgrade_from_empty_to_small() {
        storage.resize(max_small_size);
        state = Small;
    }

    void upgrade_from_empty_to_large(int n) {
        storage.resize(n);
        state = Large;
    }

    void upgrade_from_small_to_large(int n) {
        phm_assert(occupied <= max_small_size) << occupied << " " << max_small_size << "\n";
        storage_type tmp(n);
        state = Large;
        tmp.swap(storage);
        int o = occupied;
        for (int i = 0; i < o; i++) {
            emplace_large(tmp[i].first, std::move(tmp[i].second));
        }
        occupied = o;
    }

    // Methods when the map is in the empty state
    T &emplace_empty(const K *n, T &&t) {
        upgrade_from_empty_to_small();
        storage_bucket(0).first = n;
        storage_bucket(0).second = std::move(t);
        occupied = 1;
        return storage_bucket(0).second;
    }

    const T &get_empty(const K *n) const {
        phm_assert(0) << "Calling get on an empty PerfectHashMap";
        return unreachable_value();
    }

    T &get_empty(const K *n) {
        phm_assert(0) << "Calling get on an empty PerfectHashMap";
        return unreachable_value();
    }

    T &get_or_create_empty(const K *n) {
        occupied = 1;
        return emplace_empty(n, T());
    }

    bool contains_empty(const K *n) const {
        return false;
    }

    // Methods when the map is in the small state
    int find_index_small(const K *n) const {
        int i;
        for (i = 0; i < (int)occupied; i++) {
            if (storage_bucket(i).first == n) {
                return i;
            }
        }
        return i;
    }

    T &emplace_small(const K *n, T &&t) {
        int idx = find_index_small(n);
        if (idx >= max_small_size) {
            upgrade_from_small_to_large((int)(n->max_id));
            return emplace_large(n, std::move(t));
        }
        auto &p = storage_bucket(idx);
        if (p.first == nullptr) {
            occupied++;
            p.first = n;
        }
        p.second = std::move(t);
        return p.second;
    }

    const T &get_small(const K *n) const {
        int idx = find_index_small(n);
        return storage_bucket(idx).second;
    }

    T &get_small(const K *n) {
        int idx = find_index_small(n);
        return storage_bucket(idx).second;
    }

    T &get_or_create_small(const K *n) {
        int idx = find_index_small(n);
        if (idx >= max_small_size) {
            upgrade_from_small_to_large((int)(n->max_id));
            return get_or_create_large(n);
        }
        auto &p = storage_bucket(idx);
        if (p.first == nullptr) {
            occupied++;
            p.first = n;
        }
        return p.second;
    }

    bool contains_small(const K *n) const {
        int idx = find_index_small(n);
        return (idx < max_small_size) && (storage_bucket(idx).first == n);
    }

    // Methods when the map is in the large state
    T &emplace_large(const K *n, T &&t) {
        auto &p = storage_bucket(n->id);
        if (!p.first) {
            occupied++;
        }
        p.first = n;
        p.second = std::move(t);
        return p.second;
    }

    const T &get_large(const K *n) const {
        return storage_bucket(n->id).second;
    }

    T &get_large(const K *n) {
        return storage_bucket(n->id).second;
    }

    T &get_or_create_large(const K *n) {
        auto &p = storage_bucket(n->id);
        if (p.first == nullptr) {
            occupied++;
            p.first = n;
        }
        return storage_bucket(n->id).second;
    }

    bool contains_large(const K *n) const {
        return storage_bucket(n->id).first != nullptr;
    }

    void check_key(const K *n) const {
        /*
        phm_assert(n->id >= 0 && n->id < n->max_id)
            << "Invalid hash key: " << n->id << " " << n->max_id << "\n";
        phm_assert(state != Large || (int)storage.size() == n->max_id)
            << "Inconsistent key count: " << n->max_id << " vs " << storage.size() << "\n";
        */
    }

    // Helpers used to pacify compilers
    T &unreachable_value() {
        return storage_bucket(0).second;
    }

    const T &unreachable_value() const {
        return storage_bucket(0).second;
    }

public:
    // Jump straight to the large state
    void make_large(int n) {
        if (state == Empty) {
            upgrade_from_empty_to_large(n);
        } else if (state == Small) {
            upgrade_from_small_to_large(n);
        }
    }

    T &emplace(const K *n, T &&t) {
        check_key(n);
        switch (state) {
        case Empty:
            return emplace_empty(n, std::move(t));
        case Small:
            return emplace_small(n, std::move(t));
        case Large:
            return emplace_large(n, std::move(t));
        }
        return unreachable_value();
    }

    T &insert(const K *n, const T &t) {
        check_key(n);
        T tmp(t);
        switch (state) {
        case Empty:
            return emplace_empty(n, std::move(tmp));
        case Small:
            return emplace_small(n, std::move(tmp));
        case Large:
            return emplace_large(n, std::move(tmp));
        }
        return unreachable_value();
    }

    const T &get(const K *n) const {
        check_key(n);
        switch (state) {
        case Empty:
            return get_empty(n);
        case Small:
            return get_small(n);
        case Large:
            return get_large(n);
        }
        return unreachable_value();
    }

    T &get(const K *n) {
        check_key(n);
        switch (state) {
        case Empty:
            return get_empty(n);
        case Small:
            return get_small(n);
        case Large:
            return get_large(n);
        }
        return unreachable_value();
    }

    T &get_or_create(const K *n) {
        check_key(n);
        switch (state) {
        case Empty:
            return get_or_create_empty(n);
        case Small:
            return get_or_create_small(n);
        case Large:
            return get_or_create_large(n);
        }
        return unreachable_value();
    }

    bool contains(const K *n) const {
        check_key(n);
        switch (state) {
        case Empty:
            return contains_empty(n);
        case Small:
            return contains_small(n);
        case Large:
            return contains_large(n);
        }
        return false;  // Unreachable
    }

    size_t size() const {
        return occupied;
    }

    struct iterator {
        std::pair<const K *, T> *iter, *end;

        void operator++(int) {
            do {
                iter++;
            } while (iter != end && iter->first == nullptr);
        }

        void operator++() {
            (*this)++;
        }

        const K *key() const {
            return iter->first;
        }

        T &value() const {
            return iter->second;
        }

        bool operator!=(const iterator &other) const {
            return iter != other.iter;
        }

        bool operator==(const iterator &other) const {
            return iter == other.iter;
        }

        std::pair<const K *, T> &operator*() {
            return *iter;
        }
    };

    struct const_iterator {
        const std::pair<const K *, T> *iter, *end;

        void operator++(int) {
            do {
                iter++;
            } while (iter != end && iter->first == nullptr);
        }

        void operator++() {
            (*this)++;
        }

        const K *key() const {
            return iter->first;
        }

        const T &value() const {
            return iter->second;
        }

        bool operator!=(const const_iterator &other) const {
            return iter != other.iter;
        }

        bool operator==(const const_iterator &other) const {
            return iter == other.iter;
        }

        const std::pair<const K *, T> &operator*() const {
            return *iter;
        }
    };

    iterator begin() {
        if (state == Empty) {
            return end();
        }
        iterator it;
        it.iter = storage.data();
        it.end = it.iter + storage.size();
        if (it.key() == nullptr) {
            it++;
        }
        phm_assert(it.iter == it.end || it.key());
        return it;
    }

    iterator end() {
        iterator it;
        it.iter = it.end = storage.data() + storage.size();
        return it;
    }

    const_iterator begin() const {
        if (storage.empty()) {
            return end();
        }
        const_iterator it;
        it.iter = storage.data();
        it.end = it.iter + storage.size();
        if (it.key() == nullptr) {
            it++;
        }
        phm_assert(it.iter == it.end || it.key());
        return it;
    }

    const_iterator end() const {
        const_iterator it;
        it.iter = it.end = storage.data() + storage.size();
        return it;
    }
};

#endif