#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
#include <queue>
#include <map>

#ifdef __GNUC__
// c.f. http://gnuwin32.sf.net/compile.html
#include <fcntl.h>
int _CRT_fmode = _O_BINARY;
#endif

// Bitstream
// `````````

struct bstream {
	FILE* file;
	char cache;
	unsigned int cache_pos;
	bool at_eof;
	
	bstream(FILE* file) :
		file(file), cache(0), cache_pos(7), at_eof(false)
	{};
	
	~bstream() {	
		if (cache_pos != 7) flush();
	}
	
	void write(bool bit) {
		cache |= (bit?1:0) << cache_pos;
		if (cache_pos-- == 0) flush();
	};
	
	void write8(char c) {
		for (int i = 8; i--> 0;) {
			write( c & (1 << i));
		}
	}
	
	void writeb(std::string& bstr) {
		for (int i = 0, e = bstr.length(); i < e; i++) {
			if (bstr[i] == '0') {
				write(0);
			} else if (bstr[i] == '1') {
				write(1);
			}
		}
	}
	
	void flush() {
		fputc( (int)cache, file );
		cache = 0x00;
		cache_pos = 7;
	};
	
	bool read() {
		if (cache_pos == 7) {
			if (at_eof) {
				cache = 0x00;
			} else {
				int incoming = fgetc(file);
				if (incoming == EOF) {
					at_eof = true;
					cache = 0x00;
				} else {
					cache = (char)incoming;
				}
			}
		}
		bool ret = (cache & (1 << cache_pos));
		cache_pos = (cache_pos + 7) % 8; // -1 %8, but always positive
		return ret;
	};
	
	char read8() {
		char ret = 0;
		for (int i = 8; i--> 0;) {
			if (this->read()) ret |= (1<<i);
		}
		return ret;
	}
};

std::string stdin_get_contents() {
	std::freopen(NULL, "rb", stdin);

	std::string ret;
	int c;
	while((c = std::getchar()) != EOF) ret += (char)c;
	
	fprintf(stderr, "Read %d bytes\n", ret.length());
	
	return ret;
}

// Huffman Tree
// ````````````

enum node_types {
	NODE_TYPE_INVALID,
	NODE_TYPE_LITERAL,
	NODE_TYPE_PTR
};

struct huff_node {
	
	enum node_types type;
	unsigned long long weight;
	union {
		unsigned int literal;
		struct {
			huff_node *left, *right;
		} children;
	};
	
	huff_node(enum node_types type) : type(type) {};
	huff_node() {};
	
	huff_node(const huff_node& other) {
		this->type = NODE_TYPE_INVALID; // TODO
	}
	
	huff_node& operator= (const huff_node& other) {
		type = other.type;
		weight = other.weight;
		if (type == NODE_TYPE_LITERAL) {
			literal = other.literal;
		} else {
			children.left = other.children.left;
			children.right = other.children.right;
		}
		return *this;
	}
	
	~huff_node() {
		if (type == NODE_TYPE_PTR) {
			delete children.left;
			delete children.right;
		}
	}
	
	void print(int indent = 0) {
		for(int i = indent; i-->0;) printf("  ");
		printf("%d", weight);
		if (type == NODE_TYPE_LITERAL) {
			printf(": \\x%02X\n", literal);
		} else {
			printf("\n");
			children.left ->print(indent + 1);
			children.right->print(indent + 1);
		}
	}
	
	static huff_node* build(std::string& data);
};

struct huff_compare {
	bool operator() (const huff_node* a, const huff_node* b) {
		return a->weight > b->weight;
	}
};

huff_node* huff_node::build(std::string& data) {
		
	#define NUM_POSSIBILITIES (256+1)
		
	// Create nodes for each possible octet
	std::vector<huff_node*> prob(NUM_POSSIBILITIES); // maintain for raw storage
	for (int i = 0; i < NUM_POSSIBILITIES; i++) {
		prob[i] = new huff_node(NODE_TYPE_LITERAL);
		prob[i]->literal = (char)i;
		prob[i]->weight  = 0;
	}
	prob[256] = new huff_node(NODE_TYPE_LITERAL);
	prob[256]->literal = EOF;
	prob[256]->weight  = 1;
	
	// Retrieve all probabilities
	for (int i = 0, e = data.length(); i < e; i++) {
		prob[ (unsigned int)(data[i]) ]->weight++;
	}
	
	// Sort and move into queue
	std::priority_queue<huff_node*, std::vector<huff_node*>, huff_compare> q;
	for (int i = 0; i < NUM_POSSIBILITIES; i++) {	
		if (prob[i]->weight) q.push(prob[i]);
	}
	
	// Process queue into tree
	while (q.size() > 1) {
		huff_node* h = new huff_node(NODE_TYPE_PTR);
		h->children.left  = q.top(); q.pop();
		h->children.right = q.top(); q.pop();
		h->weight = h->children.left->weight + h->children.right->weight;
		q.push(h);
	}
	
	return q.top();
}

void huff_tree_write(huff_node* n, bstream& b) {
	if (n->type == NODE_TYPE_LITERAL) {
		b.write(1);
		
		// 8 bits for a numeric (if the character is \0, follow with an extra
		//    bit as a tagged character for EOF)
		if (n->literal == 0 || n->literal == EOF) {
			b.write8('\0');
			b.write( n->literal == EOF );		
			
		} else {
			b.write8(n->literal);
		}
		
		
	} else {
		b.write(0);
		huff_tree_write(n->children.left, b);
		huff_tree_write(n->children.right, b);
	}
}

huff_node* huff_tree_read(bstream& b) {
	huff_node* ret = new huff_node();
	ret->weight = 0;
	if (b.read()) {
		// literal
		ret->type = NODE_TYPE_LITERAL;
		ret->literal = b.read8();
		if (ret->literal == 0) {
			bool flag = b.read(); // take one extra bit
			if (flag) {
				ret->literal = EOF;
			}
		}
	} else {
		// tree
		ret->type = NODE_TYPE_PTR;
		ret->children.left  = huff_tree_read(b);
		ret->children.right = huff_tree_read(b);
	}
	return ret;
}

void huff_buildmap(
	huff_node* n, std::string current_prefix, std::map<char, std::string>& out
) {
	// Walk tree
	if (n->type == NODE_TYPE_LITERAL) {
		out[n->literal] = current_prefix;
	} else {
		huff_buildmap(n->children.left,  current_prefix + "0", out);
		huff_buildmap(n->children.right, current_prefix + "1", out);
	}
}

void huff_pack(huff_node* n, std::string& str, bstream& b) {
	// Write out lookup tree
	huff_tree_write(n, b);
	
	// Build map of characters to binary strings
	std::map<char, std::string> lookup;
	huff_buildmap(n, "", lookup);	
		
	// Walk str and output to b
	for (int i = 0, e = str.length(); i < e; i++) {
		b.writeb( lookup[ str[i] ] );
	}
	b.writeb( lookup[ EOF ] );
}

void huff_unpack(huff_node* n, bstream& b) {
	// We have a tree, and bstream is the data to decode
	
	huff_node* curpos = n;
	while (! b.at_eof) {
		bool bit = b.read();
		if (bit) {
			curpos = curpos->children.right;
		} else {
			curpos = curpos->children.left;
		}
		
		if (curpos->type == NODE_TYPE_LITERAL) {
			if (curpos->literal == EOF) return;
			
			fputc(curpos->literal, stdout);
			curpos = n;
			
		} else if (curpos->type == NODE_TYPE_INVALID) {
			fputs("reached invalid node", stderr);
			return;
		}
		
	}
}

// Entry point
// ```````````

int usage() {
	puts("Usage: huffenc e[ncode] | d[ecode] | p[rettyprint] | l[oad-tree]");
	return 1;
}

int main(int argc, char* argv[]) {

	if (argc != 2) {
		return usage();
	} else {
		
		// Encode text on stdin.
		if (argv[1][0] == 'e') {
			std::string message = stdin_get_contents();
			huff_node* n = huff_node::build(message);
			bstream b(stdout);
			// Write out magic
			b.write8('M');
			b.write8('H');
			huff_pack(n, message, b);
			delete n;
			return 0;
	
		// Generate and pretty-print a tree from text on stdin.
		} else if (argv[1][0] == 'p') {
			std::string message = stdin_get_contents();
			printf("%s", message.c_str());
			return 1;
			
			huff_node* n = huff_node::build(message);
			n->print();

			delete n;
			return 0;
		
		
		// Pretty-print a tree from data on stdin.
		} else if (argv[1][0] == 'l') {
			bstream b(stdin);
			if (b.read8() == 'M' && b.read8() == 'H') {
				huff_node* n = huff_tree_read(b);
				n->print();
				delete n;
				return 0;
				
			} else {
				fputs("Invalid magic header.", stderr);
				return 1;
			}
		
		// Decode data from stdin.
		} else if (argv[1][0] == 'd') {
			bstream b(stdin);
			if (b.read8() == 'M' && b.read8() == 'H') {
				huff_node* n = huff_tree_read(b);
				huff_unpack(n, b);				
				delete n;
				return 0;
				
			} else {
				fputs("Invalid header.", stderr);
				return 1;
			}
		
		// Usage
		} else if (argv[1][0] == 'h') {
			return usage();
		
		} else {
			puts("Unknown command");
			return 1;
		}
	}
}
