//
// Asura_RandTree.h
//
// Created by Kim Randell - 20/07/2004
// 
// Copyright (C) 2004 Rebellion, All Rights Reserved.
//
// This is a binary tree with randomised insertion. 
// It is much like the hash tree, and indeed much of the code is based on that
// but offers the advantage that no matter what order you insert keys, it 
// will insert at random depths, so that the tree is kept reasonably balanced.
// Insertions and deletions are slower than the hash tree, but searches are
// usually much faster.
// Obviously the randomness means no absolute guarantees can be made about the
// depth of the tree, but it is extremely unlikely that it will become
// significantly imbalanced.
// The trade-off is that nodes must keep track of the number of nodes
// in its sub-tree, so there's an extra 4 bytes on each node (this could be
// avoided, but only with a huge speed decrease for inserts and deletes)
//
// It is intended for use by systems which want to keep track of a large
// number of dynamic objects, which can be retrieved by a unique ID, and
// created or deleted at any point.
//
// It is a templated class, which allows you to use keys of any type, although
// I expect they'll usually be unsigned ints.
//
// For more details, see Algorithms in C++ by Robert Sedgewick, Chapter 13.1
//


#ifdef ASURA_USE_PRAGMA_ONCE
    #pragma once
#endif

#ifndef __ASURA_RANDTREE_H__
#define __ASURA_RANDTREE_H__

#include "Core/Asura_Handle.h"
#include "Core/Asura_Maths.h"
#include "Debugging/Asura_Assert.h"

template <class KeyType> class Asura_RandTree_Node
{
	public:
 
		KeyType m_xKey;
		int m_iIndex;
		Asura_RandTree_Node<KeyType> *m_pxLower;
		Asura_RandTree_Node<KeyType> *m_pxHigher;
		u_int m_uNumberOfNodes;

		Asura_RandTree_Node(const KeyType xKey, int iIndex)
		: m_xKey(xKey)
		, m_iIndex(iIndex)
		, m_pxLower(NULL)
		, m_pxHigher(NULL)
		, m_uNumberOfNodes(1)
		{}

		~Asura_RandTree_Node()
		{}

		// Recursive function to retrieve depth of nodes in this (sub)tree
		u_int GetDepth() const
		{
			// Default to zero sub-tree depth
			u_int uLower = 0;
			u_int uHigher = 0;
			// Retrieve depth of lower and higher sub-trees if they exist
			if(m_pxLower)	uLower  = m_pxLower->GetDepth();
			if(m_pxHigher)	uHigher = m_pxHigher->GetDepth();
			// Count this node, plus the largest of the lower and higher sub-trees
			return 1 + ((uLower > uHigher) ? uLower : uHigher);
		}
		// Returns the number of nodes below and including this one.
		inline u_int GetNumberOfNodes() const
		{
			return m_uNumberOfNodes;
		}
		// Purely for debugging
		void VerifyNumberOfNodes() const
		{
			Asura_Assert(m_uNumberOfNodes == GetNumberOfNodesRecursive(), "Mis-match in stored and actual number of nodes");
		}
		// Call this after changing subtrees (after a rotate or delete, for instance)
		void ReCalcNumberOfNodes()
		{
			// Get subtree node counts. This is not recursive - it is assumed that the sub-tree node counts are already correct.
			u_int uLower = m_pxLower ? m_pxLower->GetNumberOfNodes() : 0;
			u_int uHigher = m_pxHigher ? m_pxHigher->GetNumberOfNodes() : 0;
			m_uNumberOfNodes = 1 + uLower + uHigher;
		}
	private:
		// Only used by the VerifyNumberOfNodes function. Should never be used otherwise.
		u_int GetNumberOfNodesRecursive() const
		{
			u_int uLower = m_pxLower ? m_pxLower->GetNumberOfNodesRecursive() : 0;
			u_int uHigher = m_pxHigher ? m_pxHigher->GetNumberOfNodesRecursive() : 0;
			return 1 + uLower + uHigher;
		}
};

template <class KeyType> class Asura_RandTree
{
	public:
		// Default constructor
		Asura_RandTree()
		{
			m_iNumberOfNodes = 0;
			m_pxFirstRoot = 0;
			m_iMaximumNodes = 0x7fffffff;

			m_puHandleList = 0;
		}
		// Constructor with handles
		Asura_RandTree(int iMaximumNodes)
		{
			m_iNumberOfNodes = 0;
			m_pxFirstRoot = 0;
			m_iMaximumNodes = iMaximumNodes;

			m_puHandleList = new Asura_Handle_List(iMaximumNodes);
		}
		// Destructor
		~Asura_RandTree()
		{
			// Just check it's valid before cleaning up
			if(m_pxFirstRoot) m_pxFirstRoot->VerifyNumberOfNodes();
			FreeAll();

			if (m_puHandleList) delete m_puHandleList;
		}
		// Clears out all nodes and handles
        void Reset()
        {
			// Just check it's valid before cleaning up
			if(m_pxFirstRoot) m_pxFirstRoot->VerifyNumberOfNodes();
			// Delete all nodes
			FreeAll();
			m_iNumberOfNodes = 0;
			m_pxFirstRoot = 0;
			// Reset handle list if necessary
			if (m_puHandleList)
            {
                m_puHandleList->Reset();
            }
        }
		// To use this, GetNullKey for your key type must be implemented
		KeyType FindKeyByIndex( int __iIndex ) const
		{
			Asura_RandTree_Node<KeyType>* pxNode = m_pxFirstRoot;
            // Empty tree?
            if( !pxNode )
			{
				return GetNullKey();
			}
			// Search
			Asura_RandTree_Node<KeyType>* pxFoundNode = FindNodeByIndexRecursive(pxNode, __iIndex);
			return pxFoundNode ? pxFoundNode->m_xKey : GetNullKey();
		}
        // Returns -1 if not found.
		int FindNodeIndexByKey( KeyType xKey ) const
		{
			return FindNodeIndexByKeyRecursive( m_pxFirstRoot, xKey );
		}
		// Returns true if the node existed already
        bool InsertKey( KeyType xKey, int& iIndex )
        {
            return InsertKeyRecursive( &m_pxFirstRoot, xKey, iIndex );
        }
		// Delete a node and get its index, returns true if the node was found.
		// Note!!
		// The node with index NumberOfNodes-1 will have its index changed to the index
		// of the node being deleted. (since the indices must be consecutive)
		bool DeleteNodeByKey(KeyType xKey, int& iOutputIndex)
		{
			if(!DeleteNodeByKeyInternal(&m_pxFirstRoot, xKey, iOutputIndex)) return false;

			//need to alter the index of the last node to take the place of the deleted one
			if(m_pxFirstRoot)
			{
				Asura_RandTree_Node<KeyType>* pxNode = FindNodeByIndexRecursive(m_pxFirstRoot, m_iNumberOfNodes);
				// KOR - 22/07/2004 - it is possible that we've removed the last node, so this assert is unnecessary
				//Asura_Assert(pxNode, "Couldn't get last node");
				if(pxNode)
				{
					pxNode->m_iIndex = iOutputIndex;
				}
			}			
			return true;//search and destroy mission accomplished...
		}
		// Delete node without giving index. Can only be used if the hashtree was constructed with the maximum nodes parameter
		bool DeleteNodeByKey(KeyType xKey)
		{
			Asura_Assert(m_puHandleList!=0,"This HashTree was not constructed with a Handle List");

			int iOutputIndex;

			if(!DeleteNodeByKeyInternal(&m_pxFirstRoot, xKey, iOutputIndex)) return false;

			m_puHandleList->ReleaseHandle(iOutputIndex);

			return true;
		}
		// Constructs a list of keys in index order. Array must be pre-allocated
		void MakeKeyList(KeyType* pxKeyList) const
		{
			if (m_pxFirstRoot) MakeKeyListRecursive(m_pxFirstRoot, pxKeyList);
		}
		// Returns the number of nodes in the tree
		inline int GetNumberOfNodes(void) const
		{
			return m_iNumberOfNodes;
		}
		// Function to retrieve depth of nodes in this tree.
		// It is not recommended that you call this per frame, as it recurses down every node in the tree
		// (I can't see why you'd want to anyway) - it's mainly intended as a debug check for tree balance.
		int GetDepth() const
		{
			return m_pxFirstRoot ? m_pxFirstRoot->GetDepth() : 0;
		}
		// Another debug function
		void Verify()
		{
			if(m_pxFirstRoot) m_pxFirstRoot->VerifyNumberOfNodes();
		}
	private:
		// Deletes all nodes
		void FreeAll(void)
		{
			if (m_pxFirstRoot) FreeAllRecursive(m_pxFirstRoot);
		}
		// Deletes all nodes in this sub-tree
		void FreeAllRecursive(Asura_RandTree_Node<KeyType> *pxNode)
		{
			if (pxNode->m_pxLower)
			{
				FreeAllRecursive(pxNode->m_pxLower);
			}
			if (pxNode->m_pxHigher)
			{
				FreeAllRecursive(pxNode->m_pxHigher);
			}
			delete pxNode;
		}
		// Adds all nodes in the sub-tree to the list, in the positions given by the indices.
		void MakeKeyListRecursive(Asura_RandTree_Node<KeyType>* pxNode, KeyType* pxKeyList) const
		{
			if (pxNode->m_pxLower)
			{
				MakeKeyListRecursive(pxNode->m_pxLower, pxKeyList);
			}
			if (pxNode->m_pxHigher)
			{
				MakeKeyListRecursive(pxNode->m_pxHigher, pxKeyList);
			}
			
			pxKeyList[pxNode->m_iIndex] = pxNode->m_xKey;
		}
        // Finds the node with the given index.
        Asura_RandTree_Node<KeyType>* FindNodeByIndexRecursive(Asura_RandTree_Node<KeyType>* __pxNode, int __iIndex ) const
        {
            if( __pxNode->m_iIndex == __iIndex )
            {
                return __pxNode;
            }
            else
            {
                Asura_RandTree_Node<KeyType>* pxRetVal = NULL;

                if( __pxNode->m_pxLower )
                {
                    pxRetVal = FindNodeByIndexRecursive( __pxNode->m_pxLower, __iIndex );
                }

                if( !pxRetVal && __pxNode->m_pxHigher )
                {
                    pxRetVal = FindNodeByIndexRecursive( __pxNode->m_pxHigher, __iIndex );
                }

                return pxRetVal;
            }
        }

		int FindNodeIndexByKeyRecursive( Asura_RandTree_Node<KeyType>* pxNode, KeyType xKey ) const
		{
			if( !pxNode )
			{
				return -1;
			}

			if( pxNode->m_xKey == xKey )
			{
				return pxNode->m_iIndex;
			}

            pxNode = ( xKey < pxNode->m_xKey ? pxNode->m_pxLower : pxNode->m_pxHigher );
            return FindNodeIndexByKeyRecursive( pxNode, xKey );
		}

		bool InsertKeyRecursive( Asura_RandTree_Node<KeyType>** ppxNode, KeyType xKey, int& iIndex )
		{
			if (!(*ppxNode))
			{
				AddNode(ppxNode, xKey, iIndex);
				return false;
			}

			if((*ppxNode)->m_xKey == xKey)
			{
				// Found it
				iIndex = (*ppxNode)->m_iIndex;
				return true;
			}

            const float fRootInsertProbability = 1.f/((*ppxNode)->GetNumberOfNodes() + 1);
			// Insert it
			// This is where the randomisation is applied. Insert node at root with a probability of 1/(Number of nodes + 1)
			if( Asura_Random::GetFloat() < fRootInsertProbability )
			{
				bool bFound = InsertNodeAtRoot( ppxNode, xKey, iIndex );
				return bFound;
			}
			else
			{
				// Recurse down
				if(xKey < (*ppxNode)->m_xKey)
				{
					if( !InsertKeyRecursive(&(*ppxNode)->m_pxLower, xKey, iIndex) )
					{
						// Increment count if node has been added
						(*ppxNode)->m_uNumberOfNodes++;
						return false;
					}
					return true;
				}
				else
				{
					if( !InsertKeyRecursive(&(*ppxNode)->m_pxHigher, xKey, iIndex) )
					{
						// Increment count if node has been added
						(*ppxNode)->m_uNumberOfNodes++;
						return false;
					}
					return true;
				}
			}
		}

		// Adds a node at this position, and gives the new index or handle
		void AddNode(Asura_RandTree_Node<KeyType>** ppxNode, KeyType xKey, int& iIndex)
		{
			*ppxNode = new Asura_RandTree_Node<KeyType>(xKey, m_iNumberOfNodes++);
			if (m_puHandleList)
			{
				#ifdef ASURA_ASSERT
				bool bSuccess = 
				#endif
				m_puHandleList->GetNewHandle((*ppxNode)->m_iIndex);
				Asura_Assert(bSuccess,"No more space!");
			}
			iIndex = (*ppxNode)->m_iIndex;
		}
		// Returns true if there was already a node with that key.
		bool InsertNodeAtRoot(Asura_RandTree_Node<KeyType>** ppxNode, KeyType xKey, int& iIndex)
		{
			if(!*ppxNode)
			{
				// Null pointer. Create the node here.
				AddNode(ppxNode, xKey, iIndex);
				// Not found
				return false;
			}
			if(xKey == (*ppxNode)->m_xKey)
			{
				// Found.
				iIndex = (*ppxNode)->m_iIndex;
				return true;
			}
			if(xKey < (*ppxNode)->m_xKey)
			{
				// Insert lower
				if(!InsertNodeAtRoot( &(*ppxNode)->m_pxLower, xKey, iIndex))
				{
					// Not found, so it's been inserted. Rotate right to bring to root.
					RotateRight(ppxNode);
					return false;
				}
				// Otherwise, found. Pass on the information
				return true;
			}
			else
			{
				// Insert higher
				if(!InsertNodeAtRoot( &(*ppxNode)->m_pxHigher, xKey, iIndex))
				{
					// Not found, so it's been inserted. Rotate left to bring to root.
					RotateLeft(ppxNode);
					return false;
				}
				// Otherwise, found. Pass on the information
				return true;
			}
		}
		// Moves lower node to root
		void RotateRight(Asura_RandTree_Node<KeyType>** ppxNode)
		{
			Asura_Assert(*ppxNode, "Nothing at this node");
			if(*ppxNode)
			{
				Asura_RandTree_Node<KeyType>* pxNewRight = *ppxNode;
				Asura_RandTree_Node<KeyType>* pxMiddle = (*ppxNode)->m_pxLower->m_pxHigher;
				Asura_RandTree_Node<KeyType>* pxNewRoot = (*ppxNode)->m_pxLower;
				*ppxNode = pxNewRoot;
				pxNewRoot->m_pxHigher = pxNewRight;
				pxNewRoot->m_pxHigher->m_pxLower = pxMiddle;
				// Update number of nodes
				pxNewRoot->m_pxHigher->ReCalcNumberOfNodes();
				pxNewRoot->ReCalcNumberOfNodes();
			}
		}
		// Moves higher node to root
		void RotateLeft(Asura_RandTree_Node<KeyType>** ppxNode)
		{
			Asura_Assert(*ppxNode, "Nothing at this node");
			if(*ppxNode)
			{
				Asura_RandTree_Node<KeyType>* pxNewLeft = *ppxNode;
				Asura_RandTree_Node<KeyType>* pxMiddle = (*ppxNode)->m_pxHigher->m_pxLower;
				Asura_RandTree_Node<KeyType>* pxNewRoot = (*ppxNode)->m_pxHigher;
				*ppxNode = pxNewRoot;
				pxNewRoot->m_pxLower = pxNewLeft;
				pxNewRoot->m_pxLower->m_pxHigher = pxMiddle;
				// Update number of nodes
				pxNewRoot->m_pxLower->ReCalcNumberOfNodes();
				pxNewRoot->ReCalcNumberOfNodes();
			}
		}
		// Recursive delete
		bool DeleteNodeByKeyInternal(Asura_RandTree_Node<KeyType>** ppxRootNode, KeyType xKey,int& iOutputIndex)
		{
			if(!(*ppxRootNode))
			{
				return false;
			}

			if((*ppxRootNode)->m_xKey == xKey)
			{
				// Found it. Delete it
				iOutputIndex = (*ppxRootNode)->m_iIndex;
				Asura_RandTree_Node<KeyType>* pxNewRoot = JoinNodes(&(*ppxRootNode)->m_pxLower, &(*ppxRootNode)->m_pxHigher);
				delete (*ppxRootNode);
				m_iNumberOfNodes--;
				(*ppxRootNode) = pxNewRoot;
				if((*ppxRootNode))
				{
					(*ppxRootNode)->ReCalcNumberOfNodes();
				}
				return true;
			}
			else
			{
				if(xKey < (*ppxRootNode)->m_xKey)
				{
					if(DeleteNodeByKeyInternal(&(*ppxRootNode)->m_pxLower, xKey, iOutputIndex))
					{
						(*ppxRootNode)->ReCalcNumberOfNodes();
						return true;
					}
					return false;
				}
				else
				{
					if(DeleteNodeByKeyInternal(&(*ppxRootNode)->m_pxHigher, xKey, iOutputIndex))
					{
						(*ppxRootNode)->ReCalcNumberOfNodes();
						return true;
					}
					return false;
				}
			}
		}
		// Joins two sub-trees into one, randomly choosing which to make the root
		Asura_RandTree_Node<KeyType>* JoinNodes(Asura_RandTree_Node<KeyType>** ppxRoot1, Asura_RandTree_Node<KeyType>** ppxRoot2)
		{
			// Check for null
			if(!(*ppxRoot1)) return (*ppxRoot2);
			if(!(*ppxRoot2)) return (*ppxRoot1);
			const float fProbabilityForRoot1 = static_cast<float>((*ppxRoot1)->GetNumberOfNodes())/((*ppxRoot1)->GetNumberOfNodes() + (*ppxRoot2)->GetNumberOfNodes());
			if(Asura_Random::GetFloat() < fProbabilityForRoot1)
			{
				(*ppxRoot1)->m_pxHigher = JoinNodes(&(*ppxRoot1)->m_pxHigher, ppxRoot2);
				(*ppxRoot1)->ReCalcNumberOfNodes();
				return (*ppxRoot1);
			}
			else
			{
				(*ppxRoot2)->m_pxLower = JoinNodes(ppxRoot1, &(*ppxRoot2)->m_pxLower);
				(*ppxRoot2)->ReCalcNumberOfNodes();
				return (*ppxRoot2);
			}
		}
		KeyType GetNullKey() const;

		// Root of tree
		Asura_RandTree_Node<KeyType> *m_pxFirstRoot;
		// Only used if constructed with maximum nodes
		Asura_Handle_List *m_puHandleList;

		int m_iNumberOfNodes;
		int m_iMaximumNodes;
};

#endif // __ASURA_RANDTREE_H__
