#ifndef VTR_BIMAP
#define VTR_BIMAP
#include <map>
#include <unordered_map>
#include "vtr_flat_map.h"
#include "vtr_linear_map.h"

#include "vtr_error.h"

namespace vtr {

/*
 * A map-like class which provides a bi-directonal mapping between key and value
 *
 * Keys and values can be looked up directly by passing either the key or value.
 * the indexing operator will throw if the key/value does not exist.
 */
template<class K, class V, template<typename ...> class Map=std::map, template<typename ...> class InvMap=std::map>
class bimap {
    public: //Public types
        typedef typename Map<K,V>::const_iterator iterator;
        typedef typename InvMap<V,K>::const_iterator inverse_iterator;
    public: //Accessors

        //Iterators
        iterator begin() const { return map_.begin(); }
        iterator end() const { return map_.end(); }
        inverse_iterator inverse_begin() const { return inverse_map_.begin(); }
        inverse_iterator inverse_end() const { return inverse_map_.end(); }

        //Return an iterator to the key-value pair matching key, or end() if not found
        iterator find(const K key) const {
            return map_.find(key);
        }

        //Return an iterator to the value-key pair matching value, or inverse_end() if not found
        inverse_iterator find(const V value) const {
            return inverse_map_.find(value);
        }

        //Return an immutable reference to the value matching key
        //Will throw an exception if key is not found
        const V& operator[] (const K key) const { 
            auto iter = find(key);
            if(iter == end()) {
                throw VtrError("Invalid bimap key during look-up", __FILE__, __LINE__);
            }
            return iter->second;
        }

        //Return an immutable reference to the key matching value
        //Will throw an exception if value is not found
        const K& operator[] (const V value) const {
            auto iter = find(value);
            if(iter == inverse_end()) {
                throw VtrError("Invalid bimap value during inverse look-up", __FILE__, __LINE__);
            }
            return iter->second;
        }

        //Return the number of key-value pairs stored
        std::size_t size() const { VTR_ASSERT(map_.size() == inverse_map_.size()); return map_.size(); }

        //Return true if there are no key-value pairs stored
        bool empty() const { return (size() == 0); }

        //Return true if the specified key exists
        bool contains(const K key) const { return find(key) != end(); }

        //Return true if the specified value exists
        bool contains(const V value) const { return find(value) != inverse_end(); }

    public: //Mutators
        
        //Drop all stored key-values
        void clear() { map_.clear(); inverse_map_.clear(); }

        //Insert a key-value pair, if not already in map
        std::pair<iterator,bool> insert(const K key, const V value) {
            auto ret1 = map_.insert({key,value});
            auto ret2 = inverse_map_.insert({value,key});

            VTR_ASSERT(ret1.second == ret2.second);

            //Return true if inserted
            return ret1;
        }

        //Update a key-value pair, will insert if not already in map
        void update(const K key, const V value) {
            map_[key] = value;
            inverse_map_[value] = key;
        }

        //Remove the specified key (and it's associated value)
        void erase(const K key) {
            auto iter = map_.find(key);
            if(iter != map_.end()) {
                V val = iter->second;
                map_.erase(iter);

                auto inv_iter = inverse_map_.find(val);
                VTR_ASSERT(inv_iter != inverse_map_.end());
                inverse_map_.erase(inv_iter);
            }
        }

        //Remove the specified value (and it's associated key)
        void erase(const V val) {
            auto inv_iter = inverse_map_.find(val);
            if(inv_iter != inverse_map_.end()) {
                K key = inv_iter->second;
                inverse_map_.erase(inv_iter);

                auto iter = map_.find(key);
                VTR_ASSERT(iter != map_.end());
                map_.erase(iter);
            }
        }

        //Swap (this enables std::swap via ADL)
        friend void swap(bimap<K,V,Map,InvMap>& x, bimap<K,V,Map,InvMap>& y) {
            std::swap(x.map_, y.map_);
            std::swap(x.inverse_map_, y.inverse_map_);
        }
    private:
        Map<K,V> map_;
        InvMap<V,K> inverse_map_;
};

template<class K, class V>
using unordered_bimap = bimap<K,V,std::unordered_map,std::unordered_map>;

template<class K, class V>
using flat_bimap = bimap<K,V,vtr::flat_map,vtr::flat_map>;

template<class K, class V>
using linear_bimap = bimap<K,V,vtr::linear_map,vtr::linear_map>;

}

#endif