#pragma once

#include <cassert>

namespace containers {

template<typename ValueT, typename KeyT>
class NTree {
public:
    enum {
        Dimensions = KeyT::Dimensions,
        Branches = 1<<Dimensions,
    };
    typedef ValueT ValueType;
    typedef KeyT KeyType;

    struct Node {
        Node **nodes;
        KeyType key;
        ValueType value;

        Node() : nodes(NULL) {
        }

        Node(const KeyType &key, const ValueType &value)
            : key(key), value(value), nodes(NULL) {
        }

        ~Node() {
            for (int i = 0; i < Branches; ++i) {
                delete nodes[i];
            }
            delete[] nodes;
        }

        const Node *get(const KeyType &key) const {
            if (this->key == key)
                return this;

            if (!nodes)
                return NULL;

            int index = key > this->key;
            assert((index >= 0)&&(index < Branches));

            return nodes[index]->get(key);
        }

        Node &insert(const KeyType &key, const ValueType &value) {
            if (this->key == key) {
                this->value = value;
                return *this;
            }

            int index = key > this->key;
            assert((index >= 0)&&(index < Branches));

            if (!nodes) {
                nodes = new Node*[Branches];
                for (int i = 0; i < Branches; ++i) {
                    nodes[i] = NULL;
                }
            }

            if (!nodes[index]) {
                Node *node = new Node(key, value);
                nodes[index] = node;
                return *node;
            }

            return nodes[index]->insert(key, value);
        }
    };

    Node *root;

    NTree() : root(NULL) {
    }

    const ValueType *get(const KeyType &key) const {
        if (!root)
            return NULL;
        const Node *node = root->get(key);
        if (!node)
            return NULL;
        return &node->value;
    }

    void insert(const KeyType &key, const ValueType &value) {
        if (!root) {
            root = new Node(key, value);
            return;
        }
        root->insert(key, value);
    }
};

namespace test {

struct CharV3 {
    enum {
        Dimensions = 3,
    };
    char x;
    char y;
    char z;

    bool operator ==(const CharV3 &other) const {
        return (x == other.x) && (y == other.y) && (z == other.z);
    }

    int operator >(const CharV3 &other) const {
        return ((x > other.x)?1:0) | ((y > other.y)?1:0)<<1 | ((z > other.z)?1:0)<<2;
    }
};

inline void test_ntree() {
    NTree<int, CharV3> tree;

    CharV3 c;
    c.x = 0; c.y = 1; c.z = 25;
    tree.insert(c, 10);
    c.x = 2; c.y = 0; c.z = 24;
    tree.insert(c, 11);
    c.x = -5; c.y = -1; c.z = 100;
    tree.insert(c, 12);

    c.x = 0; c.y = 1; c.z = 25;
    assert(*tree.get(c) == 10);
    c.x = 2; c.y = 0; c.z = 24;
    assert(*tree.get(c) == 11);
    c.x = -5; c.y = -1; c.z = 100;
    assert(*tree.get(c) == 12);
    c.x = -6; c.y = -1; c.z = 100;
    assert(tree.get(c) == NULL);


}

} // namespace test

} // namespace containers