#ifndef PREFIX_SEARCH_H_A6672A1D
#define PREFIX_SEARCH_H_A6672A1D

#include <unordered_map>
#include <vector>
#include <stdexcept>

template <typename T>
class prefix_search {
public:

  prefix_search(size_t depth = 0) : depth_(depth), count_(0)
  {
  }

  size_t getDepth() const
  {
    return depth_;
  }

  size_t getCount() const
  {
    return count_;
  }

  void add(std::string& key, T val, size_t i = 0)
  {
    count_++;

    if (i == key.length())
    {
      elements_.push_back(val);
    } else {
      char next = key.at(i);

      if (!children_.count(next))
      {
        children_.emplace(next, prefix_search<T> { depth_ + 1 } );
      }

      children_.at(next).add(key, val, i + 1);
    }
  }

  T at(size_t i) const
  {
    if (i < elements_.size())
    {
      return elements_.at(i);
    } else {
      i -= elements_.size();

      for (const auto& mapping : children_)
      {
        if (mapping.second.count_ > i)
        {
          return mapping.second.at(i);
        } else {
          i -= mapping.second.count_;
        }
      }
    }

    throw std::out_of_range("Index out of range");
  }

  const prefix_search<T>& find(const std::string& val, size_t i) const
  {
    if (i == val.length() || !children_.count(val.at(i)))
    {
      return *this;
    }

    return children_.at(val.at(i)).find(val, i+1);
  }

private:

  size_t depth_;
  size_t count_;
  std::unordered_map<char, prefix_search<T>> children_;
  std::vector<T> elements_;
};

#endif /* end of include guard: PREFIX_SEARCH_H_A6672A1D */