#include <stdio.h>
#include <assert.h>
#include <stdlib.h>

#include "bitstreams.h"
#include "StaticHuffman.h"
#include "HuffmanHeap.h"

void Code::ReverseSequence (void)
{
	unsigned int i;
	assert(length>0); // otherwise this'd be weird

	for (i=0; i<length/2; i++) {
		// swap them
		char temp;
		temp=sequence[i];
		sequence[i]=sequence[length-1-i];
		sequence[length-1-i]=temp;
	}
}
void Code::WriteBits (FILE *file)
{
	unsigned int i;
	for (i=0; i<length; i++) {
		WriteOneBit(sequence[i],file);
	}
}


StaticHuffmanTree::StaticHuffmanTree (void)
{
}

StaticHuffmanTree::~StaticHuffmanTree ()
{
}

void StaticHuffmanTree::CountFrequencies(FILE *normalFile)
{
	// allocate the nodes, and initialize them
	unsigned long i;
	for (i=0; i<257; i++) {
		codes[i].node=new StaticHuffmanNode;
		codes[i].node->isLeaf=true;
		codes[i].node->theChar=i;
		codes[i].node->weight=0;
		codes[i].node->parent=NULL;
		codes[i].node->child[0]=NULL;
		codes[i].node->child[1]=NULL;
	}
	// special case for null character
	codes[256].node->weight=1;

	// read every character from file & adjust count
	rewind(normalFile);
	while (!feof(normalFile)) {
		int ch=fgetc(normalFile);
		if (ch==-1) {
			if (feof(normalFile)) {
				break;
			}

			// using exceptions might be nice...
			printf("Error while reading normal file.  Didn't expect that.\n");
			exit(0);
		}

		assert(ch>=0 && ch<256);
		codes[ch].node->weight++;
	}

	// done.
}

void StaticHuffmanTree::BuildTreeFromForest(void)
{
	HuffmanHeap heap;
	int i;
	for (i=0; i<257; i++) {
		HuffmanHeapData data;
		data.priority=codes[i].node->weight;
		data.node=codes[i].node;
		heap.Insert(data);
	}

	unsigned long lastPriority=0;

	// now pull 'em out until there's nothing
#if _DEBUG
	int loops=0;
#endif

	while (1) {
		HuffmanHeapData tree0,tree1,newParent;

		assert(!heap.IsEmpty()); // can't be empty now
		tree0=heap.DeleteMin();
#ifdef _DEBUG
		//assert(tree0.priority<20000);
		assert(lastPriority<=tree0.priority);
		lastPriority=tree0.priority;
#endif

		// here's where it'd be empty if we got the root
		if (heap.IsEmpty()) {
			root=tree0.node;
			break;
		}

		tree1=heap.DeleteMin();
#ifdef _DEBUG
		//assert(tree0.priority<20000);
		assert(lastPriority<=tree1.priority);
		lastPriority=tree1.priority;
#endif

		newParent.node=new StaticHuffmanNode;
		newParent.node->isLeaf=false;
		newParent.node->child[0]=tree0.node;
		newParent.node->child[1]=tree1.node;
		newParent.node->weight=tree0.node->weight + tree1.node->weight;
		newParent.node->parent=NULL;
		tree0.node->parent=newParent.node;
		tree1.node->parent=newParent.node;
		newParent.priority=newParent.node->weight;
		heap.Insert(newParent);

#if _DEBUG
		loops++;
#endif
	}
}

void StaticHuffmanTree::FindEncodingsFromTreeStrucure (void)
{
	int i;
	unsigned long length;
	for (i=0; i<257; i++) {

		length=0;

		// walk up the tree to get the sequence in reverse
		StaticHuffmanNode *node,*parent;
		node=codes[i].node;
		while (node->parent!=NULL) {

			assert(length<256); // otherwise we were wrong about the maximum

			parent=node->parent;
			if (parent->child[0]==node) {
				codes[i].sequence[length]=0;
			}
			else {
				assert(parent->child[1]==node);
				codes[i].sequence[length]=1;
			}
			length++;

			node=parent;
		}

		codes[i].length=length;
		codes[i].ReverseSequence();
	}
}

void StaticHuffmanTree::WriteEncodedFile(FILE *normalFile,FILE *compressFile)
{
	rewind(normalFile);
	while (!feof(normalFile)) {
		int ch=fgetc(normalFile);
		if (ch==-1) {
			if (feof(normalFile)) {
				break;
			}

			// using exceptions might be nice...
			printf("Error while reading normal file.  Didn't expect that.\n");
			exit(0);
		}
		assert(ch>=0 && ch<256);

		codes[ch].WriteBits(compressFile);
	}

	// write end-of-file
	codes[256].WriteBits(compressFile);
}

void StaticHuffmanTree::Encode (FILE *normalFile,FILE *compressFile)
{
	InitBitStream ();

	// we start with a bunch of trees - in the codes member variable

	CountFrequencies(normalFile);

	// put the frequencies in the Heap & build the full static Huffman tree
	BuildTreeFromForest();

	// make the code sequences
	FindEncodingsFromTreeStrucure();
	SanityCheck();

	// write the tree into the file
	WriteTreeToFile(compressFile);
	SanityCheck();

	long treeLen=ftell(compressFile);

	// and write the actual file
	WriteEncodedFile(normalFile,compressFile);

	long fileLen=ftell(compressFile);
	printf("Compressed: %ld (Tree: %ld, Other: %ld)\n",
		fileLen,
		treeLen,
		fileLen-treeLen);

	// make sure it's there
	FlushBits(compressFile);
	fflush(compressFile);
}

void StaticHuffmanTree::Decode (FILE *compressFile,FILE *normalFile)
{
	InitBitStream ();

	// get the tree back
	ReadTreeFromFile(compressFile);

	// and read it

	StaticHuffmanNode *node;

	node=root;

	while (1) {
		int aBit;
		while (!node->isLeaf) {

			aBit=ReadOneBit(compressFile);
			switch (aBit) {
			case 0:
				node=node->child[0];
				break;
			case 1:
				node=node->child[1];
				break;
			}
		}

		if (node->theChar==256) {
			// end-of-file character
			break;
		}

		fputc(node->theChar,normalFile);

		node=root;
	}
}

void StaticHuffmanTree::SanityCheck (void)
{
	// verify that the code sequences are correct
	unsigned long i,l;
	for (i=0; i<257; i++) {
		StaticHuffmanNode *node;

		node=root;
		for (l=0; l<codes[i].length; l++) {
			assert(!node->isLeaf);
			assert(node->child[0]->parent==node);
			assert(node->child[1]->parent==node);
			assert(node->weight==node->child[0]->weight + node->child[1]->weight);
			node=node->child[codes[i].sequence[l]];
		}
		assert(node->isLeaf);
		assert(node->theChar==(unsigned int)i);
	}
}