/* Copyright (c) <2003-2011> <Julio Jerez, Newton Game Dynamics>
*
* This software is provided 'as-is', without any express or implied
* warranty. In no event will the authors be held liable for any damages
* arising from the use of this software.
*
* Permission is granted to anyone to use this software for any purpose,
* including commercial applications, and to alter it and redistribute it
* freely, subject to the following restrictions:
*
* 1. The origin of this software must not be misrepresented; you must not
* claim that you wrote the original software. If you use this software
* in a product, an acknowledgment in the product documentation would be
* appreciated but is not required.
*
* 2. Altered source versions must be plainly marked as such, and must not be
* misrepresented as being the original software.
*
* 3. This notice may not be removed or altered from any source distribution.
*/

#ifndef __dgTree__
#define __dgTree__

#include "dgStdafx.h"
#include "dgRef.h"
#include "dgDebug.h"
#include "dgMemory.h"



// Note: this is a low level class for dgTree use only
// unpredictable result will happen if you attempt to manipulate
// any member of this class
class dgRedBackNode {
public:
	enum REDBLACK_COLOR {
		RED = true,
		BLACK = false
	};

	DG_CLASS_ALLOCATOR(allocator)

	dgRedBackNode() {
	}

	virtual ~dgRedBackNode() {
	}

	void RemoveAllLow();
	void RotateLeft(dgRedBackNode **const head);
	void RotateRight(dgRedBackNode **const head);
	void RemoveFixup(dgRedBackNode *const node, dgRedBackNode * * const head);

	dgRedBackNode *GetLeft() const;
	dgRedBackNode *GetRight() const;
	dgRedBackNode *GetParent() const;

	dgRedBackNode(dgRedBackNode *const parent);
	inline void Initdata(dgRedBackNode *const parent);
	inline void SetColor(REDBLACK_COLOR color);
	REDBLACK_COLOR GetColor() const;
	dgUnsigned32 IsInTree() const;
	inline void SetInTreeFlag(dgUnsigned32 flag);

	void RemoveAll();
	const dgRedBackNode *Prev() const;
	const dgRedBackNode *Next() const;
	const dgRedBackNode *Minimum() const;
	const dgRedBackNode *Maximum() const;
	dgRedBackNode *Prev();
	dgRedBackNode *Next();
	dgRedBackNode *Minimum();
	dgRedBackNode *Maximum();
	void Remove(dgRedBackNode **const head);
	void Unlink(dgRedBackNode **const head);
	void InsertFixup(dgRedBackNode **const head);

	dgRedBackNode *m_left;
	dgRedBackNode *m_right;
	dgRedBackNode *m_parent;
	dgUnsigned32  m_color   : 1;
	dgUnsigned32  m_inTree  : 1;
//	dgUnsigned32  m_pad[2];
};

template<class OBJECT, class KEY>
class dgTree {
public:
	class dgTreeNode: public dgRedBackNode {
		dgTreeNode(
		    const OBJECT &info,
		    const KEY &key,
		    dgTreeNode *parentNode)
			: dgRedBackNode(parentNode), m_info(info), m_key(key) {
//			NEWTON_ASSERT ((dgUnsigned64 (&m_info) & 0x0f) == 0);
		}

		~dgTreeNode() {
		}

		dgTreeNode *GetLeft() const {
			return (dgTreeNode *)dgRedBackNode::m_left;
		}

		dgTreeNode *GetRight() const {
			return (dgTreeNode *)dgRedBackNode::m_right;
		}

		dgTreeNode *GetParent() {
			return (dgTreeNode *)dgRedBackNode::m_parent;
		}

		void SetLeft(dgTreeNode *const node) {
			dgRedBackNode::m_left = node;
		}

		void SetRight(dgTreeNode *const node) {
			dgRedBackNode::m_right = node;
		}

		void SetParent(dgTreeNode *const node) {
			dgRedBackNode::m_parent = node;
		}

	public:
		const KEY &GetKey() const {
			return m_key;
		}

		OBJECT &GetInfo() {
			return m_info;
		}

		const OBJECT &GetInfo() const {
			return m_info;
		}

	private:
		OBJECT m_info;
		KEY m_key;
		friend class dgTree<OBJECT, KEY>;

	};

	class Iterator {

	public:
		Iterator(const dgTree<OBJECT, KEY> &me) {
			m_ptr = NULL;
			m_tree = &me;
		}

		~Iterator() {
		}

		void Begin() {
			m_ptr = m_tree->Minimum();
		}

		void End() {
			m_ptr = m_tree->Maximum();
		}

		void Set(dgTreeNode *const node) {
			m_ptr = node;
		}

		operator dgInt32() const {
			return m_ptr != NULL;
		}

		void operator++ () {
			NEWTON_ASSERT(m_ptr);
			m_ptr = m_ptr->Next();
		}

		void operator++ (int) {
			NEWTON_ASSERT(m_ptr);
			m_ptr = m_ptr->Next();
		}

		void operator-- () {
			NEWTON_ASSERT(m_ptr);
			m_ptr = m_ptr->Prev();
		}

		void operator-- (int) {
			NEWTON_ASSERT(m_ptr);
			m_ptr = m_ptr->Prev();
		}

		OBJECT &operator* () const {
			return ((dgTreeNode *)m_ptr)->GetInfo();
		}

		dgTreeNode *GetNode() const {
			return (dgTreeNode *)m_ptr;
		}

		KEY GetKey() const {
			dgTreeNode *const tmp = (dgTreeNode *)m_ptr;
			return tmp ? tmp->GetKey() : KEY(0);
		}

	private:
		dgRedBackNode *m_ptr;
		const dgTree *m_tree;

	};


	// ***********************************************************
	// member functions
	// ***********************************************************
public:
	DG_CLASS_ALLOCATOR(allocator)

//	dgTree ();
	dgTree(dgMemoryAllocator *const allocator);
	virtual ~dgTree();

	dgMemoryAllocator *GetAllocator() const;
	void SetAllocator(dgMemoryAllocator *const allocator);


	operator dgInt32() const;
	dgInt32 GetCount() const;

	dgTreeNode *GetRoot() const;
	dgTreeNode *Minimum() const;
	dgTreeNode *Maximum() const;

	dgTreeNode *Find(KEY key) const;
	dgTreeNode *FindGreater(KEY key) const;
	dgTreeNode *FindGreaterEqual(KEY key) const;
	dgTreeNode *FindLessEqual(KEY key) const;

	dgTreeNode *GetNodeFromInfo(OBJECT &info) const;

	dgTreeNode *Insert(const OBJECT &element, KEY key, bool &elementWasInTree);
	dgTreeNode *Insert(const OBJECT &element, KEY key);
	dgTreeNode *Insert(dgTreeNode *const node, KEY key);

	dgTreeNode *Replace(OBJECT &element, KEY key);
	dgTreeNode *ReplaceKey(KEY oldKey, KEY newKey);
	dgTreeNode *ReplaceKey(dgTreeNode *const node, KEY key);

	void Remove(KEY key);
	void Remove(dgTreeNode *const node);
	void RemoveAll();

	void Unlink(dgTreeNode *const node);
	void SwapInfo(dgTree &tree);

	bool SanityCheck() const;

	// ***********************************************************
	// member variables
	// ***********************************************************
private:
	dgInt32 m_count;
	dgTreeNode *m_head;
	dgMemoryAllocator *m_allocator;

	dgInt32 CompareKeys(const KEY &key0, const KEY &key1) const;
	bool SanityCheck(dgTreeNode *const ptr, dgInt32 height) const;

	friend class dgTreeNode;
};


inline dgRedBackNode::dgRedBackNode(dgRedBackNode *const parent) {
	Initdata(parent);
}

inline void dgRedBackNode::Initdata(dgRedBackNode *const parent) {
	SetColor(RED);
	SetInTreeFlag(true);
	m_left = NULL;
	m_right = NULL;
	m_parent = parent;
}

inline void dgRedBackNode::SetColor(dgRedBackNode::REDBLACK_COLOR color) {
	m_color = color;
}



inline dgRedBackNode::REDBLACK_COLOR  dgRedBackNode::GetColor() const {
	return REDBLACK_COLOR(m_color);
}


inline void dgRedBackNode::SetInTreeFlag(dgUnsigned32 flag) {
	m_inTree = flag;
}

inline dgUnsigned32 dgRedBackNode::IsInTree() const {
	return m_inTree;
}

/*
template<class OBJECT, class KEY>
dgTree<OBJECT, KEY>::dgTree ()
{
    m_count = 0;
    m_head = NULL;
    m_allocator = NULL;
}
*/

template<class OBJECT, class KEY>
dgTree<OBJECT, KEY>::dgTree(dgMemoryAllocator *const allocator) {
	m_count = 0;
	m_head = NULL;
	m_allocator = allocator;
}


template<class OBJECT, class KEY>
dgTree<OBJECT, KEY>::~dgTree() {
	RemoveAll();
}

template<class OBJECT, class KEY>
dgMemoryAllocator *dgTree<OBJECT, KEY>::GetAllocator() const {
	return m_allocator;
}

template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::SetAllocator(dgMemoryAllocator *const allocator) {
	if ((m_count == 0) && (m_allocator == NULL)) {
		m_allocator = allocator;
	}
}


template<class OBJECT, class KEY>
dgTree<OBJECT, KEY>::operator dgInt32() const {
	return m_head != NULL;
}

template<class OBJECT, class KEY>
dgInt32 dgTree<OBJECT, KEY>::GetCount() const {
	return m_count;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Minimum() const {
	return m_head ? (dgTreeNode *)m_head->Minimum() : NULL;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Maximum() const {
	return m_head ? (dgTreeNode *)m_head->Maximum() : NULL;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::GetRoot() const {
	return m_head;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Find(KEY key) const {
	if (m_head == NULL) {
		return NULL;
	}

	dgTreeNode *ptr = m_head;
	while (ptr != NULL) {
		if (key < ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == -1) ;
			ptr = ptr->GetLeft();
		} else if (key > ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 1) ;
			ptr = ptr->GetRight();
		} else {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 0) ;
			break;
		}
	}
	return ptr;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::GetNodeFromInfo(OBJECT &info) const {
	dgTreeNode *const node = (dgTreeNode *) &info;
	dgInt64 offset = ((char *) &node->m_info) - ((char *) node);
	dgTreeNode *const retnode = (dgTreeNode *)(((char *) node) - offset);

	NEWTON_ASSERT(retnode->IsInTree());
	NEWTON_ASSERT(&retnode->GetInfo() == &info);
	return (retnode->IsInTree()) ? retnode : NULL;

}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::FindGreater(KEY key) const {
	if (m_head == NULL) {
		return NULL;
	}

	dgTreeNode *prev = NULL;
	dgTreeNode *ptr = m_head;

	while (ptr != NULL) {
		if (key < ptr->m_key) {
			prev = ptr;
			ptr = ptr->GetLeft();
		} else {
			ptr = ptr->GetRight();
		}
	}

#ifdef __ENABLE_SANITY_CHECK
	if (prev) {
		Iterator iter(*this);
		for (iter.Begin(); iter.GetNode() != prev; iter ++) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 <= key);
		}
		for (; iter.GetNode(); iter ++) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 > key);
		}
	}
#endif

	return (dgTreeNode *)prev;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::FindGreaterEqual(KEY key) const {
	if (m_head == NULL) {
		return NULL;
	}

	dgTreeNode *prev = NULL;
	dgTreeNode *ptr = m_head;

	while (ptr != NULL) {
		if (key == ptr->m_key) {
			return ptr;
		}
		if (key < ptr->m_key) {
			prev = ptr;
			ptr = ptr->GetLeft();
		} else {
			ptr = ptr->GetRight();
		}
	}

#ifdef __ENABLE_SANITY_CHECK
	if (prev) {
		Iterator iter(*this);
		for (iter.Begin(); iter.GetNode() != prev; iter ++) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 <= key);
		}
		for (; iter.GetNode(); iter ++) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 >= key);
		}
	}
#endif

	return (dgTreeNode *)prev;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::FindLessEqual(KEY key) const {
	if (m_head == NULL) {
		return NULL;
	}

	dgTreeNode *prev = NULL;
	dgTreeNode *ptr = m_head;

	while (ptr != NULL) {
		if (key == ptr->m_key) {
			return ptr;
		}

		if (key < ptr->m_key) {
			ptr = ptr->GetLeft();
		} else {
			prev = ptr;
			ptr = ptr->GetRight();
		}

	}

#ifdef __ENABLE_SANITY_CHECK
	if (prev) {
		Iterator iter(*this);
		for (iter.End(); iter.GetNode() != prev; iter --) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 >= key);
		}
		for (; iter.GetNode(); iter --) {
			KEY key1 = iter.GetKey();
			NEWTON_ASSERT(key1 < key);
		}
	}
#endif

	return (dgTreeNode *)prev;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Insert(const OBJECT &element, KEY key, bool &elementWasInTree) {
	dgTreeNode *parent = NULL;
	dgTreeNode *ptr = m_head;
	dgInt32 val = 0;
	elementWasInTree = false;
	while (ptr != NULL) {
		parent = ptr;

		if (key < ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == -1) ;
			val = -1;
			ptr = ptr->GetLeft();
		} else if (key > ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 1) ;
			val = 1;
			ptr = ptr->GetRight();
		} else {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 0) ;
			elementWasInTree = true;
			return ptr;
		}
	}

	m_count ++;
	NEWTON_ASSERT(m_allocator);
	ptr = new (m_allocator) dgTreeNode(element, key, parent);
	if (!parent) {
		m_head = ptr;
	} else {
		if (val < 0) {
			parent->m_left = ptr;
		} else {
			parent->m_right = ptr;
		}
	}

	dgTreeNode **const headPtr = (dgTreeNode **) &m_head;
//	ptr->InsertFixup ((dgRedBackNode**)&m_head);
	ptr->InsertFixup((dgRedBackNode **)headPtr);
	return ptr;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Insert(const OBJECT &element, KEY key) {
	bool foundState;

	dgTreeNode *const node = Insert(element, key, foundState);
	if (foundState) {
		return NULL;
	}
	return node;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Insert(typename dgTree<OBJECT, KEY>::dgTreeNode *const node, KEY key) {
	dgInt32 val = 0;
	dgTreeNode *ptr = m_head;
	dgTreeNode *parent = NULL;
	while (ptr != NULL) {
		parent = ptr;
//		val = CompareKeys (ptr->m_key, key);
//		if (val < 0) {
//			ptr = ptr->GetLeft();
//		} else if (val > 0) {
//			ptr = ptr->GetRight();
//		} else {
//			return NULL;
//		}

		if (key < ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == -1) ;
			val = -1;
			ptr = ptr->GetLeft();
		} else if (key > ptr->m_key) {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 1) ;
			val = 1;
			ptr = ptr->GetRight();
		} else {
			NEWTON_ASSERT(CompareKeys(ptr->m_key, key) == 0) ;
			return NULL;
		}
	}

	m_count ++;

	ptr = node;
	ptr->m_key = key;
	ptr->Initdata(parent);

	if (!parent) {
		m_head = ptr;
	} else {
		if (val < 0) {
			parent->m_left = ptr;
		} else {
			parent->m_right = ptr;
		}
	}

	dgTreeNode **const headPtr = (dgTreeNode **) &m_head;
//	ptr->InsertFixup ((dgRedBackNode**)&m_head);
	ptr->InsertFixup((dgRedBackNode **)headPtr);
	return ptr;
}


template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::Replace(OBJECT &element, KEY key) {
	dgTreeNode *parent = NULL;
	dgTreeNode *ptr = m_head;
	dgInt32 val = 0;
	while (ptr != NULL) {
		parent = ptr;

		NEWTON_ASSERT(0);
		val = CompareKeys(ptr->m_key, key);
		if (val == 0) {
			ptr->m_info = element;
			return ptr;
		}
		if (val < 0) {
			ptr = ptr->GetLeft();
		} else {
			ptr = ptr->GetRight();
		}
	}

	NEWTON_ASSERT(m_allocator);
	ptr = new (m_allocator) dgTreeNode(element, key, parent);
	if (!parent) {
		m_head = ptr;
	} else {
		if (val < 0) {
			parent->m_left = ptr;
		} else {
			parent->m_right = ptr;
		}
	}

	dgTreeNode **const headPtr = (dgTreeNode **) &m_head;
//	ptr->InsertFixup ((dgRedBackNode**)&m_head);
	ptr->InsertFixup((dgRedBackNode **)headPtr);
	return ptr;
}


template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::ReplaceKey(typename dgTree<OBJECT, KEY>::dgTreeNode *const node, KEY key) {
	Unlink(node);
	dgTreeNode *const ptr = Insert(node, key);
	NEWTON_ASSERT(ptr);
	return ptr;
}

template<class OBJECT, class KEY>
typename dgTree<OBJECT, KEY>::dgTreeNode *dgTree<OBJECT, KEY>::ReplaceKey(KEY oldKey, KEY newKey) {
	dgTreeNode *const node = Find(oldKey);
	return node ? ReplaceKey(node, newKey) : NULL;
}

template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::Unlink(typename dgTree<OBJECT, KEY>::dgTreeNode *const node) {
	m_count --;

	dgTreeNode **const headPtr = (dgTreeNode **) &m_head;
	node->Unlink((dgRedBackNode **)headPtr);
	NEWTON_ASSERT(!Find(node->GetKey()));
}


template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::Remove(typename dgTree<OBJECT, KEY>::dgTreeNode *const node) {
	m_count --;
	dgTreeNode **const headPtr = (dgTreeNode **) &m_head;
	node->Remove((dgRedBackNode **)headPtr);
}

template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::Remove(KEY key) {
	dgTreeNode *node;

	// find node in tree
	node = Find(key);
	if (node == NULL) {
		return;
	}
	Remove(node);
}

template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::RemoveAll() {
	if (m_head) {
		m_count  = 0;
		m_head->RemoveAll();
		m_head = NULL;
	}
}

template<class OBJECT, class KEY>
bool dgTree<OBJECT, KEY>::SanityCheck() const {
	return SanityCheck(m_head, 0);
}


template<class OBJECT, class KEY>
bool dgTree<OBJECT, KEY>::SanityCheck(typename dgTree<OBJECT, KEY>::dgTreeNode *const ptr, dgInt32 height) const {
	if (!ptr) {
		return true;
	}

	if (!ptr->IsInTree()) {
		return false;
	}

	if (ptr->m_left) {
		if (CompareKeys(ptr->m_key, ptr->GetLeft()->m_key) > 0) {
			return false;
		}
	}

	if (ptr->m_right) {
		if (CompareKeys(ptr->m_key, ptr->GetRight()->m_key) < 0) {
			return false;
		}
	}

	if (ptr->GetColor() == dgTreeNode::BLACK) {
		height ++;
	} else if (!((!ptr->m_left  || (ptr->m_left->GetColor() == dgTreeNode::BLACK)) &&
	             (!ptr->m_right || (ptr->m_right->GetColor() == dgTreeNode::BLACK)))) {
		return false;
	}

	if (!ptr->m_left && !ptr->m_right) {
		dgInt32 bh = 0;
		for (dgTreeNode *x = ptr; x; x = x->GetParent()) {
			if (x->GetColor() == dgTreeNode::BLACK) {
				bh ++;
			}
		}
		if (bh != height) {
			return false;
		}
	}

	if (ptr->m_left && !SanityCheck(ptr->GetLeft(), height)) {
		return false;
	}

	if (ptr->m_right && !SanityCheck(ptr->GetRight(), height)) {
		return false;
	}
	return true;
}


template<class OBJECT, class KEY>
dgInt32 dgTree<OBJECT, KEY>::CompareKeys(const KEY &key0, const KEY &key1) const {
	if (key1 < key0) {
		return - 1;
	}
	if (key1 > key0) {
		return 1;
	}
	return 0;
}

template<class OBJECT, class KEY>
void dgTree<OBJECT, KEY>::SwapInfo(dgTree<OBJECT, KEY> &tree) {
	Swap(m_head, tree.m_head);
	Swap(m_count, tree.m_count);
}

//template<class OBJECT, class KEY> dgInt32 dgTree<OBJECT,KEY>::m_size = 0;
//template<class OBJECT, class KEY> dgMemoryAllocator* dgTree<OBJECT,KEY>::m_staticAllocator = NULL;


#endif