/*! \file BitTree.h
	\brief Declaration and implementation of the sps::BitTree class. 

	Include this file if you want to use the sps::BitTree class.
	
	\sa
	- BitList.h
	- DenseBitVector.h
	- SparseBitVector.h
	- RiceSet.h
	\sa
*/

#ifndef _BitTree_h_included_
#define _BitTree_h_included_

#include "Helpers.h"
#include "Iterator.h"
#include <vector>

namespace sps
{
	/*! \brief Implementation of the BitSet interface based on a binary tree.
		
		In this implementation of the BitSet interface the data is stored in the
		leaves of a binary tree. The position of the leaves in the tree encode the
		most significant bits of each set member, the least significant bits are 
		stored in the leaves. Starting at the root and the most significant bit,
		the left subtree will contain all nodes where the next bit is 0, the right
		subtree all where the next bit is 1. For the leaves no special node type
		is used, but the links of the parents of the leave nodes are used to
		store the least significant 5-6 bits (depending on your machine word length).

		Due to the fact that this implementation is based on a binary tree
		the runtime of the methods Add, Remove and Contains is in O(lg(n)) and
		the runtime of the methods Union, Intersect and Difference in O(n*lg(n)),
		where n is the biggest value ever stored in this instance of the tree.

		\sa
			- sps::BitList
			- sps::DenseBitVector
			- sps::SparseBitVector
			- sps::RiceSet
		\sa
	*/
	class BitTree
	{
	protected:
		struct Node
		{
			Node( size_t nLeft, size_t nRight )
			{
				m_arrData[ 0 ] = nLeft;
				m_arrData[ 1 ] = nRight;
			}

			size_t m_arrData[ 2 ];
		};
		typedef std::vector< Node > TNodeVector;

	public:
		/*! \brief Implementation of the sps::Iterator interface for a sps::BitTree. 
			
			Example: 
			\code
			sps::BitTree set;
			sps::BitTree::Iterator it = set.GetIterator();
			for ( it.Begin(); it.HasMoreElement(); it.Next() )
			{
				cout << *it << ", ";
			}
			\endcode
		*/
		class Iterator : public sps::Iterator
		{
		public:
			Iterator();

			/*! \brief Constructs an iterator for rBitTree. 
			*/
			Iterator( const BitTree& rBitTree );
			
			/*! \brief Checks if there are elements left. 

				If this method returns true it's save to call operator*() and Next().
				The runtime of this method is in O(1).
			*/
			virtual bool HasMoreElements() const;
			
			/*! \brief Sets the iterator on the first bit in the sps::BitTree.

				This method should be called before any other method of the Iterator
				is called. 
				The runtime of Begin() is in O(lg(n)). 
			*/
			virtual void Begin();
			
			/*! \brief Sets the iterator to the next bit in the sps::BitTree.

				Don't call this method if HasMoreElements() returned false. 
				The runtime of Next() is in O(lg(n)). 
			*/
			virtual void Next();
			
			/*! \brief Returns the current value(bit) the iterator points to.

				Don't call this method if HasMoreElements() returned false. 
				The runtime of operator*() is in O(1).
			*/
			virtual size_t operator*() const;

			/*! \brief Overwrites this instance of Iterator with a copy of \c rOther. 
			
				The runtime of this method is in O(lg(n)). 
			*/
			Iterator& operator=( const Iterator& rOther );

		private:
			struct StackFrame
			{
				enum TState
				{
					k_nNew = 0,
					k_nLeft,
					k_nRight,
					k_nFinished
				};

				StackFrame( size_t nIndex, size_t nOffset, size_t nDepth )
				: m_nIndex( nIndex )
				, m_nOffset( nOffset )
				, m_nCurrentDepth( nDepth )
				, m_nState( k_nNew )
				{}

				size_t m_nIndex;
				size_t m_nCurrentDepth;
				size_t m_nState;
				size_t m_nOffset;
			};

			typedef std::vector< StackFrame > TStack;

			const BitTree* m_pBitTree;
			TStack m_Stack;
			size_t m_nBit;
		};

		friend class Iterator;

		BitTree();

		/*! \brief Adds bit \c nBit to the set.
		
			The runtime of this method is in O(lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set;
				set.Add(3);
				set.Add(42); // the set now contains the elements 3, 42
			\endcode
		*/
		void Add( size_t nBit );

		/*! \brief Removes bit \c nBit from the set.
		
			The runtime of this method is in O(lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set;
				set.Add(3);
				set.Add(42); 
				set.Remove(3); // the set now contains only the element 42
			\endcode
		*/
		void Remove( size_t nBit );

		/*! \brief Checks if the set is empty.
		
			Returns true if the set contains no elements, false otherwise.
			The runtime of this method is in O(1). 

			Example:
			\code
				BitTree set;
				set.IsEmpty(); // will return true
				set.Add(3);
				set.Add(42); 
				set.IsEmpty(); // will return false
			\endcode
		*/
		bool IsEmpty() const;

		/*! \brief Checks if bit \c nBit is a member of the set.
		
			Returns true if the set contains \c nBit, false otherwise.
			The runtime of this method is in O(lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set;
				set.Add(3);
				set.Add(42); 
				set.Contains(42); // will return true
			\endcode
		*/
		bool Contains( size_t nBit ) const;

		/*! \brief Checks the two sets for equalness.
		
			Returns true if \c this and \c rOther contain exactly the same bits, false otherwise.
			The runtime of this method is in O(n*lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set1, set2;
				set1.Add(3);
				set1.Add(42); 

				set2.Add(3);
				set1.IsEqual(set2); // will return false

				set2.Add(42);
				set1.IsEqual(set2); // will return true
			\endcode
		*/
		bool IsEqual( const BitTree& rOther ) const;

		/*! \brief Builds the union of the two sets in place.
		
			The result set (\c this) will contain all bits that are set in \c this
			\b or in \c rOther.
			The runtime of this method is in O(n*lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set1, set2;
				set1.Add(3);
				set1.Add(42); 

				set2.Add(3);
				set2.Add(43);

				set1.Union(set2); // set1 now contains 3, 42 and 43
			\endcode
		*/
		void Union( const BitTree& rOther );

		/*! \brief Builds the intersection of the two sets in place.
		
			The result set (\c this) will contain all bits that are set in \c this
			\b and in \c rOther.
			The runtime of this method is in O(n*lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set1, set2;
				set1.Add(3);
				set1.Add(42); 

				set2.Add(3);
				set2.Add(43);

				set1.Intersect(set2); // set1 now contains only 3
			\endcode
		*/
		void Intersect( const BitTree& rOther );

		/*! \brief Builds the difference of the two sets in place.
		
			The result set (\c this) will contain all bits that are set in \c this
			\b and \b not set in \c rOther.
			The runtime of this method is in O(n*lg(n)). Where n is the biggest value
			ever stored in this instance of the sps::BitTree.

			Example:
			\code
				BitTree set1, set2;
				set1.Add(3);
				set1.Add(42); 

				set2.Add(3);
				set2.Add(43);

				set1.Difference(set2); // set1 now contains only 42
			\endcode
		*/
		void Difference( const BitTree& rOther );

		/*! \brief Clears the set.
		
			The result set (\c this) will be empty.
			The runtime of this method is in O(n). Where n is the number of nodes
			in the tree.

			Example:
			\code
				BitTree set;
				set.Add(3);
				set.Add(42); 
				set.Clear(); // set is now empty
			\endcode
		*/
		void Clear();

		/*! \brief Creates an iterator for the set.
		
			An sps::BitTree::Iterator for the members of this set is returned.
			The runtime of this method is in O(1).

			Example:
			\code
				BitTree set;
				set.Add(3);
				set.Add(42); 
				BitTree::Iterator it = set.GetIterator();
				for(it.Begin(); it.HasMoreElement(); it.Next() )
				{
					cout << *it << endl; 
				}
			\endcode
		*/
		Iterator GetIterator() const;

	protected:
		// Searchs for the leaf containing nBit in the tree,
		// missing nodes on the way are created.
		// Returns a pointer to bits-value of the leaf.
		size_t* CreateLeaf( size_t nBit );

		// Searchs for the leaf containing nBit in the tree.
		// Returns a pointer to the bits-value of the leaf, if found,
		// NULL otherwise.
		size_t* FindLeaf( size_t nBit );

		// Searchs for the leaf containing nBit in the tree.
		// Returns a pointer to the bits-value of the leaf, if found,
		// NULL otherwise.
		const size_t* FindLeaf( size_t nBit ) const;

		TNodeVector m_arrNodes;
		size_t m_nDepth;
		size_t m_nWordSize;
		size_t m_nMask;
		size_t m_nRoot;
	};

	inline BitTree::BitTree()
	{
		Clear();
	}

	inline void BitTree::Add( size_t nBit )
	{
		size_t* pnLeaf = CreateLeaf( nBit );
		*pnLeaf |= 1 << ( nBit & m_nMask );
	}

	inline void BitTree::Remove( size_t nBit )
	{
		size_t* pnLeaf = FindLeaf( nBit );
		if ( pnLeaf )
		{
			*pnLeaf &= ~( 1 << ( nBit & m_nMask ) );
		}
	}

	inline bool BitTree::IsEmpty() const
	{
		Iterator it = GetIterator();
		return !it.HasMoreElements();
	}

	inline bool BitTree::Contains( size_t nBit ) const
	{
		const size_t* pnLeaf = FindLeaf( nBit );
		if ( pnLeaf )
		{
			return ( *pnLeaf & ( 1 << ( nBit & m_nMask ) ) ) != 0;
		}
		return false;
	}

	inline bool BitTree::IsEqual( const BitTree& rOther ) const
	{
		Iterator it = GetIterator();
		Iterator itOther = rOther.GetIterator();
		for (	it.Begin(), itOther.Begin(); 
				it.HasMoreElements() && itOther.HasMoreElements(); 
				it.Next(), itOther.Next() 
		)
		{
			if ( *it != *itOther )
			{
				return false;
			}
		}
		return !it.HasMoreElements() && !itOther.HasMoreElements();
	}

	inline void BitTree::Union( const BitTree& rOther )
	{
		BitTree::Iterator itOther = rOther.GetIterator();
		for ( itOther.Begin(); itOther.HasMoreElements(); itOther.Next() )
		{
			Add( *itOther );
		}
	}

	inline void BitTree::Intersect( const BitTree& rOther )
	{
		BitTree::Iterator it = GetIterator();
		for ( it.Begin(); it.HasMoreElements(); it.Next() )
		{
			if ( !rOther.Contains( *it ) )
			{
				Remove( *it );
			}
		}
	}

	inline void BitTree::Difference( const BitTree& rOther )
	{
		BitTree::Iterator it = rOther.GetIterator();
		for ( it.Begin(); it.HasMoreElements(); it.Next() )
		{
			Remove( *it );
		}
	}

	inline void BitTree::Clear()
	{
		m_arrNodes.clear();
		m_arrNodes.push_back( Node( 0, 0 ) ); // NULL-Node
		m_arrNodes.push_back( Node( 0, 0 ) ); // Root-Node
		m_nRoot = 1;
		m_nDepth = 1;
		m_nWordSize = m_nDepth + lg( sizeof_bits< size_t >() );
		m_nMask = sizeof_bits< size_t >() - 1;
	}

	inline BitTree::Iterator BitTree::GetIterator() const
	{
		return BitTree::Iterator( *this );
	}

	inline size_t* BitTree::CreateLeaf( size_t nBit )
	{
		// let the tree grow if necessary
		while ( nBit > size_t( -1 ) >> ( sizeof_bits< size_t >() - m_nWordSize ) )
		{
			m_arrNodes.push_back( Node( m_nRoot, 0 ) );
			m_nRoot = m_arrNodes.size() - 1;
			++m_nWordSize;
			++m_nDepth;
		}

		size_t nNextNode = m_nRoot;
		size_t* pnNextNode = &nNextNode;
		for ( size_t i = 0; i != m_nDepth; ++i )
		{
			Node& rNode = m_arrNodes[ nNextNode ];
			if ( ( ( nBit << i ) & ( 1 << ( m_nWordSize - 1 ) ) ) == 0 )
			{
				pnNextNode = &rNode.m_arrData[ 0 ];
			}
			else
			{
				pnNextNode = &rNode.m_arrData[ 1 ];
			}

			if ( i + 1 < m_nDepth )
			{
				if ( !*pnNextNode )
				{
					*pnNextNode = m_arrNodes.size();
					nNextNode = *pnNextNode;
					m_arrNodes.push_back( Node( 0, 0 ) );
				}
				else
				{
					nNextNode = *pnNextNode;
				}
			}
		}
		return pnNextNode;
	}

	inline size_t* BitTree::FindLeaf( size_t nBit )
	{
		size_t nNextNode = m_nRoot;
		size_t* pnNextNode = &nNextNode;
		for ( size_t i = 0; i != m_nDepth; ++i )
		{
			Node& rNode = m_arrNodes[ *pnNextNode ];
			if ( ( ( nBit << i ) & ( 1 << ( m_nWordSize - 1 ) ) ) == 0 )
			{
				pnNextNode = &rNode.m_arrData[ 0 ];
			}
			else
			{
				pnNextNode = &rNode.m_arrData[ 1 ];
			}

			if ( i + 1 < m_nDepth )
			{
				if ( !*pnNextNode )
				{
					return 0;
				}
			}
		}
		return pnNextNode;
	}

	inline const size_t* BitTree::FindLeaf( size_t nBit ) const
	{
		size_t nNextNode = m_nRoot;
		const size_t* pnNextNode = &nNextNode;
		for ( size_t i = 0; i != m_nDepth; ++i )
		{
			const Node& rNode = m_arrNodes[ *pnNextNode ];
			if ( ( ( nBit << i ) & ( 1 << ( m_nWordSize - 1 ) ) ) == 0 )
			{
				pnNextNode = &rNode.m_arrData[ 0 ];
			}
			else
			{
				pnNextNode = &rNode.m_arrData[ 1 ];
			}

			if ( i + 1 < m_nDepth )
			{
				if ( !*pnNextNode )
				{
					return 0;
				}
			}
		}
		return pnNextNode;
	}

	///////////////////////////////////////////////////////////////////////////
	// BitTree::Iterator implementation
	///////////////////////////////////////////////////////////////////////////
	
	inline BitTree::Iterator::Iterator()
	: m_pBitTree( NULL ), m_nBit( 0 )
	{
	}

	inline BitTree::Iterator::Iterator( const BitTree& rBitTree )
	: m_pBitTree( &rBitTree ), m_nBit( 0 )
	{
	}

	inline bool BitTree::Iterator::HasMoreElements() const
	{
		return !m_Stack.empty();
	}

	inline void BitTree::Iterator::Begin()
	{
		m_Stack.clear();
		m_Stack.push_back( StackFrame( m_pBitTree->m_nRoot, 0, 0 ) );
		Next();
	}

	inline void BitTree::Iterator::Next()
	{
		while ( !m_Stack.empty() )
		{
			StackFrame& rFrame = m_Stack.back();
			const BitTree::Node rNode = m_pBitTree->m_arrNodes[ rFrame.m_nIndex ];

			if ( rFrame.m_nState == StackFrame::k_nNew )
			{
				++rFrame.m_nState;
				m_nBit = -1;
			}

			if ( rFrame.m_nCurrentDepth + 1 == m_pBitTree->m_nDepth )
			{
				do
				{
					++m_nBit;
					if ( m_nBit >= sizeof_bits< size_t >() )
					{
						m_nBit = 0;
						++rFrame.m_nState;
					}
				} 
				while ( 
					rFrame.m_nState != StackFrame::k_nFinished &&
					( rNode.m_arrData[ rFrame.m_nState - 1 ] & ( 1 << m_nBit ) ) == 0
				);

				if ( rFrame.m_nState != StackFrame::k_nFinished )
				{
					break;
				}
			}

			if ( rFrame.m_nState == StackFrame::k_nFinished )
			{
				m_Stack.pop_back();
			}
			else
			{
				++rFrame.m_nState;
				if ( rNode.m_arrData[ rFrame.m_nState - 2 ] )
				{
					size_t nDepth = rFrame.m_nCurrentDepth + 1;
					size_t nOffset = rFrame.m_nOffset | ( ( rFrame.m_nState - 2 ) << ( m_pBitTree->m_nWordSize - nDepth ) );
					m_Stack.push_back( StackFrame( rNode.m_arrData[ rFrame.m_nState - 2 ], nOffset,	nDepth ) );
				}
			}
		}
	}

	inline size_t BitTree::Iterator::operator*() const
	{
		if ( m_Stack.empty() )
		{
			return -1;
		}

		const StackFrame& rFrame = m_Stack.back();
		size_t nBit = rFrame.m_nOffset + m_nBit;
		if ( rFrame.m_nState == StackFrame::k_nRight )
		{
			nBit += 1 << ( m_pBitTree->m_nWordSize - m_pBitTree->m_nDepth );
		}
		return nBit;
	}

	inline BitTree::Iterator& BitTree::Iterator::operator=( const BitTree::Iterator& rOther )
	{
		m_nBit = rOther.m_nBit;
		m_pBitTree = rOther.m_pBitTree;
		m_Stack = rOther.m_Stack;
		return *this;
	}
}

#endif //_BitTree_h_included_