#ifndef _TREEMAP_H_
#define _TREEMAP_H_

#include <iostream>

template <class K, class V>
struct treemap_t {
  struct pair_t {
    K key;
    V value;
  };
  
  struct node_t : pair_t {
    node_t* parent;
    node_t* left;
    node_t* right;

    node_t(node_t* parent) : parent { parent }, left { nullptr }, right { nullptr } {
    }
  };

  struct iterator_t {
    node_t* cur;

    iterator_t(node_t* root) : cur { root } {
      if(cur != nullptr)
        for(; cur->left != nullptr; cur = cur->left) { }
    }
    
    void next() {
      if(cur->right != nullptr) {
        for(cur = cur->right; cur->left != nullptr; cur = cur->left) { }
      } else {
        while(cur->parent != nullptr && cur->parent->right == cur)
          cur = cur->parent;

        cur = cur->parent;
      }
    }
    
    pair_t& operator *() {
      return *cur;
    }
    
    iterator_t& operator ++() {
      next();
      return *this;
    }
    
    bool operator != (const iterator_t& it) const {
      return cur != it.cur;
    }
  };
  
  node_t* root = nullptr;

  node_t* lookup(node_t** pcur, node_t* parent, const K* key, bool create) {
    if(*pcur == nullptr) {
      if(create) {
        *pcur = new node_t { parent };
        (*pcur)->key = *key;
      }
      else
        return nullptr;
    }

    node_t* cur = *pcur;
    
    if(*key == cur->key)
      return cur;
    else if(*key < cur->key)
      return lookup(&cur->left, cur, key, create);
    else
      return lookup(&cur->right, cur, key, create);
  }

  void set_at(const K& key, const V& value) {
    node_t* res = lookup(&root, nullptr, &key, true);
    res->value = value;
  }

  V& at(const K& key) {
    node_t* res = lookup(&root, nullptr, &key, false);
    if(res == nullptr)
      throw std::out_of_range("index out of range");
    return res->value;
  }
  
  iterator_t begin() {
    return iterator_t { root };
  }
  
  iterator_t end() {
    return iterator_t { nullptr };
  }

  void display(node_t* cur, std::string prefix) {
    if(cur != nullptr) {
      display(cur->left, prefix + "  ");
      std::cout << prefix << "(" << cur->key << ", " << cur->value << ")" << std::endl;
      display(cur->right, prefix + "  ");
    }
  }
  
  void display() {
    display(root, "");
  }

  V& operator[] (const K& key) {
    return lookup(&root, nullptr, &key, true)->value;
  }  
};

#endif