/**
 * @file
 * @brief Persistent(Purely functional) Red-Black Tree
 * @author AKIYAMA Kouhei
 * @since 2011-05-18
 */
#ifndef PRBTREE_H
#define PRBTREE_H

#include <vector>
#include <boost/intrusive_ptr.hpp>
#include <boost/iterator/iterator_facade.hpp>


namespace prbtree
{
    // see: http://www.sampou.org/cgi-bin/haskell.cgi?Programming%3a%b6%cc%bc%ea%c8%a2%3a%a4%bd%a4%ce%c2%be#H-179mxtt
    // see: google: Okasaki-style red-black tree

    // example:
    //   Node<int>::Ptr root;
    //   for(int i = 0; i < 10000; ++i){
    //      root = insert(root, std::rand());
    //   }
    //   Node<int>::Ptr seven = lookup(root, 7);

    // example:
    //   void f(const KeyValue<int, int> &kv) {
    //       std::cout << '(' << kv.key << ',' << kv.value << ')';
    //   }
    //   Node<KeyValue<int, int> >::Ptr root;
    //   root = insert_multi(root, KeyValue<int, int>(5, 1));
    //   root = insert_multi(root, KeyValue<int, int>(10, 1));
    //   root = insert_multi(root, KeyValue<int, int>(10, 2));
    //   root = insert_multi(root, KeyValue<int, int>(28, 1));
    //   root = insert_multi(root, KeyValue<int, int>(10, 3));
    //   root = insert_multi(root, KeyValue<int, int>(10, 4));
    //   root = insert_multi(root, KeyValue<int, int>(10, 5));
    //   root = insert_multi(root, KeyValue<int, int>(28, 2));
    //   std::for_each(begin(root), end(root), f); std::cout << std::endl;
    //   std::for_each(lower_bound(root, 5), upper_bound(root, 5), f); std::cout << std::endl;
    //   std::for_each(lower_bound(root, 10), upper_bound(root, 10), f); std::cout << std::endl;
    //   std::for_each(lower_bound(root, 28), upper_bound(root, 28), f); std::cout << std::endl;
    // output:
    //   (5,1)(10,1)(10,2)(10,3)(10,4)(10,5)(10,6)(10,7)(28,1)(28,2)
    //   (5,1)
    //   (10,1)(10,2)(10,3)(10,4)(10,5)(10,6)(10,7)
    //   (28,1)(28,2)

    
    template<typename T>
    class Node
    {
        int ref_count;
    public:
        enum Color { RED, BLACK };
        typedef boost::intrusive_ptr<Node> Ptr;
        typedef T Value;

        const Color color;
        const Ptr l; // left
        const Ptr r; // right
        const Value v; // value

        Node(Color color, const Ptr &left, const T &value, const Ptr &right)
                : ref_count(1)
                , color(color), l(left), r(right), v(value)
        {}

        friend void intrusive_ptr_add_ref(Node *p)
        {
            ++(p->ref_count);
        }
        friend void intrusive_ptr_release(Node *p)
        {
            if(--(p->ref_count) == 0){
                delete p;
            }
        }

        // lookup

        template<typename Key>
        friend Ptr lookup(const Ptr &node, const Key &v)
        {
            if(!node){
                return NULL;
            }
            else{
                if(v < node->v){
                    return lookup(node->l, v);
                }
                else if(node->v < v){
                    return lookup(node->r, v);
                }
                else{
                    return node;
                }
            }
        }

        
        // insertion
        
        friend Ptr insert(const Ptr &node, const Value &v)
        {
            return insert_root<false>(node, v);
        }
        
        friend Ptr insert_multi(const Ptr &node, const Value &v)
        {
            return insert_root<true>(node, v);
        }
        

    private:
        template<bool ALLOW_SAME_VALUE>
        static Ptr insert_root(const Ptr &node, const Value &v)
        {
            const Ptr result = insert_non_root<ALLOW_SAME_VALUE>(node, v);
            if(result->color == RED){
                return new_node(BLACK, result->l, result->v, result->r);
            }
            else{
                return result;
            }
        }
        
        template<bool ALLOW_SAME_VALUE>
        static Ptr insert_non_root(const Ptr &node, const Value &v)
        {
            if(!node){
                return new_node(RED, NULL, v, NULL);
            }
            else{
                if(v < node->v){
                    return blance(node->color, insert_non_root<ALLOW_SAME_VALUE>(node->l, v), node->v, node->r);
                }
                else if(ALLOW_SAME_VALUE || node->v < v){
                    return blance(node->color, node->l, node->v, insert_non_root<ALLOW_SAME_VALUE>(node->r, v));
                }
                else{
                    return new_node(node->color, node->l, v, node->r);
                }
            }
        }
        
        static Ptr blance(Color color, const Ptr &l, const Value &v, const Ptr &r)
        {
            //                    v
            //          l                   r
            //    ll        lr        rl        rr
            // lll  llr  lrl  lrr  rll  rlr  rrl  rrr
            if(color == BLACK){
                if(l && l->color == RED){
                    if(l->l && l->l->color == RED ){ //ll
                        return new_node(RED, new_node(BLACK, l->l->l, l->l->v, l->l->r), l->v, new_node(BLACK, l->r, v, r));
                    }
                    else if(l->r && l->r->color == RED ){ //lr
                        return new_node(RED, new_node(BLACK, l->l, l->v, l->r->l), l->r->v, new_node(BLACK, l->r->r, v, r));
                    }
                }
                else if(r && r->color == RED){
                    if(r->l && r->l->color == RED ){ //rl
                        return new_node(RED, new_node(BLACK, l, v, r->l->l), r->l->v, new_node(BLACK, r->l->r, r->v, r->r));
                    }
                    else if(r->r && r->r->color == RED ){ //rr
                        return new_node(RED, new_node(BLACK, l, v, r->l), r->v, new_node(BLACK, r->r->l, r->r->v, r->r->r));
                    }
                }
            }
            return new_node(color, l, v, r);
        }

        static Ptr new_node(Color color, const Ptr &left, const Value &value, const Ptr &right)
        {
            return Ptr(new Node(color, left, value, right), false);
        }
    };

    //
    // struct KeyValue
    //
    
    template<typename Key, typename Value>
    struct KeyValue
    {
        Key key;
        Value value;
        
        KeyValue(const Key &k, const Value &v)
                : key(k), value(v)
        {}

        friend bool operator<(const KeyValue &lhs, const KeyValue &rhs) {return lhs.key < rhs.key;}
        friend bool operator<(const Key &lhs, const KeyValue &rhs) {return lhs < rhs.key;}
        friend bool operator<(const KeyValue &lhs, const Key &rhs) {return lhs.key < rhs;}
    };


    //
    // iterator
    //
    template<typename T>
    class Iterator 
        : public boost::iterator_facade<Iterator<T> , const T, boost::forward_traversal_tag>
    {
    public:
        typedef typename Node<T>::Ptr Ptr;
        std::vector<Ptr> nodes;

        Ptr to_ptr() const
        {
            return nodes.empty() ? Ptr() : nodes.back();
        }
        
    private:
        friend class boost::iterator_core_access;
        
        void increment()
        {
            if(nodes.empty()){
            }
            else if(Ptr right = nodes.back()->r){
                Ptr next = right;
                do{
                    nodes.push_back(next);
                    next = next->l;
                }
                while(next);
            }
            else{
                for(; nodes.size() >= 2 && (*(nodes.end()-2))->l != nodes.back(); nodes.pop_back()) ;
                nodes.pop_back();
            }
        }

        const T &dereference() const
        {
            assert(!nodes.empty());
            return nodes.back()->v;
        }

        bool equal(const Iterator &rhs) const
        {
            return (nodes.empty() && rhs.nodes.empty()) ||
                    (!nodes.empty() && !rhs.nodes.empty() && nodes.back() == rhs.nodes.back());
        }
        
    };

    template<typename T>
    Iterator<T> begin(const boost::intrusive_ptr<Node<T> > &node)
    {
        boost::intrusive_ptr<Node<T> > np = node;
        Iterator<T> it;
        while(np){
            it.nodes.push_back(np);
            np = np->l;
        }
        return it;
    }

    template<typename T>
    Iterator<T> last(const boost::intrusive_ptr<Node<T> > &node)
    {
        boost::intrusive_ptr<Node<T> > np = node;
        Iterator<T> it;
        while(np){
            it.nodes.push_back(np);
            np = np->r;
        }
        return it;
    }

    template<typename T>
    Iterator<T> end(const boost::intrusive_ptr<Node<T> > &node)
    {
        return Iterator<T>();
    }
    
    
    template<typename T, typename Key>
    Iterator<T> lower_bound(const boost::intrusive_ptr<Node<T> > &node, const Key &v)
    {
        boost::intrusive_ptr<Node<T> > np = node;
        Iterator<T> it;
        std::size_t last = 0;
        
        while(np){
            it.nodes.push_back(np);
            if(np->v < v){
                np = np->r;
            }
            else{
                last = it.nodes.size();
                np = np->l;
            }
        }

        if(it.nodes.size() > last){
            it.nodes.erase(it.nodes.begin() + last, it.nodes.end());
        }
        return it;
    }

    template<typename T, typename Key>
    Iterator<T> upper_bound(const boost::intrusive_ptr<Node<T> > &node, const Key &v)
    {
        boost::intrusive_ptr<Node<T> > np = node;
        Iterator<T> it;
        std::size_t last = 0;
        
        while(np){
            it.nodes.push_back(np);
            if(v < np->v){
                last = it.nodes.size();
                np = np->l;
            }
            else{
                np = np->r;
            }
        }

        if(it.nodes.size() > last){
            it.nodes.erase(it.nodes.begin() + last, it.nodes.end());
        }
        return it;
    }



    //
    // equal range
    //

    template<typename Func, typename T>
    void for_each_equal_range(const boost::intrusive_ptr<Node<T> > &node, Func func)
    {
        if(!node){
            return;
        }
        std::for_each<Iterator<T>, Func>(
                lower_bound(node, node->v),
                upper_bound(node, node->v),
                func);
    }
    
    template<typename Pred, typename T>
    boost::intrusive_ptr<Node<T> > find_in_equal_range(const boost::intrusive_ptr<Node<T> > &node, Pred func)
    {
        if(!node){
            return NULL;
        }
        const Iterator<T> it = std::find_if<Iterator<T>, Pred>(
                lower_bound(node, node->v),
                upper_bound(node, node->v),
                func);
        return it.to_ptr();
    }
    
}
#endif