//
//  bucket_hash_map.hpp
//  bucket_hash_map
//
//  Created by Sam Jaffe on 2/7/17.
//

#pragma once

#include <algorithm>
#include <list>
#include <memory>
#include <utility>
#include <vector>

#include "iterator/recursive_iterator.hpp"

template <typename K, typename V, typename Hash = std::hash<K>, typename KeyEqual = std::equal_to<K> >
class bucket_hash_map;

template <typename K, typename V, typename Hash, typename KeyEqual>
class bucket_hash_map {
public: // typedefs
  using key_type = K;
  using mapped_type = V;
  using value_type = std::pair<key_type const, mapped_type>;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;
  
  using hasher = Hash;
  using key_equal = KeyEqual;
  
  using reference = value_type &;
  using const_reference = value_type const &;
  using pointer = value_type *;
  using const_pointer = value_type const *;
  
  using bucket_type = std::list<value_type>;
  using impl_type = std::vector<bucket_type>;
  
  template <typename It>
  using iter_wrap = iterator::recursive_iterator<It>;
  using iterator = iter_wrap<typename impl_type::iterator>;
  using const_iterator = iter_wrap<typename impl_type::const_iterator>;
  using local_iterator = typename bucket_type::iterator;
  using local_const_iterator = typename bucket_type::const_iterator;
  
  static const constexpr size_type default_buckets = 10;
  static const constexpr float default_load_factor = 1.0F;
  static const constexpr float growth_factor = 2.0F;
public:
  // Construction
  bucket_hash_map() : bucket_hash_map(default_buckets) {}
  
  bucket_hash_map(size_type num_buckets,
                  hasher const & hash = hasher(),
                  key_equal const & keq = key_equal()) :
  buckets_(num_buckets), hasher_(hash), key_equals_(keq) {}
  
  template <typename InputIt>
  bucket_hash_map(InputIt first, InputIt last,
                  size_type num_buckets = default_buckets,
                  hasher const & hash = hasher(),
                  key_equal const & keq = key_equal())
  : bucket_hash_map(num_buckets, hash, keq) {
    insert(first, last);
  }
  
  bucket_hash_map(std::initializer_list<value_type> && ilist,
                  size_type num_buckets = default_buckets,
                  hasher const & hash = hasher(),
                  key_equal const & keq = key_equal())
  : bucket_hash_map(num_buckets, hash, keq) {
    insert(ilist);
  }
  
  // Metadata
  bool empty() const { return size() == 0; }
  size_type size() const { return size_; }

  // Iterators
  iterator begin() { return { make_end_aware_iterator( buckets_ ) }; }
  const_iterator begin() const { return cbegin(); }
  const_iterator cbegin() const { return { make_end_aware_iterator( buckets_ ) }; }
  iterator end() { return iterator{}; }
  const_iterator end() const { return cend(); }
  const_iterator cend() const { return const_iterator{}; }
  
  // Create
  std::pair<iterator, bool> insert(value_type const & vt) {
    maybe_expand(1);
    auto lookup = lookup_impl(buckets_, vt.first);
    auto const end = lookup.first->end();
    bool const create = lookup.second == end;
    if ( create ) {
      ++size_;
      lookup.second = lookup.first->insert(end, vt);
    }
    return { {::iterator::in_place, make_end_aware_iterator(lookup.first, buckets_.end()), make_end_aware_iterator(lookup.second, end)}, create };
  }
  iterator insert(const_iterator, value_type const & vt) {
    return insert(vt).first;
  }
  template <typename InputIt>
  void insert(InputIt first, InputIt last) {
    maybe_expand(static_cast<std::size_t>(std::distance(first, last)));
    while (first != last) { insert(*first++); }
  }
  void insert(std::initializer_list<value_type> ilist) { insert(ilist.begin(), ilist.end()); }
  
  // Access
  mapped_type & operator[](key_type const & key) {
    auto lookup = lookup_impl(buckets_, key);
    auto end = lookup.first->end();
    if (lookup.second == end) {
      ++size_;
      return lookup.first->insert(lookup.second, {key, mapped_type{}})->second;
    } else {
      return lookup.second->second;
    }
  }
  mapped_type & at(key_type const & key) {
    auto it = find(key);
    if (it == end()) { throw std::out_of_range{"no element at key"}; }
    return it->second;
  }
  mapped_type const & at(key_type const & key) const {
    auto it = find(key);
    if (it == end()) { throw std::out_of_range{"no element at key"}; }
    return it->second;
  }

  iterator find(key_type const & key) {
    return find_impl<iterator>(buckets_, key);
  }
  std::pair<iterator, iterator> equal_range(key_type const & key) {
    auto it = find(key);
    return { it, ++iterator(it) };
  }
  const_iterator find(key_type const & key) const {
    return find_impl<const_iterator>(buckets_, key);
  }
  std::pair<const_iterator, const_iterator> equal_range(key_type const & key) const {
    auto it = find(key);
    return { it, ++const_iterator(it) };
  }
  size_type count(key_type const & key) const { return find(key) != end(); }
  
  // Remove
  iterator erase(const_iterator it) {
    return erase_impl(unconst_iterator(it));
  }
  iterator erase(const_iterator first, const_iterator last) {
    iterator it = unconst_iterator(first);
    while (last != it) { it = erase_impl(it); }
    return it;
  }
  size_type erase(key_type const & key) {
    size_type old_size = size_;
    erase(find(key));
    return old_size - size_; // 1 or 0
  }
  void clear() { erase(begin(), end()); }
  
  // Bucket Interaction Functions
  local_iterator begin(size_type bkt) { return buckets_[bkt].begin(); }
  local_const_iterator begin(size_type bkt) const { return buckets_[bkt].cbegin(); }
  local_const_iterator cbegin(size_type bkt) const { return buckets_[bkt].cbegin(); }
  local_iterator end(size_type bkt) { return buckets_[bkt].end(); }
  local_const_iterator end(size_type bkt) const { return buckets_[bkt].cend(); }
  local_const_iterator cend(size_type bkt) const { return buckets_[bkt].cend(); }
  size_type bucket_count() const { return buckets_.size(); }
  size_type bucket_size(size_type bkt) const { return buckets_[bkt].size(); }
  size_type bucket(key_type const & key) const { return hasher_(key) % bucket_count(); }
  
  // Hash Policy
  float load_factor() const { return static_cast<float>(size()) / bucket_count(); }
  float max_load_factor() const { return max_load_; }
  void max_load_factor(float max_load) { max_load_ = max_load; }
  void rehash(size_type buckets) {
    buckets = std::max({1UL, buckets, static_cast<size_type>(size() / max_load_factor())});
    bucket_hash_map next{buckets, hasher_, key_equals_};
    next.max_load_factor(max_load_factor());
    next.insert(begin(), end());
    swap(next);
  }
  void reserve(size_type count) {
    rehash(static_cast<size_type>(count / max_load_factor()));
  }
  
  void swap( bucket_hash_map & other ) {
    using std::swap;
    swap(size_, other.size_);
    swap(buckets_, other.buckets_);
    swap(max_load_, other.max_load_);
    swap(hasher_, other.hasher_);
    swap(key_equals_, other.key_equals_);
  }
  friend void swap(bucket_hash_map & lhs, bucket_hash_map & rhs) { lhs.swap(rhs); }
  friend bool operator==(bucket_hash_map const & lhs, bucket_hash_map const & rhs) {
    if (lhs.size() != rhs.size()) { return false; }
    auto end = rhs.end();
    for (auto & pair : lhs) {
      auto it = rhs.find(pair.first);
      if (it == end || it->second != pair.second) {
        return false;
      }
    }
    return true;
  }
  friend bool operator!=(bucket_hash_map const & lhs, bucket_hash_map const & rhs) {
    return !(lhs == rhs);
  }
private:
  void maybe_expand(size_type add) {
    if ( static_cast<float>(size() + add) / bucket_count() > max_load_factor() ) {
      reserve(std::max(size() + add, static_cast<size_type>(size() * growth_factor)));
    }
  }
  
  template <typename Bucket>
  auto lookup_impl(Bucket & bkt, key_type const & key) const -> std::pair<decltype(bkt.begin()), decltype(bkt.begin()->begin())> {
    auto listit = bkt.begin() + static_cast<std::ptrdiff_t>(bucket(key));
    auto it = search(listit, key);
    return {listit, it};
  }
  
  template <typename Iterator>
  auto search(Iterator lit, key_type const & key) const -> decltype(lit->begin()) {
    for ( auto it = lit->begin(); it != lit->end(); ++it ) {
      if ( key_equals_(key, it->first) ) { return it; }
    }
    return lit->end();
  }
  
  template <typename Iterator, typename Bucket>
  Iterator find_impl(Bucket & bkt, key_type const & key) const {
    auto lookup = lookup_impl(bkt, key);
    if (lookup.second == lookup.first->end()) {
      return Iterator{};
    } else {
      return { ::iterator::in_place, make_end_aware_iterator(lookup.first, bkt.end()), make_end_aware_iterator(lookup.second, lookup.first->end()) };
    }
  }
  
  iterator unconst_iterator(const_iterator it) {
    auto lit = std::get<0>(it);
    auto iter = std::get<1>(it);
    
    if ( lit.done() ) { return end(); }
    auto nit = buckets_.begin();
    std::advance(nit, std::distance(buckets_.cbegin(), lit.current()));
    
    return { ::iterator::in_place, make_end_aware_iterator(nit, buckets_.end()), make_end_aware_iterator(nit->erase(iter.current(), iter.current()), nit->end()) };
  }
  
  iterator erase_impl(iterator it) {
    auto b = std::get<0>(it);
    auto l = std::get<1>(it);
    if ( b.done() || l.done() ) { return it; }
    --size_;
    return { ::iterator::in_place, make_end_aware_iterator(b.current(), b.end()), make_end_aware_iterator(b->erase(l.current()), l.end()) };
  }
private: // members
  impl_type buckets_{default_buckets};
  hasher hasher_;
  key_equal key_equals_;
  //allocator_type alloc_;
  size_type size_{0};
  float max_load_{default_load_factor};
};