//
//  trie.hpp
//  trie
//
//  Created by Sam Jaffe on 6/16/17.
//

#pragma once

#include "pointers/const_propogating_ptr.hpp"
#include "pointers/not_null.hpp"
#include "pointers/value_ptr.hpp"
#include <map>

template <typename K, typename V, typename Compare = std::less<K>> class trie;

namespace detail {
  template <typename Trie, typename Iter> class trie_iterator_base;
  template <typename Iter, typename Trie, typename KS>
  Iter find_impl(Trie * tr, KS const & keys);
}

template <typename Trie, typename Iter> class trie_iterator;
template <typename Trie, typename Iter> class trie_reverse_iterator;

/*
 * \brief A trie is a type of associative container where each element is
 * locatable by a string of keys, including a zero-length string. Sometimes
 * known as a prefix-tree.
 * \tparam K The key type. The container is indexable by both single keys, and
 * iterable collections of keys e.g. std::initializer_list<K>.
 * \tparam V The value type of stored data.
 * \tparam Compare The comparator function, used by the backing std::map object
 * to order this container.
 *
 * In principle, a trie<K, V> is equal to a std::map<std::vector<K>, V> with the
 * additional features of being able to locate all elements whose keys start
 * with the key fragment { K1, K2, ..., Kn }.
 *
 * Because the conceptual structure of a trie is a tree, both pre-order and
 * port-order iteration are supported. In-order traversal is not supported,
 * because it doesn't translate well with non-binary trees.
 */
template <typename K, typename V, typename Compare> class trie {
public:
  using key_type = K;
  using mapped_type = V;
  using key_compare = Compare;
  using map_type =
      std::map<K, pointers::const_propogating_ptr<pointers::value_ptr<trie>>,
               Compare>;

  using local_iterator = typename map_type::iterator;
  using local_const_iterator = typename map_type::const_iterator;
  using local_reverse_iterator = typename map_type::reverse_iterator;
  using local_const_reverse_iterator =
      typename map_type::const_reverse_iterator;

  using iterator = trie_iterator<trie, local_iterator>;
  using const_iterator = trie_iterator<trie const, local_const_iterator>;
  using post_iterator = trie_reverse_iterator<trie, local_iterator>;
  using const_post_iterator =
      trie_reverse_iterator<trie const, local_const_iterator>;
  using reverse_iterator = trie_reverse_iterator<trie, local_reverse_iterator>;
  using const_reverse_iterator =
      trie_reverse_iterator<trie const, local_const_reverse_iterator>;

private:
  template <typename KS>
  using element_type = std::decay_t<decltype(*std::begin(std::declval<KS>()))>;
  template <typename KS>
  using is_collection_t =
      std::enable_if_t<std::is_same<K, element_type<KS>>::value>;
  using impl_iterator = detail::trie_iterator_base<trie, local_iterator>;
  using impl_const_iterator =
      detail::trie_iterator_base<trie const, local_const_iterator>;

private:
  mapped_type value_{};
  map_type impl_{};

public:
  trie() = default;
  trie(trie const & other);
  trie(trie && other) : trie() { swap(*this, other); }
  trie(mapped_type const & root, std::map<K, trie, Compare> && children);
  ~trie() { clear(); }
  trie & operator=(mapped_type const & value);
  trie & operator=(trie const & value);
  trie & operator=(trie && value);

  operator mapped_type &() { return value(); }
  operator mapped_type const &() const { return value(); }
  mapped_type & value() { return value_; }
  mapped_type const & value() const { return value_; }

  template <typename KS, typename = is_collection_t<KS>>
  trie & operator[](KS const & keys) {
    return *emplace(keys).first.parent_trie_.top();
  }
  trie & operator[](key_type const & key) {
    return *emplace(key).first.parent_trie_.top();
  }
  trie & operator[](std::initializer_list<key_type> keys) {
    return operator[]<>(keys);
  }

  template <typename KS>
  std::pair<iterator, bool> insert(KS const & keys, mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  std::pair<iterator, bool> insert(key_type const & key,
                                   mapped_type const & value) {
    return emplace_impl({key}, value);
  }
  std::pair<iterator, bool> insert(std::initializer_list<key_type> keys,
                                   mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  template <typename KS, typename... Args>
  std::pair<iterator, bool> emplace(KS && keys, Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type const & key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type && key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(std::initializer_list<key_type> keys,
                                    Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }

  iterator begin() { return {this}; }
  iterator end() { return {}; }
  const_iterator begin() const { return {this}; }
  const_iterator end() const { return {}; }
  const_iterator cbegin() const { return begin(); }
  const_iterator cend() const { return end(); }
  reverse_iterator rbegin() { return {this}; }
  reverse_iterator rend() { return {}; }
  const_reverse_iterator rbegin() const { return {this}; }
  const_reverse_iterator rend() const { return {}; }
  const_reverse_iterator crbegin() const { return rbegin(); }
  const_reverse_iterator crend() const { return rend(); }

  local_iterator local_begin() { return impl_.begin(); }
  local_iterator local_end() { return impl_.end(); }
  local_const_iterator local_begin() const { return impl_.begin(); }
  local_const_iterator local_end() const { return impl_.end(); }

  local_reverse_iterator local_rbegin() { return impl_.rbegin(); }
  local_reverse_iterator local_rend() { return impl_.rend(); }
  local_const_reverse_iterator local_rbegin() const { return impl_.rbegin(); }
  local_const_reverse_iterator local_rend() const { return impl_.rend(); }

  template <typename KS> iterator find(KS const & keys) {
    return detail::find_impl<impl_iterator>(this, keys);
  }
  iterator find(key_type const & key) {
    return find_impl<impl_iterator>(this, {key});
  }
  iterator find(std::initializer_list<key_type> keys) {
    return find_impl<impl_iterator>(this, keys);
  }
  template <typename KS> const_iterator find(KS const & keys) const {
    return detail::find_impl<impl_const_iterator>(this, keys);
  }
  const_iterator find(key_type const & key) const {
    return find_impl<impl_const_iterator>(this, {key});
  }
  const_iterator find(std::initializer_list<key_type> keys) const {
    return find_impl<impl_const_iterator>(this, keys);
  }

  template <typename KS> void erase(KS const & keys) { drop(find(keys)); }
  void erase(key_type const & key) { drop(find(key)); }
  void erase(std::initializer_list<key_type> keys) { drop(find(keys)); }

  void clear();

private:
  void drop(iterator it);

  template <typename Iter, typename Trie, typename KS>
  friend Iter detail::find_impl(Trie * tr, KS const & keys);
  template <typename Iter, typename Trie>
  static Iter find_impl(Trie * tr, std::initializer_list<key_type> keys) {
    return detail::find_impl<Iter>(tr, keys);
  }

  template <typename KS, typename... Args>
  std::pair<impl_iterator, bool> emplace_impl(KS && keys, Args &&... args);
  template <typename... Args>
  std::pair<impl_iterator, bool>
  emplace_impl(std::initializer_list<key_type> keys, Args &&... args) {
    return emplace_impl<std::initializer_list<key_type>>(
        std::move(keys), std::forward<Args>(args)...);
  }

  template <typename Key>
  void insert_impl(impl_iterator & out, bool & create, Key && key);

  friend bool operator==(trie const & lhs, trie const & rhs) {
    const_iterator it1 = lhs.begin(), it2 = rhs.begin();
    const_iterator const end1 = lhs.end(), end2 = lhs.end();
    for (; it1 != end1 && it2 != end2; ++it1, ++it2) {
      if (it1 != it2) { return false; }
    }
    return it1 == end1 && it2 == end2;
  }

  friend bool operator!=(trie const & lhs, trie const & rhs) {
    return !(lhs == rhs);
  }

  friend void swap(trie & lhs, trie & rhs) {
    using std::swap;
    swap(lhs.value_, rhs.value_);
    swap(lhs.impl_, rhs.impl_);
  }
};

#include "trie.tpp"