// NOTE: this code was written for educational purposes, so excuse the multitude of comments

#include <iostream>
#include <stdlib.h>
#include <stdio.h>
#include <cstring>
#include <cstdint>
#include <string>

using namespace std;

enum {DIR_LEFT, DIR_RIGHT};
enum {DATA_FREQ, DATA_CHAR, DATA_BALANCE, DATA_CHILD_TYPE};
enum {BALANCE_SUB = -1, BALANCE_ADD = 1};
enum {LESS_CHILD, GREAT_CHILD};
enum {ROT_RIGHT, ROT_LEFT}; // it makes sense to define the enum like this, a right bias (DATA_BALANCE > 0) is corrected with a left rotation
enum {ROT_RIGHT_LEFT, ROT_LEFT_RIGHT}; 
enum {ENCODE_BIN = DATA_CHILD_TYPE, ENCODE_LEN = DATA_BALANCE};

string stringToCompress = "An idiot admires complexity, a genius admires simplicity.";

int uniqueChars;

void print_bin(int hex, int len)
{
	for (int i = len - 1; i >= 0; i--)
		if ((hex >> i) & 0x01)
			printf("1");
		else
			printf("0");
}

struct avlNode
{
	int data[4];
	avlNode* parent = NULL;
	avlNode* children[2] = {NULL, NULL}; // to designate nodes without children
};

void insert_avl_child(avlNode* nodeToInsert, avlNode* rootNode, bool childType) // only concerned with updating pointers to make it more general purpose
{
	rootNode->children[childType] = nodeToInsert;
	nodeToInsert->parent = rootNode;
	
	nodeToInsert->data[DATA_CHILD_TYPE] = childType;
}

avlNode* freqRoot = new(avlNode); // Originally, I kept track of character frequencies in an array equal in size to the total number of possible characters 
		   		  // (mitigated somewhat by excluding the first 32 entries in the ascii table) which ultimately yielded time and space complexities of O(1).
	         		  // The big downside was the absolute space cost of 2^b for strings of words b bits in size; this produced many blank frequency entries 
		   		  // that would incur their own time cost to deal with. Using an avl tree to track duplicates introduces a space complexity of O(n) and a 
		   		  // time complexity of O(log(n)) for strings with n unique characters with the upside being the dramatically lower absolute time and space costs.

void avl_rotate_balancing(avlNode* newRoot, bool single, bool rotType)
{
	switch ((single << 1) | (newRoot->data[DATA_BALANCE] != 0))
	{
		case 0:	// double, new root's old balance zero	
			newRoot->children[GREAT_CHILD]->data[DATA_BALANCE] = 0;
			newRoot->children[LESS_CHILD]->data[DATA_BALANCE] = 0;
			newRoot->data[DATA_BALANCE] = 0;			
			break;

		case 1:	// double, new root's old balance non zero
			if (newRoot->data[DATA_BALANCE] > 0)
			{
				newRoot->children[GREAT_CHILD]->data[DATA_BALANCE] = 0;
				newRoot->children[LESS_CHILD]->data[DATA_BALANCE] = -1;			
			}
			else
			{
				newRoot->children[GREAT_CHILD]->data[DATA_BALANCE] = 1;
				newRoot->children[LESS_CHILD]->data[DATA_BALANCE] = 0;	
			}
			newRoot->data[DATA_BALANCE] = 0;		
			break;
		
		case 2:	// single, new root's old balance zero
			if (rotType)
			{
				newRoot->data[DATA_BALANCE] = -1;
				newRoot->children[!rotType]->data[DATA_BALANCE] = 1;		
			}
			else
			{
				newRoot->data[DATA_BALANCE] = 1;
				newRoot->children[!rotType]->data[DATA_BALANCE] = -1;
			}
			break;

		case 3:	// single, new root's old balance non zero
			newRoot->data[DATA_BALANCE] = 0;
			newRoot->children[!rotType]->data[DATA_BALANCE] = 0;	
			break;

		default:
			break;
	}

	if (newRoot->parent == NULL)
		freqRoot = newRoot;
}

avlNode* avl_rotate_single(avlNode* unbalNode, bool rotType)
{
	avlNode* parent = unbalNode->parent;
	avlNode* neighbor = unbalNode->children[rotType]; // REMEMBER, ROT_LEFT yields a rotType of 1, despite what your most 'basic instincts' might tell you

	if (parent != NULL)
		insert_avl_child(neighbor, parent, unbalNode->data[DATA_CHILD_TYPE]);
	else
		neighbor->parent = NULL;

	avlNode* neighborChild = neighbor->children[!rotType];
	
	if (neighborChild != NULL)
		insert_avl_child(neighborChild, unbalNode, rotType);
	else
		unbalNode->children[rotType] = NULL; // rotation causes the child slot formerly occupied by neighbor to now be empty in the absence of an 'heir'
	
	insert_avl_child(unbalNode, neighbor, !rotType);
	
	return neighbor;
}

avlNode* avl_rotate_double(avlNode* unbalNode, bool rotType)
{
	avl_rotate_single(unbalNode->children[!rotType], rotType);
	return avl_rotate_single(unbalNode, !rotType);
}

void avl_rebalance_tree(avlNode* unbalNode)
{
	bool rootBias = (unbalNode->data[DATA_BALANCE] > 0); 
	int childBias = unbalNode->children[rootBias]->data[DATA_BALANCE];

	if (childBias == 0)
	{
		avl_rotate_single(unbalNode, rootBias);
		return;
	}
	else
	{
		switch ((rootBias << 1) | (childBias > 0))
		{
			case 0: // root and child both have left bias
				avl_rotate_balancing(avl_rotate_single(unbalNode, ROT_RIGHT), true, ROT_RIGHT);
				break;

			case 1: // root has left bias and child has right bias 
				avl_rotate_balancing(avl_rotate_double(unbalNode, ROT_LEFT_RIGHT), false, ROT_LEFT_RIGHT);
				break;

			case 2: // root has right bias and child has left bias
				avl_rotate_balancing(avl_rotate_double(unbalNode, ROT_RIGHT_LEFT), false, ROT_RIGHT_LEFT);
				break;

			case 3: // root and child both have right bias
				avl_rotate_balancing(avl_rotate_single(unbalNode, ROT_LEFT), true, ROT_LEFT);
				break;

			default:
				break;
		}
	}
}

void update_parent_balances(avlNode* balNode, int isAddition)
{
	avlNode* currNode = balNode->parent;

	if (currNode == NULL)
		return;

	currNode->data[DATA_BALANCE] += isAddition * ((balNode->data[DATA_CHILD_TYPE] << 1) - 1);

	if ((currNode->data[DATA_BALANCE] == 2) || (currNode->data[DATA_BALANCE] == -2))
		avl_rebalance_tree(currNode);
	
	switch (((currNode->data[DATA_BALANCE] != 0) << 1) | isAddition) // it's helpful to think of isAddition as just telling us if its associated subtree
	{ 							         // is gaining or losing depth, not referring to literal insertions or deletions
		case 0: // subtraction and smaller abs balance		 // the only legal way for a node to have a balance factor of zero after one of its		      
			update_parent_balances(currNode, BALANCE_SUB);	 // subtrees changes in depth is for the balance factor's absolute value to decrease
			break;

		case 1: // addition and smaller abs balance
			// no change in subtree depth
			break;

		case 2: // subtraction and larger abs balance
			// no change in subtree depth
			break;

		case 3: // addition and larger abs balance
			update_parent_balances(currNode, BALANCE_ADD);
			break;

		default:
			break;
	}
}

avlNode* search_avl(int checkVal, avlNode* rootNode, bool compType, bool freqCount = false)
{	
	if (checkVal == rootNode->data[compType]) // found a match!
	{
		if (freqCount)
			rootNode->data[DATA_FREQ]++;
		return rootNode;
	}
	
	bool comp = (checkVal > rootNode->data[compType]);

	if (rootNode->children[comp] == NULL) // didn't find a match...
	{
		if (freqCount) // ...but if you want we can append it to the tree
		{
			avlNode* insertNode = new(avlNode);
			insert_avl_child(insertNode, rootNode, comp);

			insertNode->data[DATA_BALANCE] = 0;
			insertNode->data[DATA_CHAR] = checkVal;
			insertNode->data[DATA_FREQ] = 1;

			update_parent_balances(insertNode, BALANCE_ADD);
			uniqueChars++;

			return insertNode;
		}
		
		return NULL;
	}
	else
		return search_avl(checkVal, rootNode->children[comp], compType, freqCount); // let's search the next subtree our node could be in
}

struct llNode
{
	int data[2];
	llNode* neighbors[2] = {NULL, NULL}; // to designate nodes at the ends of the linked list	
};

llNode* charMsort;

void remove_llnode(llNode* nodeToRemove)
{
	if (nodeToRemove->neighbors[DIR_RIGHT] != NULL)
		nodeToRemove->neighbors[DIR_RIGHT]->neighbors[DIR_LEFT] = nodeToRemove->neighbors[DIR_LEFT];

	
	if (nodeToRemove->neighbors[DIR_LEFT] != NULL)
		nodeToRemove->neighbors[DIR_LEFT]->neighbors[DIR_RIGHT] = nodeToRemove->neighbors[DIR_RIGHT];
		
	nodeToRemove->neighbors[DIR_LEFT] = NULL;
	nodeToRemove->neighbors[DIR_RIGHT] = NULL;
}

void delete_llnode(llNode* nodeToDelete)
{
	remove_llnode(nodeToDelete);
	delete(nodeToDelete);
}

void insert_llnode(llNode* nodeToInsert, llNode* indexNode, bool toRight)
{
	nodeToInsert->neighbors[toRight] = indexNode->neighbors[toRight];

	if (indexNode->neighbors[toRight] != NULL)
		indexNode->neighbors[toRight]->neighbors[!toRight] = nodeToInsert;	// if toRight is false, the function instead inserts a node to the left of the index node

	nodeToInsert->neighbors[!toRight] = indexNode;
	indexNode->neighbors[toRight] = nodeToInsert;
}

avlNode* decodeTree;

void add_node_to_decode_tree(llNode* mergeNode, int freeNodeIndex, bool insertDir)
{
	if (mergeNode->data[DATA_CHAR] < 0) // a single bit to denote a tree index entry won't hurt, right?
		decodeTree[freeNodeIndex].children[insertDir] = &decodeTree[(-1 * mergeNode->data[DATA_CHAR]) - 1];
	else
		decodeTree[freeNodeIndex].data[insertDir] = mergeNode->data[DATA_CHAR]; // if a LESS_CHILD or a GREAT_CHILD is NULL, then data[0] or data[1], repsectively, contains a terminal character
}

llNode* merge_llnodes(llNode* mergeNode1, llNode* mergeNode2) // net decrease in the size of the list is -1, so no special considerations for even or odd list sizes need to be taken unless the list is only one entry long (how ridiculous!)
{
	static int freeNodeIndex = 0;

	int node1Freq = mergeNode1->data[DATA_FREQ];
	int node2Freq = mergeNode2->data[DATA_FREQ];
	llNode* returnNode;

	remove_llnode(mergeNode1); // this allow the exit condition for the while() loop over in generate_decode_tree() to work properly and mergeNode2's deletion to not mess with mergeNode1's pointers

	add_node_to_decode_tree(mergeNode1, freeNodeIndex, LESS_CHILD);
	add_node_to_decode_tree(mergeNode2, freeNodeIndex, GREAT_CHILD);

	int freqSum = node1Freq + node2Freq;
	
	mergeNode1->data[DATA_CHAR] = -1 * (freeNodeIndex + 1);
	mergeNode1->data[DATA_FREQ] = freqSum;
	freeNodeIndex++;

	if (mergeNode2->neighbors[DIR_RIGHT] != NULL) // if we just merged the last two nodes in the list, the insertion part of this function is rendered unnecessary
	{
		llNode* index = mergeNode2;

		while ((index->neighbors[DIR_RIGHT] != NULL) && (index->neighbors[DIR_RIGHT]->data[DATA_FREQ] < freqSum))
			index = index->neighbors[DIR_RIGHT];

		insert_llnode(mergeNode1, index, DIR_RIGHT);

		returnNode = mergeNode2->neighbors[DIR_RIGHT];
	}
	else
		returnNode = mergeNode1;
	
	delete_llnode(mergeNode2);

	return returnNode;
}

int currIndex;

void copy_chars_into_list(avlNode* root)
{
	if (root == NULL)
		return;

	charMsort[currIndex].neighbors[0] = new(llNode);
	charMsort[currIndex].data[0] = 1;
	charMsort[currIndex].neighbors[0]->data[DATA_FREQ] = root->data[DATA_FREQ];				
	charMsort[currIndex].neighbors[0]->data[DATA_CHAR] = root->data[DATA_CHAR];				
	if (currIndex != (uniqueChars - 1))
		charMsort[currIndex].neighbors[1] = &charMsort[currIndex + 1];
	currIndex++;	

	copy_chars_into_list(root->children[LESS_CHILD]);
	copy_chars_into_list(root->children[GREAT_CHILD]);
}

void init_list()
{
	freqRoot->data[DATA_BALANCE] = 0;
	freqRoot->data[DATA_CHAR] = (int)stringToCompress.at(0);
	freqRoot->data[DATA_FREQ] = 1;
	uniqueChars = 1;

	for (int i = 1; i < stringToCompress.length(); i++)	
		search_avl((int)stringToCompress.at(i), freqRoot, DATA_CHAR, true);
		
	charMsort = (llNode*)malloc(uniqueChars * sizeof(llNode)); 

	currIndex = 0;
	copy_chars_into_list(freqRoot);
}

void merge_lists(llNode* list1Meta, llNode* list2Meta)
{
	int l1Size = list1Meta->data[0];
	int l2Size = list2Meta->data[0];

	list1Meta->data[0] = l1Size + l2Size;

	llNode* l1 = list1Meta->neighbors[0];
	llNode* l2;
	llNode* l2Temp;
	
	int newl2Size = l2Size;

	for (int i = 0; i < l1Size; i++)
	{
		l2 = list2Meta->neighbors[0];
		
		for (int j = 0; j < l2Size; j++) // let's search for entries in the second list less than the current l1
		{
			l2Temp = l2->neighbors[DIR_RIGHT];
			
			if (l2->data[DATA_FREQ] < l1->data[DATA_FREQ])
			{
				if (l2 == list2Meta->neighbors[0])
					list2Meta->neighbors[0] = list2Meta->neighbors[0]->neighbors[DIR_RIGHT]; // update starting pointer for l2's list
			
				remove_llnode(l2);
				insert_llnode(l2, l1, DIR_LEFT);
				newl2Size--;

				if (l1 == list1Meta->neighbors[0])
					list1Meta->neighbors[0] = list1Meta->neighbors[0]->neighbors[DIR_LEFT]; // update l1's list starting pointer if need be
			}
			else		// no point in searching anymore, everything from here on out is going to be equal to or larger than the current l1
				break;
			
			l2 = l2Temp;
		}
		
		if (i != (l1Size - 1))
			l1 = l1->neighbors[DIR_RIGHT]; // preserve l1 in case l2Size isn't zero by the end of the loop
		
		l2Size = newl2Size;
	
		if (l2Size == 0) // l1 already contains the complete sorted list if l2 is empty
			break;
	}

	if (l2Size != 0) // append the values of l2 larger than the greatest value in l1 to the end of the latter
	{
		l2 = list2Meta->neighbors[0];

		for (int j = 0; j < l2Size; j++)
		{
			l2Temp = l2->neighbors[DIR_RIGHT];

			insert_llnode(l2, l1, DIR_RIGHT);
			l1 = l1->neighbors[DIR_RIGHT];	

			l2 = l2Temp;
		}
	}

	list1Meta->neighbors[1] = list2Meta->neighbors[1];
}

void msort()
{
	llNode* currList;
	llNode* rightNeighbor;

	while (charMsort->neighbors[1] != NULL)
	{
		currList = charMsort;
		rightNeighbor = charMsort->neighbors[1];

		while(rightNeighbor != NULL)
		{
			merge_lists(currList, rightNeighbor);
			
			if (currList->neighbors[1] != NULL)		// this part doesn't use rightNeighbor as currList's neighbors[1] was altered in merge_lists()
				currList = currList->neighbors[1];
			
			rightNeighbor = currList->neighbors[1];	// do we still have another node left to sort after jumping?
		}
	}
}

void generate_decode_tree()
{
	decodeTree = (avlNode*)malloc((uniqueChars - 1) * sizeof(avlNode));

	llNode* firstNode = charMsort->neighbors[0];

	while (firstNode->neighbors[DIR_RIGHT] != NULL)
		firstNode = merge_llnodes(firstNode, firstNode->neighbors[DIR_RIGHT]); // we generate a max heap tree where only terminal leaves are associated with characters as a consequence of huffman being a prefix code

	delete(firstNode); // clean up
}

int compressedSize = 0;

void generate_encode_tree(avlNode* root = &decodeTree[uniqueChars - 2], int encodeBin = 0, int encodeLen = 0)
{
	static avlNode* curr;

	if (root->children[LESS_CHILD] != NULL)
	{
		generate_encode_tree(root->children[LESS_CHILD], encodeBin << 1, encodeLen + 1);
	}
	else
	{
		curr = search_avl(root->data[0], freqRoot, DATA_CHAR);
		compressedSize += (encodeLen + 1) * curr->data[DATA_FREQ];
	
		curr->data[ENCODE_BIN] = encodeBin << 1;	// already a zero in this spot, bitwise OR not required
		curr->data[ENCODE_LEN] = encodeLen + 1;

		printf("%c - ", curr->data[DATA_CHAR]);
		print_bin(curr->data[ENCODE_BIN], curr->data[ENCODE_LEN]);
		printf("\n");
	}
	
	if (root->children[GREAT_CHILD] != NULL)
	{
		generate_encode_tree(root->children[GREAT_CHILD], (encodeBin << 1) | 0x01, encodeLen + 1);
	}
	else
	{
		curr = search_avl(root->data[1], freqRoot, DATA_CHAR);
		compressedSize += (encodeLen + 1) * curr->data[DATA_FREQ];

		curr->data[ENCODE_BIN] = (encodeBin << 1) | 0x01;
		curr->data[ENCODE_LEN] = encodeLen + 1;

		printf("%c - ", curr->data[DATA_CHAR]);
		print_bin(curr->data[ENCODE_BIN], curr->data[ENCODE_LEN]);
		printf("\n");
	}
}

void print_message()
{
	printf("original message (%d bits in size):\n%s\n\n", stringToCompress.length() * 7, stringToCompress.c_str());
}

uint8_t* binaryStream;

void print_stream()
{
	for (int i = 0; i < compressedSize; i++)
	{
		if ((binaryStream[i / 8] >> (7 - (i % 8))) & 0x01)
			printf("1");
		else
			printf("0");

		if ((((i + 1) % 40) == 0) && (i != (compressedSize - 1))) // the last condition prevents an extra \n from being added to the output if the last character to be printed on the current line is also the last character of the stream
			printf("\n");
	}
}

void insert_token(int token, int tokenLen)
{
	static int bitsLeft = 8;
	static int currentInt = 0;
	static int temp;

	if (tokenLen > bitsLeft)
	{
		temp = tokenLen - bitsLeft;

		binaryStream[currentInt] |= (token >> temp) & 0xff; // the & 0xff isn't really necessary here, but having it there makes me feel better, kinda like a teddy bear
		
		bitsLeft = 8;
		currentInt++;
		insert_token(token & ((1 << temp) - 1), temp);
	}
	else // degenerate case, your function's over mutie, time to ret.
	{
		binaryStream[currentInt] |= (token << (bitsLeft - tokenLen)) & 0xff;
		
		bitsLeft -= tokenLen;

		if (bitsLeft == 0)
		{
			bitsLeft = 8;
			currentInt++;
		}
	}
}

void encode_message()
{
	binaryStream = (uint8_t*)calloc((compressedSize / 8) + ((compressedSize % 8) > 0), sizeof(uint8_t));
	
	avlNode* curr;

	for (int i = 0; i < stringToCompress.length(); i++)
	{
		curr = search_avl((int)stringToCompress.at(i), freqRoot, DATA_CHAR);

		insert_token(curr->data[ENCODE_BIN], curr->data[ENCODE_LEN]);
	}
}

string decode_message()
{
	avlNode* curr = &decodeTree[uniqueChars - 2];
	avlNode* next;
	string ret;

	bool dir;

	for (int i = 0; i < compressedSize; i++)
	{
		dir = (binaryStream[i / 8] >> (7 - (i % 8))) & 0x01;
		next = curr->children[dir];
	
		if (next == NULL)
		{
			ret += (char)curr->data[dir];
			curr = &decodeTree[uniqueChars - 2];
		}
		else
			curr = next;
	}

	return ret;
}

void annihilate_avl(avlNode* root)
{
	if (root == NULL)
		return;

	annihilate_avl(root->children[LESS_CHILD]);	
	annihilate_avl(root->children[GREAT_CHILD]);

	delete(root);
}

void free_mem()
{
	free(charMsort);
	free(binaryStream);
	free(decodeTree);

	annihilate_avl(freqRoot);
}

int main ()
{
	printf("\n");
	print_message();

	init_list();
	msort();
	generate_decode_tree();

	printf("translation table produced by counting duplicate characters and assigning codewords with fewer bits to frequently used characters:\n");
	generate_encode_tree();

	printf("\ncompressed data stream (%d bits in size):\n", compressedSize);
	encode_message();
	print_stream();

	printf("\n\noriginal message as decoded from the compressed data stream:\n%s", decode_message().c_str());

	printf("\n\na compression ratio of %.4f was achieved", (float)(stringToCompress.length() * 7) / (float)compressedSize);
	
	printf("\n\n");

	free_mem();
}
