//
// Urho3D Engine
// Copyright (c) 2008-2011 Lasse Öörni
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#pragma once

#include "TreeBase.h"

// Based on Red Black Trees by Julienne Walker
// http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx

/// %Set template class using a red-black tree.
template <class T> class Set : public TreeBase
{
public:
    /// %Set node.
    struct Node : public TreeNodeBase
    {
        // Construct undefined.
        Node()
        {
        }
        
        // Construct with key.
        Node(const T& key) :
            key_(key)
        {
        }
        
        /// Key.
        T key_;
        
        /// Return parent node.
        Node* Parent() const { return static_cast<Node*>(parent_); }
        /// Return the left or right child.
        Node* Child(unsigned dir) const { return static_cast<Node*>(link_[dir]); }
    };
    
    /// %Set iterator.
    class Iterator : public TreeIteratorBase
    {
    public:
        /// Construct.
        Iterator()
        {
        }
        
        /// Construct with a node pointer.
        Iterator(Node* ptr) :
            TreeIteratorBase(ptr)
        {
        }
        
        /// Preincrement the pointer.
        Iterator& operator ++ () { GotoNext(); return *this; }
        /// Postincrement the pointer.
        Iterator operator ++ (int) { Iterator it = *this; GotoNext(); return it; }
        /// Predecrement the pointer.
        Iterator& operator -- () { GotoPrev(); return *this; }
        /// Postdecrement the pointer.
        Iterator operator -- (int) { Iterator it = *this; GotoPrev(); return it; }
        
        /// Point to the key.
        const T* operator -> () const { return &(static_cast<Node*>(ptr_))->key_; }
        /// Dereference the key.
        const T& operator * () const { return (static_cast<Node*>(ptr_))->key_; }
    };
    
    /// %Set const iterator.
    class ConstIterator : public TreeIteratorBase
    {
    public:
        /// Construct.
        ConstIterator()
        {
        }
        
        /// Construct with a node pointer.
        ConstIterator(Node* ptr) :
            TreeIteratorBase(ptr)
        {
        }
        
        // Construct from a non-const iterator.
        ConstIterator(const Iterator& it) :
            TreeIteratorBase(it.ptr_)
        {
        }
        
        /// Assign from a non-const iterator.
        ConstIterator& operator = (const Iterator& rhs) { ptr_ = rhs.ptr_; return *this; }
        /// Preincrement the pointer.
        ConstIterator& operator ++ () { GotoNext(); return *this; }
        /// Postincrement the pointer.
        ConstIterator operator ++ (int) { ConstIterator it = *this; GotoNext(); return it; }
        /// Predecrement the pointer.
        ConstIterator& operator -- () { GotoPrev(); return *this; }
        /// Postdecrement the pointer.
        ConstIterator operator -- (int) { ConstIterator it = *this; GotoPrev(); return it; }
        
        /// Point to the key.
        const T* operator -> () const { return &(static_cast<Node*>(ptr_))->key_; }
        /// Dereference the key.
        const T& operator * () const { return (static_cast<Node*>(ptr_))->key_; }
    };
    
    /// Construct empty.
    Set()
    {
    }
    
    /// Construct from another set.
    Set(const Set<T>& set)
    {
        allocator_ = AllocatorInitialize(sizeof(Node), set.Size());
        *this = set;
    }
    
    /// Destruct.
    ~Set()
    {
        Clear();
        AllocatorUninitialize(allocator_);
    }
    
    /// Assign a set.
    Set& operator = (const Set<T>& set)
    {
        Clear();
        Insert(set);
        
        return *this;
    }
    
    /// Add-assign a value.
    Set& operator += (const T& rhs)
    {
        Insert(rhs);
        return *this;
    }
    
    /// Add-assign a set.
    Set& operator += (const Set<T>& rhs)
    {
        Insert(rhs);
        return *this;
    }
    
    /// Test for equality with another set.
    bool operator == (const Set<T>& rhs) const
    {
        if (rhs.size_ != size_)
            return false;
        
        ConstIterator i = Begin();
        ConstIterator j = rhs.Begin();
        while (i != End())
        {
            if (*i != *j)
                return false;
            ++i;
            ++j;
        }
        
        return true;
    }
    
    /// Test for inequality with another set.
    bool operator != (const Set<T>& rhs) const
    {
        if (rhs.size_ != size_)
            return true;
        
        ConstIterator i = Begin();
        ConstIterator j = rhs.Begin();
        while (i != End())
        {
            if (*i != *j)
                return true;
            ++i;
            ++j;
        }
        
        return false;
    }
    
    /// Clear the set.
    void Clear()
    {
        Node* root = Root();
        if (!root)
            return;
        
        EraseNodes(root);
        root_ = 0;
    }
    
    /// Insert a key. Return iterator to the value.
    Iterator Insert(const T& key) { return Iterator(InsertNode(key)); }
    /// Insert a set.
    void Insert(const Set<T>& set) { Insert(set.Begin(), set.End()); }
    /// Insert a key by iterator. Return iterator to the value.
    Iterator Insert(const ConstIterator& it) { return Iterator(InsertNode(*it)); }
    
    /// Insert a range by iterators.
    void Insert(const ConstIterator& start, const ConstIterator& end)
    {
        ConstIterator it = start;
        while (it != end)
            InsertNode(*it++);
    }
    
    /// Erase a key. Return true if was found.
    bool Erase(const T& key) { return EraseNode(key); }
    /// Erase a key by iterator.
    void Erase(const Iterator& it) { EraseNode(*it); }
    
    /// Erase a range by iterators.
    void Erase(const Iterator& start, const Iterator& end)
    {
        Iterator it = start;
        while (it != end)
        {
            Iterator current = it++;
            Erase(current);
        }
    }
    
    /// Return whether contains a key.
    bool Contains(const T& key) const { return FindNode(key) != 0; }
    /// Return iterator to the node with key, or end iterator if not found.
    Iterator Find(const T& key) { Node* node = FindNode(key); return node ? Iterator(node) : End(); }
    /// Return const iterator to the node with key, or null iterator if not found.
    ConstIterator Find(const T& key) const { Node* node = FindNode(key); return node ? ConstIterator(node) : End(); }
    /// Return iterator to the beginning.
    Iterator Begin() { return Iterator(FindFirst()); }
    /// Return const iterator to the beginning.
    ConstIterator Begin() const { return ConstIterator(FindFirst()); }
    /// Return iterator to the end.
    Iterator End() { return ++Iterator(FindLast()); }
    /// Return const iterator to the end.
    ConstIterator End() const { return ++ConstIterator(FindLast()); }
    /// Return first key.
    const T& Front() { return FindFirst()->key_; }
    /// Return last key.
    const T& Back() { return FindLast()->key_; }
    /// Return number of keys.
    unsigned Size() const { return size_; }
    /// Return whether set is empty.
    bool Empty() const { return size_ == 0; }
    
private:
    /// Return the root pointer with correct type.
    Node* Root() const { return reinterpret_cast<Node*>(root_); }
    
    /// Find the node with smallest key.
    Node* FindFirst() const
    {
        Node* node = Root();
        while (node && node->link_[0])
            node = node->Child(0);
        return node;
    }
    
    /// Find the node with largest key.
    Node* FindLast() const
    {
        Node* node = Root();
        while (node && node->link_[1])
            node = node->Child(1);
        return node;
    }
    
    /// Find a node with key. Return null if not found.
    Node* FindNode(const T& key) const
    {
        Node* node = Root();
        while (node)
        {
            if (node->key_ == key)
                return node;
            else
                node = node->Child(node->key_ < key);
        }
        return 0;
    }
    
    /// Insert a node and return a pointer to it.
    Node* InsertNode(const T& key)
    {
        Node* ret = 0;
        
        if (!root_)
        {
            root_ = ret = ReserveNode(key);
            ++size_;
        }
        else
        {
            Node head;
            Node* g, * t, * p, * q;
            
            unsigned dir = 0;
            unsigned last = 0;
            
            t = &head;
            g = p = 0;
            q = Root();
            t->SetChild(1, Root());
            
            for (;;)
            {
                if (!q)
                {
                    p->SetChild(dir, q = ret = ReserveNode(key));
                    ++size_;
                }
                else if (IsRed(q->link_[0]) && IsRed(q->link_[1]))
                {
                    q->isRed_ = true;
                    q->link_[0]->isRed_ = false;
                    q->link_[1]->isRed_ = false;
                }
                
                if (IsRed(q) && IsRed(p))
                {
                    unsigned dir2 = (t->link_[1] == g);
                    if (q == p->link_[last])
                        t->SetChild(dir2, RotateSingle(g, !last));
                    else
                        t->SetChild(dir2, RotateDouble(g, !last));
                }
                
                if (q->key_ == key)
                {
                    ret = q;
                    break;
                }
                
                last = dir;
                dir = q->key_ < key;
                
                if (g)
                    t = g;
                g = p;
                p = q;
                q = q->Child(dir);
            }
            
            root_ = head.Child(1);
        }
        
        root_->isRed_ = false;
        root_->parent_ = 0;
        
        return ret;
    }
    
    /// Erase a node. Return true if was erased.
    bool EraseNode(const T& key)
    {
        if (!root_)
            return false;
        
        Node head;
        Node* q, * p, *g;
        Node* f = 0;
        unsigned dir = 1;
        bool removed = false;
        
        q = &head;
        g = p = 0;
        q->SetChild(1, Root());
        
        while (q->link_[dir])
        {
            unsigned last = dir;
            g = p;
            p = q;
            q = q->Child(dir);
            dir = q->key_ < key;
            
            if (q->key_ == key)
                f = q;
             
            if (!IsRed(q) && !IsRed(q->link_[dir]))
            {
                if (IsRed(q->link_[!dir]))
                {
                    p->SetChild(last, RotateSingle(q, dir));
                    p = p->Child(last);
                }
                else if (!IsRed(q->link_[!dir]))
                {
                    Node* s = p->Child(!last);
                    
                    if (s)
                    {
                        if (!IsRed(s->link_[!last]) && !IsRed(s->link_[last]))
                        {
                            p->isRed_ = false;
                            s->isRed_ = true;
                            q->isRed_ = true;
                        }
                        else
                        {
                            int dir2 = (g->link_[1] == p);
                            if (IsRed(s->link_[last]))
                                g->SetChild(dir2, RotateDouble(p, last));
                            else if (IsRed(s->link_[!last]))
                                g->SetChild(dir2, RotateSingle(p, last));
                            
                            Node* t = g->Child(dir2);
                            q->isRed_ = t->isRed_ = true;
                            t->link_[0]->isRed_ = false;
                            t->link_[1]->isRed_ = false;
                        }
                    }
                }
            }
        }
        
        if (f)
        {
            f->key_ = q->key_;
            p->SetChild(p->link_[1] == q, q->link_[q->link_[0] == 0]);
            FreeNode(q);
            --size_;
            removed = true;
        }
        
        root_ = head.Child(1);
        if (root_)
        {
            root_->isRed_ = false;
            root_->parent_ = 0;
        }
        
        return removed;
    }
    
    /// Erase the nodes recursively.
    void EraseNodes(Node* node)
    {
        Node* left = node->Child(0);
        Node* right = node->Child(1);
        FreeNode(node);
        --size_;
        
        if (left)
            EraseNodes(left);
        if (right)
            EraseNodes(right);
    }
    
    /// Reserve a node.
    Node* ReserveNode()
    {
        if (!allocator_)
            allocator_ = AllocatorInitialize(sizeof(Node));
        Node* newNode = static_cast<Node*>(AllocatorReserve(allocator_));
        new(newNode) Node();
        
        return newNode;
    }
    
    /// Reserve a node with specified key.
    Node* ReserveNode(const T& key)
    {
        if (!allocator_)
            allocator_ = AllocatorInitialize(sizeof(Node));
        Node* newNode = static_cast<Node*>(AllocatorReserve(allocator_));
        new(newNode) Node(key);
        return newNode;
    }
    
    /// Free a node.
    void FreeNode(Node* node)
    {
        (node)->~Node();
        AllocatorFree(allocator_, node);
    }
};