#ifndef HUFFMAN_H_
#define HUFFMAN_H_

/*
 * (c) Markus Lumpe, 2007
 *	v 1.1
 */

#include <string>
#include <iostream>
#include <map>
#include <set>
#include <functional>

using namespace std;

typedef unsigned char byte;

//---- HuffmanException ----

class HuffmanException
{
private:
	string fMessage;
	
public:
	HuffmanException( string aMessage );
	HuffmanException( int aLine, string aFileName, string aMessage );
	
	string getMessage();	
};

#define HUFFMAN_EXCEPTION( aMessage ) HuffmanException( __LINE__, __FILE__, aMessage )

//---- HuffmanBitCode ----

class HuffmanBitCode
{
private:
	byte* fBitArray;
	unsigned int fCodeLength;

	void copy( const HuffmanBitCode& aCode );
	void extend( int aToCodeLength );
	
	void copy( byte aBitArray[], int aCodeLength );

public:
	HuffmanBitCode();
	HuffmanBitCode( byte aBitArray[], int aCodeLength );
	HuffmanBitCode( const HuffmanBitCode& aBitCode );
	~HuffmanBitCode();

	unsigned int size() const;
	byte* getBits() const;
	int lengthInByte() const;
	
	void add0();
	void add1();
	
	// required bc copy constructor
	HuffmanBitCode& operator=( const HuffmanBitCode& aBitCode );
	
	int at( unsigned int aIndex ) const;
	int operator[]( unsigned int aIndex ) const;
	
	// HuffmanBitCode iterator 
	
	class iterator
	{
	friend class HuffmanBitCode;
	
	private:
		HuffmanBitCode* fCode;
		unsigned int fIndex;

	// follow singleton pattern (constructor should not be public)	
	protected:
		iterator();
		iterator( HuffmanBitCode* aCode );
		iterator( HuffmanBitCode* aCode, unsigned int aPosition );
		
	public:
		
		// Iterator behavior
		int operator*() const;
		iterator& operator++();		// prefix
		iterator operator++(int);	// postix (extra unused argument)
		bool operator==( const iterator& ) const; 
		bool operator!=( const iterator& ) const;
		
		// non-standard iterator method to obtain a cloned end
		iterator end();
		iterator next();
	};
	
	iterator begin();
	iterator end();
};

bool operator<( const HuffmanBitCode& aLeft, const HuffmanBitCode& aRight );

ostream& operator<<( ostream& aOS, const HuffmanBitCode& aCode );

//---- T ----

template <class T>
T init() { return T(); }

//---- IHuffmanElement<T> ----

template <class T> class HuffmanCode;

template <class T>
class IHuffmanElement
{
public:
	virtual ~IHuffmanElement() {}
	
	virtual long getFrequency() const = 0;
	virtual long getDataSize() const = 0;
	virtual long getDictionarySize() const = 0;
	virtual int getDegree() const = 0;
	virtual void setHuffmanCode( HuffmanBitCode aHuffmanCode ) = 0;
	virtual void collectHuffmanCodes( set< HuffmanCode<T> >& aSet ) const = 0;
};

template <class T>
ostream& operator<<( ostream& aOS, const IHuffmanElement<T>* aElement );

//---- HuffmanCode<T> ----

template <class T>
class HuffmanCode
{
private:
	T fMappedCharacter;
	HuffmanBitCode fCode;
	
	static unsigned int fMaxLength;
	static unsigned int fMaxSize;
	
public:
	HuffmanCode();
	HuffmanCode( T aMappedCharacter );
	HuffmanCode( T aMappedCharacter, HuffmanBitCode aCode );

	T getMappedCharacter() const;
	const HuffmanBitCode getCode() const;
	
	void setHuffmanMapping( HuffmanBitCode aCode );

	static void setMaxLength( unsigned int aMaxLength );
	static unsigned int getMaxLength();
	static unsigned int getMaxSize();
};

template <class T>
bool operator<( const HuffmanCode<T>& aLeft, const HuffmanCode<T>& aRight );

template <class T>
ostream& operator<<( ostream& aOS, const HuffmanCode<T>& aCode );

//---- HuffmanLeaf<T> ----

template <class T>
class HuffmanLeaf : public IHuffmanElement<T>
{
private:
	HuffmanCode<T> fCharacter;
	long fFrequency;
	int fDegree;
	
public:
	HuffmanLeaf() {}
	HuffmanLeaf( T aCharacter, long aFrequency );
	HuffmanLeaf( HuffmanCode<T> aCode );
	virtual ~HuffmanLeaf() {}
	
	virtual T getCharacter() const;
	virtual long getFrequency() const;
	virtual long getDataSize() const;
	virtual long getDictionarySize() const;
	virtual int getDegree() const;
	
	const HuffmanBitCode getCode() const;
	virtual void setHuffmanCode( HuffmanBitCode aHuffmanCode );
	
	virtual void collectHuffmanCodes( set< HuffmanCode<T> >& aSet ) const;
};

template <class T>
ostream& operator<<( ostream& aOS, const HuffmanLeaf<T>& aLeave );

//---- HuffmanNode<T> ----

template <class T>
class HuffmanNode : public IHuffmanElement<T>
{
private:
	IHuffmanElement<T>* fLeft;
	IHuffmanElement<T>* fRight;
	long fSum;
	int fNumber;
	int fDegree;
	
	static int fCount;
	
public:
	HuffmanNode();
	HuffmanNode( IHuffmanElement<T>* aLeft, IHuffmanElement<T>* aRight );
	
	~HuffmanNode();
	
	IHuffmanElement<T>* getLeft() const;
	IHuffmanElement<T>* getRight() const;
	virtual long getFrequency() const;
	virtual long getDataSize() const;
	virtual long getDictionarySize() const;
	int getCount() const;
	virtual int getDegree() const;
	virtual void setHuffmanCode( HuffmanBitCode aHuffmanCode );
	virtual void collectHuffmanCodes( set< HuffmanCode<T> >& aSet ) const;
	
	void buildSubtree( HuffmanCode<T> aCode, HuffmanBitCode::iterator aRemainingBits );
};

template <class T>
ostream& operator<<( ostream& aOS, const HuffmanNode<T>* aNode );

//---- HuffmanTree<T> ----

template <class T>
struct LessIHuffmanElementPtr : public binary_function<IHuffmanElement<T>*, IHuffmanElement<T>*, bool>
{
	bool isLeave( const IHuffmanElement<T>* aElement ) const
	{
		return typeid(*aElement) == typeid(HuffmanLeaf<T>);	
	}
	
	bool isNode( const IHuffmanElement<T>* aElement ) const
	{
		return typeid(*aElement) == typeid(HuffmanNode<T>);	
	}
	
	bool operator()(const IHuffmanElement<T>* aLeft, const IHuffmanElement<T>* aRight) const
    { 
    	// We need to guarantee antisymmetry
    	
		// Leave < Leave
		if ( isLeave( aLeft ) && isLeave( aRight ) )
		{
			bool Result = aLeft->getFrequency() == aRight->getFrequency();
			if ( Result )
				return ((HuffmanLeaf<T>*)aLeft)->getCharacter() < ((HuffmanLeaf<T>*)aRight)->getCharacter();
			else
				return aLeft->getFrequency() < aRight->getFrequency();
		}

		// Leave < Node
		if ( isLeave( aLeft ) && isNode( aRight ) )
		{
			long lFL = aLeft->getFrequency();
			long lFR = aRight->getFrequency();
			
			// compare frequency only
			return lFL <= lFR;
		}

		// Node < Leave
		if ( isNode( aLeft ) && isLeave( aRight ) )
		{
			long lFL = aLeft->getFrequency();
			long lFR = aRight->getFrequency();
			
			// compare frequency only
			return lFL < lFR;
		}

		// Node < Node
		bool Result = aLeft->getFrequency() == aRight->getFrequency();
		if ( Result )
		{
			// continue with count (order of creation)
			return ((HuffmanNode<T>*)aLeft)->getCount() < ((HuffmanNode<T>*)aRight)->getCount();
		}
		else
			return aLeft->getFrequency() < aRight->getFrequency();
    }
};

template <class T>
class HuffmanTree
{
private:
	IHuffmanElement<T> *fElements;
	set< HuffmanCode<T> > fDictionary;
	long fNumberOfCharacters;

	set< IHuffmanElement<T>*, LessIHuffmanElementPtr<T> > analyzeFrequencies( istream& aIS ) const;

	void buildTree( istream& aIS );
	void buildTreeFromHeaderOfStream( istream& aIS );

	void buildHuffmanCodes();	
	void buildDictionary();
	void writeHeaderToStream( ostream& aOS );
	void writeHeaderToStreamB( ostream& aOS );

public:
	HuffmanTree();
	~HuffmanTree();

	void compressStream( ostream& aOS, istream& aIS );
	void uncompressStream( ostream& aOS, istream& aIS );
	
	IHuffmanElement<T>* getNodes() const;
	int getDegree() const;
	long getNumberOfCharacters() const;
	int getEffectiveRatio() const;
};

template <class T>
ostream& operator<<( ostream& aOS, const HuffmanTree<T>* aTree );

//---- HuffmanCharacter<T> ----

template <class T>
class HuffmanCharacter
{
public:
	typedef 
	  union
	  {
	  	T field1;
	  	unsigned char field2[sizeof( T )];
	  } character;

	static bool read( istream& aIS, T& aValue );
	static bool write( ostream& aOS, T aValue );
};

//---- Huffman<T> ----

template <class T>
class Huffman
{
public:
	static void compress( string aFileName );
	static void uncompress( string aFileName );
};

//---- BitStreamWriter ----

class BitStreamWriter
{
private:
	ostream& fOS;
	byte fBuffer[32];
	int fByteIndex;
	int fBitIndex;
	
	void init();
	
public:
	BitStreamWriter( ostream& aOS );
	
	void writeBits( HuffmanBitCode& aBits );
	void finish();
};

//---- BitStreamReader<T> ----

template <class T>
class BitStreamReader
{
private:
	istream& fIS;
	IHuffmanElement<T>* fTree;
	long fNumberOfCharacters;

	int fLastRead;	
	byte fBuffer[32];
	int fByteIndex;
	int fBitIndex;
	
	void init();
	
public:
	BitStreamReader( istream& aIS, HuffmanTree<T>* aTree );
	
	bool getCharacter( T& aCharacter );
};

#include "HuffmanImpl.h"

#endif /*HUFFMAN_H_*/