// HuffmanDict.cpp: implementation of the HuffmanDict class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "HuffmanTest.h"
#include "HuffmanDict.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

HuffmanDict::HuffmanDict()
{
	CachedNTT=NULL;
}

HuffmanDict::~HuffmanDict()
{
	delete[] CachedNTT;
}

HuffmanDict::NTTEntry *HuffmanDict::GetNTTable(unsigned int &OutStateCount)
{
	if (!CachedNTT)
		CachedNTT=GenerateNTT(GetEncodeDict(),CachedStateCount);

	OutStateCount=CachedStateCount;
	return CachedNTT;
}

unsigned short HuffmanDict::GetStats(unsigned char &MaxCodeLen, unsigned char &MinCodeLen)
{
	EncEntry *DictA=GetEncodeDict();
	
	unsigned int x;
	unsigned short SymbolCount=0;
	unsigned char MaxLen=0,MinLen=255;
	for (x=0; x<256; x++)
	{
		if (DictA[x].CodeLen)
		{
			SymbolCount++;

			if (DictA[x].CodeLen>MaxLen)
				MaxLen=DictA[x].CodeLen;
			else if (DictA[x].CodeLen<MinLen)
				MinLen=DictA[x].CodeLen;
		}
	}

	MaxCodeLen=MaxLen;
	MinCodeLen=MinLen;
	return SymbolCount;
}

HuffmanDict::NTTEntry *HuffmanDict::GenerateNTT(EncEntry *EncodeDict, unsigned int &OutStateCount)
{
	HNode *HuffRoot=new HNode;
	unsigned int NameCount=BuildHuffTree(EncodeDict,HuffRoot)+1;

	NTTEntry *RetTable=new NTTEntry[ NameCount << 8 ];

	BuildTTable(HuffRoot,NameCount,RetTable);
	delete HuffRoot;
	
	OutStateCount=NameCount;
	return RetTable;
}

unsigned char HuffmanDict::BuildHuffTree(EncEntry *SourceTable, HNode *TargetRoot)
{
	//Needed to name every node.
	TargetRoot->Symbol=0;

	unsigned char CurrName=0;
	unsigned int x;
	for (x=0; x<256; x++)
	{
		if (SourceTable[x].CodeLen)
			CurrName=BuildHuffTreeInternal(CurrName,x,&(SourceTable[x]),TargetRoot);
	}

	return CurrName;
}

unsigned char HuffmanDict::BuildHuffTreeInternal(unsigned char CurrName, unsigned char CurrSymbol, EncEntry *SourceEntry, HNode *TargetNode)
{
	if (SourceEntry->CodeLen==1)
	{
		//Create a leaf node.
		if (SourceEntry->Code >> 31)
		{
			TargetNode->Right=new HNode;
			TargetNode->Right->Symbol=CurrSymbol;
		}
		else
		{
			TargetNode->Left=new HNode;
			TargetNode->Left->Symbol=CurrSymbol;
		}

		return CurrName;
	}
	else
	{
		//Create an internal node.
		HNode *NewNode;
		if (SourceEntry->Code >> 31)
		{
			if (TargetNode->Right)
			{
				SourceEntry->Code<<=1;
				SourceEntry->CodeLen--;
				return BuildHuffTreeInternal(CurrName,CurrSymbol,SourceEntry,TargetNode->Right);
			}
			else
			{
				NewNode=new HNode;
				TargetNode->Right=NewNode;
			}
		}
		else
		{
			if (TargetNode->Left)
			{
				SourceEntry->Code<<=1;
				SourceEntry->CodeLen--;
				return BuildHuffTreeInternal(CurrName,CurrSymbol,SourceEntry,TargetNode->Left);
			}
			else
			{
				NewNode=new HNode;
				TargetNode->Left=NewNode;
			}
		}
		
		CurrName++;
		//NewNode->Parent=TargetNode;
		NewNode->Symbol=CurrName;

		SourceEntry->Code<<=1;
		SourceEntry->CodeLen--;
		return BuildHuffTreeInternal(CurrName,CurrSymbol,SourceEntry,NewNode);
	}
}

/*To access a table entry, address the table as follows:
OutTable[(CurrState << 8) + CodedByte].*/
void HuffmanDict::BuildTTable(HNode *SourceNode, unsigned char NameCount, NTTEntry *OutTable)
{
	unsigned int x,y;

	HNode *CurrStartNode;
	for (y=0; y<NameCount; y++)
	{
		CurrStartNode=LookupInternalNode(SourceNode,y);
		for (x=0; x<256; x++)
			FillTTEntry(SourceNode, CurrStartNode, x, &(OutTable[(y << 8)+x]) );
	}
}

void HuffmanDict::FillTTEntry(HNode *RootNode, HNode *StartNode, unsigned char EncodedByte, NTTEntry *OutEntry)
{
	for (char x=7; x>=0; x--)
	{
		//Traverse the tree further.
		if ( (EncodedByte >> x) & 1 )
			StartNode=StartNode->Right;
		else
			StartNode=StartNode->Left;

		if (!(StartNode->Left))		
		{
			//Output a symbol, then jump to the root.
			if (OutEntry->SymbolA)
			{
				unsigned char *TmpBuff=new unsigned char[OutEntry->SymbolCount+1];
				memcpy(TmpBuff,OutEntry->SymbolA,OutEntry->SymbolCount);

				TmpBuff[OutEntry->SymbolCount]=StartNode->Symbol;
				
				delete[] OutEntry->SymbolA;
				OutEntry->SymbolA=TmpBuff;
				OutEntry->SymbolCount++;
			}
			else
			{
				OutEntry->SymbolA=new unsigned char;			
				*(OutEntry->SymbolA)=StartNode->Symbol;
				OutEntry->SymbolCount=1;
			}

			StartNode=RootNode;
		}
	}
	OutEntry->NextNodeName=StartNode->Symbol;
}

HuffmanDict::HNode * HuffmanDict::LookupInternalNode(HNode *RootNode, unsigned char NodeName)
{
	if (RootNode->Symbol==NodeName)
		return RootNode;
	else
	{
		HNode *RetNode;
		//Call only if the child is an internal node
		if (RootNode->Left->Left)
		{
			RetNode=LookupInternalNode(RootNode->Left,NodeName);
			if (RetNode)
				return RetNode;
		}
		if (RootNode->Right->Left)
		{
			RetNode=LookupInternalNode(RootNode->Right,NodeName);
			if (RetNode)
				return RetNode;
		}

		return NULL;
	}
}