#pragma once

/*
    Author: fractalgong
    Date: 2013-04-07
    a rbtree based rank tree, interprocess traits is supported.
*/

//#include <boost/pointer_to_other.hpp>
#include <boost/interprocess/offset_ptr.hpp>
#include <utility>

namespace qday
{

template <typename key_t, typename value_t, typename allocator_t>
class rank_tree
{
public:
    typedef std::pair<key_t, value_t>   value_type;
    typedef value_type*                 pointer;
    typedef const value_type*           const_pointer;
    typedef value_type&                 reference;
    typedef const value_type&           const_reference;

    typedef size_t rank_t;

private:

    template <typename other_pointer_t>
    struct node
    {
        typedef typename boost::pointer_to_other<other_pointer_t, node>::type pointer_t;
        static const int red = 0;
        static const int black = 1;
        node():parent(NULL), size(1), color(red) 
        {
            children[0] = NULL;
            children[1] = NULL;
        }
        pointer_t children[2];
        pointer_t parent;
        rank_t size;
        int color;
        value_type value;
        void construct()
        {
            new(this) node;
        }

        static bool is_red(pointer_t n)
        {
            return n && n->color == red;
        }

        static int dir(pointer_t p, pointer_t c)
        {
            return c == p->children[1];
        }

        static void update_size(pointer_t h)
        {
            rank_t size = 1;
            if (h->children[0])
            {
                size += h->children[0]->size;
            }
            if (h->children[1])
            {
                size += h->children[1]->size;
            }
            h->size = size;
        }

        static void update_size_to_root(pointer_t h)
        {
            while (h)
            {
                node::update_size(h);
                h = h->parent;
            }
        }

        static pointer_t rotate(pointer_t h, int dir)
        {
            int opp_dir = !dir;
            pointer_t x = h->children[opp_dir];
            h->children[opp_dir] = x->children[dir];
            x->children[dir] = h;
            std::swap(x->color, h->color);
            pointer_t parent = h->parent;
            if (parent)
            {
                int pdir = node::dir(parent, h);
                parent->children[pdir] = x;
            }
            x->parent = parent;
            h->parent = x;
            if (h->children[opp_dir])
            {
                h->children[opp_dir]->parent = h;
            }

            node::update_size(h);
            node::update_size(x);
            return x;
        }
        static pointer_t rotate_left(pointer_t h)
        {
            return node::rotate(h, 0);
        }

        static pointer_t rotate_right(pointer_t h)
        {
            return node::rotate(h, 1);
        }

        static void color_flip(pointer_t h)
        {
            h->color = !h->color;
            h->children[0]->color = !h->children[0]->color;
            h->children[1]->color = !h->children[1]->color;
        }

        static pointer_t spine(pointer_t h, int dir)
        {
            if (!h)
            {
                return h;
            }
            while (h->children[dir])
            {
                h = h->children[dir];
            }
            return h;
        }

        static pointer_t move_pointer(pointer_t h, int dir)
        {
            int opp_dir = !dir;
            if (!h)
            {
                return h;
            }
            else if (h->children[dir])
            {
                return node::spine(h->children[dir], opp_dir);
            }
            else
            {
                while (h->parent)
                {
                    if (node::dir(h->parent, h) == opp_dir)
                    {
                        return h->parent;
                    }
                    h = h->parent;
                }
                return NULL;
            }
        }
    };

    typedef node<typename allocator_t::pointer> node_t;
    typedef typename allocator_t::template rebind<node_t>::other node_allocator_t;
    typedef typename node_allocator_t::pointer node_pointer_t;
    typedef typename node_allocator_t::const_pointer const_node_pointer_t;
    
    static node_pointer_t minimum(node_pointer_t h)
    {
        return node_t::spine(h, 0);
    }

    static node_pointer_t maximum(node_pointer_t h)
    {
        return node_t::spine(h, 1);
    }

    static node_pointer_t prev(node_pointer_t h)
    {
        return node_t::move_pointer(h, 0);
    }

    static node_pointer_t next(node_pointer_t h)
    {
        return node_t::move_pointer(h, 1);
    }

public:

    class iterator
    {
    public:
        iterator(node_pointer_t p):m_p(p){}
        iterator():m_p(NULL){}
        iterator operator--()
        {
            m_p = prev(m_p);
            return *this;
        }
        iterator operator++()
        {
            m_p = next(m_p);
            return *this;
        }
        iterator operator--(int)
        {
            iterator tmp = *this;
            --*this;
            return tmp;
        }
        iterator operator++(int)
        {
            iterator tmp = *this;
            ++*this;
            return tmp;
        }
        reference operator*()
        {
            return m_p->value;
        }
        pointer operator->()
        {
            return pointer(&(m_p->value));
        }
        bool operator==(iterator other) const
        {
            return m_p == other.m_p;
        }
        bool operator !=(iterator other) const
        {
            return !(*this == other);
        }
    private:
        node_pointer_t m_p;
        friend class rank_tree;
    };

    rank_tree(const allocator_t& a):m_root(NULL), m_alloc((node_allocator_t)a)
    {

    }

    rank_t rank(iterator i) const
    {
        return __rank(i.m_p);
    }
    
    iterator find_by_rank(rank_t r) const
    {
        node_pointer_t p = __find_by_rank(r);
        return iterator(p);
    }

    iterator insert(const_reference v)
    {
        node_pointer_t p = __insert(v, std::less<key_t>());
        return iterator(p);
    }

    template <typename comp_t>
    iterator insert(const_reference v, comp_t comp)
    {
        node_pointer_t p = __insert(v, comp);
        return iterator(p);
    }

    void erase(iterator i)
    {
        return __erase(i.m_p);
    }

    template <typename comp_t>
    iterator lower_bound(const key_t& k, comp_t comp)
    {
        return iterator(__lower_bound(k, comp));
    }

    template <typename comp_t>
    iterator upper_bound(const key_t& k, comp_t comp)
    {
        return iterator(__upper_bound(k, comp));
    }

    iterator lower_bound(const key_t& k)
    {
        return iterator(__lower_bound(k, std::less<key_t>()));
    }

    iterator upper_bound(const key_t& k)
    {
        return iterator(__upper_bound(k, std::less<key_t>()));
    }

    rank_t size() const
    {
        return m_root ? m_root->size : 0;
    }

    // TODO cache iterator begin.
    iterator begin()
    {
        return iterator(minimum(m_root));
    }

    iterator last()
    {
        return iterator(maximum(m_root));
    }

    iterator end()
    {
        return iterator();
    }

private:

    template <typename comp_t>
    node_pointer_t __lower_bound(const key_t& k, comp_t comp)
    {   
        // invariant left < key <= right
        // left <= root <= right
        node_pointer_t n = m_root;
        node_pointer_t p = NULL;
        while (n) 
        {   
            if (comp(n->value.first, k)) 
                // k > root, so root as left, turn right 
                n = n->children[1];
            else
            {   
                // k <= root, so root as right, turn left, and cache root
                p = n;
                n = n->children[0]; 
            }   
        }   
        return p;
    }   

    template <typename comp_t>
    node_pointer_t __upper_bound(const key_t& k, comp_t comp)
    {   
        // invariant left <= key < right
        node_pointer_t n = m_root;
        node_pointer_t p = NULL;
        while (n) 
        {   
            if (comp(k, n->value.first))
            {   
                p = n;
                n = n->children[0];
            }   
            else
                n = n->children[1]; 
        }   
        return p;
    }  

    rank_t __rank(node_pointer_t n) const
    {
        // the reverse process of __find_by_rank.
        node_pointer_t left = n->children[0];
        rank_t r = left ? left->size : 0;
        node_pointer_t p = n->parent;
        while (p)
        {
            if (node_t::dir(p, n) == 1)
            {
                r++;    // parent in the left.
                node_pointer_t left = p->children[0];
                if (left)
                    r += left->size;    // left children of parent in the left.
            }
            n = p;
            p = n->parent; 
        }
        return r;
    }

    node_pointer_t __find_by_rank(rank_t r) const
    {
        // the reverse process of __rank.
        if (r >= size())
        {
            return NULL;
        }
        node_pointer_t n = m_root;

        while (n)
        {
            node_pointer_t left = n->children[0];
            rank_t left_size = left ? left->size : 0;
            if (r < left_size)
            {
                n = left;
            }
            else if (r == left_size)
            {
                return n;
            }
            else
            {
                r -= left_size + 1;
                n = n->children[1];
            }
        }
        return NULL;
    }

    template <typename comp_t>
    node_pointer_t __insert(const_reference v, comp_t comp)
    {
        node_pointer_t n = __allocate_node();
        n->value = v;
        if (!m_root)
        {
            m_root = n;
            m_root->color = node_t::black;
            return n;
        }
        node_pointer_t h = m_root;
        node_pointer_t p = NULL;
        while (h)
        {
            int dir = !comp(v.first, h->value.first);
            p = h;
            h = h->children[dir];
        }
        int dir = !comp(v.first, p->value.first);
        p->children[dir] = n;
        n->parent = p;
        node_pointer_t saved_n = n;
        while (node_t::is_red(p))
        {
            node_t::update_size(p);
            node_pointer_t gp = p->parent;
            node_t::update_size(gp);
            int pdir = node_t::dir(gp, p);
            node_pointer_t u = gp->children[!pdir];
            if (!node_t::is_red(u))
            {
                if (node_t::dir(p, n) != pdir)
                {
                    p = node_t::rotate(p, pdir);
                }
                gp = node_t::rotate(gp, !pdir);
                if (!gp->parent)
                {
                    m_root = gp;
                }
                break;
            }
            node_t::color_flip(gp);
            n = gp;
            p = n->parent;
        }
        m_root->color = node_t::black;
        node_t::update_size_to_root(p);
        return saved_n;
    }

    void __erase(node_pointer_t n)
    {
        assert(m_root);
        if (n->children[0] && n->children[1])
        {
            // swap all things except value.
            node_pointer_t successor = next(n);
            node_pointer_t c = NULL;
            if (successor == n->children[1])
            {
                n->children[1] = successor->children[1];    // n child1 set
                c = n->children[1];
                if (c)
                {
                    c->parent = n;                          // n child1 fix
                }
                successor->children[1] = n;                 // successor child1 set && n parent fix.
                successor->parent = n->parent;              // successor parent set
                n->parent = successor;                      // n parent set && successor child1 fix.
                if (successor->parent)
                {
                    successor->parent->children[node_t::dir(successor->parent, n)] = successor;// successor parent fix.
                }
                // n child1 done, n parent done, successor child1 done, successor parent done.
            }
            else
            {
                std::swap(n->parent, successor->parent);    // n parent set && successor parent set.
                if (n->parent)
                {
                    n->parent->children[node_t::dir(n->parent, successor)] = n; // n parent fix.
                }
                if (successor->parent)
                {
                    successor->parent->children[node_t::dir(successor->parent, n)] = successor;// successor parent fix.
                }
                std::swap(n->children[1], successor->children[1]);// n child1 set && successor set.
                c = n->children[1];
                if (c)
                {
                    c->parent = n;                               // n child1 fix.
                }
                c = successor->children[1];
                if (c)
                {
                    c->parent = successor;                      // successor child1 fix.
                }
                // n child1 done, n parent done, successor child1 done, successor parent done.
            }
            std::swap(n->children[0], successor->children[0]);  // n child0 set && successor child0 set.
            c = n->children[0];
            if (c)
            {
                c->parent = n;                                  // n child0 fix.
            }
            c = successor->children[0];
            if (c)
            {
                c->parent = successor;                          // successor fix.
            }
            std::swap(n->color, successor->color);
            std::swap(n->size, successor->size);
            if (!successor->parent)
            {
                m_root = successor;
            }
        }
        node_pointer_t c = NULL;
        node_pointer_t p = n->parent;
        if (n->children[0])
        {
            c = n->children[0];
        }
        else if (n->children[1])
        {
            c = n->children[1];
        }
        node_pointer_t saved_n = n;
        if (c)
        {
            c->parent = p;
            if (node_t::is_red(c))
            {
                c->color = node_t::black;
                goto finished;
            }
        }

        if (!p || node_t::is_red(n))
        {
            goto finished;
        }
        
        n->size = 0;
        for (;!node_t::is_red(n) && p; n = p, p = n->parent)
        {
            node_t::update_size(p);
            int dir = node_t::dir(p, n);
            node_pointer_t sibling = p->children[!dir];
            if (node_t::is_red(sibling))
            {
                p = node_t::rotate(p, dir);
                if (!p->parent)
                {
                    m_root = p;
                }
                p = n->parent;
                dir = node_t::dir(p, n);
                sibling = p->children[!dir];
            }
            if(node_t::is_red(sibling->children[dir]) && !node_t::is_red(sibling->children[!dir]))
            {
                sibling = node_t::rotate(sibling, !dir);
            }
            if (node_t::is_red(sibling->children[!dir]))
            {
                sibling->children[!dir]->color = node_t::black;
                p = node_t::rotate(p, dir);
                if (!p->parent)
                {
                    m_root = p;
                }
                break;
            }
            sibling->color = node_t::red;
        }
        if (node_t::is_red(n))
        {
            n->color = node_t::black;
        }
finished:
        node_pointer_t saved_p = saved_n->parent;
        if (saved_p)
        {
            saved_p->children[node_t::dir(saved_p, saved_n)] = c;
        }
        else
        {
            m_root = c;
        }
        if (m_root)
        {
            m_root->color = node_t::black;
        }
        node_t::update_size_to_root(p);
        __deallocate_node(saved_n);
    }

#ifdef DEBUG
    int __height(node_pointer_t n)
    {
        if (!n)
        {
            return 0;
        }
        int h = n->color;
        int lh = __height(n->children[0]);
        int rh = __height(n->children[1]);
        if (!node_t::is_red(n))
        {
            if (lh != rh)
            {
                __print_r(n, 0);
            }
            assert(lh == rh);
            return h + lh;
        }
        else
        {
            return std::max(lh, rh);
        }
    }

    void __print_r(const_node_pointer_t h, int indent)
    {
        if (!h)
        {
            std::cout << std::endl;
            return;
        }
        std::cout << std::string(indent, '\t') << h->value.first << (h->color == node_t::red ? 'r': 'b') << h->size << std::endl;
        __print_r(h->children[0], indent + 1);
        __print_r(h->children[1], indent + 1);
    }

    void __print()
    {
        __print_r(m_root, 0);
    }

#endif

    node_pointer_t __root(){return m_root;}

    node_pointer_t __allocate_node()
    {
        node_pointer_t n = m_alloc.allocate_one();
        n->construct();
        return n;
    }
    void __deallocate_node(node_pointer_t n)
    {
        n->~node_t();
        m_alloc.deallocate_one(n);
    }

    node_pointer_t m_root;
    node_allocator_t m_alloc;
};

}