#include <math>
#include <stdlib.h>
#include <fstream>
#include <iostream>
#include <map>
#include <deque>
#include <string>
#include <queue>
#include <functional>
using namespace std;

class ibstream {
public:
	ibstream(istream& is);
	int readBit();
	istream& getInputStream() const;
private:
	istream& in;
	char buffer;
	int bufferPos;
};


class obstream {
public:
	obstream(ostream& os);
	~obstream();
	void writeBit(int val);
	void writeBits(const deque<int>& val);
	ostream& getOutputStream() const;
	void flush();
private:
	ostream& out;
	char buffer;
	int bufferPos;
};

class CharCounter {
public:
	CharCounter();
	CharCounter(istream& input);
	int getCount(char ch);
	void setCount(char ch, int count);
private:
	map<char, int> theCounts;
};


class HuffmanTree {
public:
	HuffmanTree();
	HuffmanTree(const CharCounter& cc);
	enum {ERROR=-3, INCOMPLETE_CODE=-2, END=-1};
	deque<int> getCode(int ch);
	int getChar(const deque<int>& code) const;
	void writeEncodingTable(ostream& out);
	void readEncodingTable(istream& in);
private:
	class HuffNode {
	public:
		 HuffNode* left;
		 HuffNode* right;
		 HuffNode* parent;
		 int value;
		 int weight;
		 HuffNode(int v, int w, HuffNode* lt, HuffNode* rt, HuffNode* pt):
				value(v), weight(w), left(lt), right(rt), parent(pt) {
		}
	};
	CharCounter theCounts;
	map<int, HuffNode*> theNodes;
	HuffNode* root;
	void createTree();
	deque<int> getCode(HuffNode* current);
friend bool greaterWeight(HuffNode* lhs, HuffNode* rhs);
};

class Compressor {
public:
	static void compress(const string& inFile);
	static void uncompress(const string& compressedFile);
};

static const int BITS_PER_CHAR = 8;
static const int DIFF_CHARS = 256;
static const int READ_MODE = ios::in | ios::binary;
static const int WRITE_MODE = ios::out | ios::binary;

int getBit(char pack, int pos) {
	return (pack&(1<<pos)) ? 1 : 0;
}

void setBit(char& pack, int pos, int val) {
	if (val==1)
		pack|=static_cast<char>(val<<pos);
}

ibstream::ibstream(istream& is): bufferPos(BITS_PER_CHAR), in(is) {
}

int ibstream::readBit() {
	if (bufferPos==BITS_PER_CHAR) {
		in.get(buffer);
		if(in.eof())
			return EOF;
		bufferPos=0;
	}
	return getBit(buffer, bufferPos++);
}

istream& ibstream::getInputStream() const {
	return in;
}


obstream::obstream(ostream& os): bufferPos(0), buffer(0), out(os) {
}

obstream::~obstream() {
	flush();
}

void obstream::flush() {
	if (bufferPos==0)
		return;
	out.put(buffer);
	bufferPos=0;
	buffer=0;
}

void obstream::writeBit(int val) {
	setBit(buffer, bufferPos++, val);
	if (bufferPos==BITS_PER_CHAR)
		flush();
}

void obstream::writeBits(const deque<int>& val) {
	for (deque<int>::size_type i(0); i<val.size(); ++i)
		writeBit(val[i]);
}

ostream& obstream::getOutputStream() const {
	return out;
}

CharCounter::CharCounter() {
}

CharCounter::CharCounter(istream& input) {
	char ch;
	while(!input.get(ch).eof())
		theCounts[ch]++;
}


int CharCounter::getCount(char ch) {
	return theCounts[ch];
}

void CharCounter::setCount(char ch, int count) {
	theCounts[ch]=count;
}


HuffmanTree::HuffmanTree(const CharCounter& cc): theCounts(cc), root(0) {
	createTree();
}


HuffmanTree::HuffmanTree(): root(0) {
}

deque<int> HuffmanTree::getCode(int ch) {
	if (root==0)
		return deque<int>();
	return getCode(theNodes[ch]);
}

deque<int> HuffmanTree::getCode(HuffNode* current) {
	deque<int> v;
	HuffNode* par(current->parent);
	while (par != 0) {
		if(par->left==current)
			v.push_front(0);
		else
			v.push_front(1);
		current=current->parent;
		par=current->parent;
	}
	return v;
}

int HuffmanTree::getChar(const deque<int>& code) const {
	HuffNode* p(root);
	for (deque<int>::size_type i(0); p!=0 && i<code.size(); ++i)
		if (code[i]==0)
			p=p->left;
		else
			p=p->right;
	if (p==0)
		return ERROR;
	return p->value;
}

void HuffmanTree::writeEncodingTable(ostream& out) {
	for (int i(0); i<DIFF_CHARS; ++i) {
		char ci(static_cast<char>(i));
		if (theCounts.getCount(ci)>0)
			out<<ci<<theCounts.getCount(ci)<<endl;
	}
	out<<'\0'<<0<<endl;
}

void HuffmanTree::readEncodingTable(istream& in) {
	for(int i(0); i<DIFF_CHARS; ++i)
		theCounts.setCount(static_cast<char>(i), 0);
	char ch;
	int num;
	char nl;
	while (true) {
		in.get(ch);
		in >> num;
		in.get(nl);
		if (num==0)
			break;
		theCounts.setCount(ch, num);
	}
	createTree();
}

bool greaterWeight(HuffmanTree::HuffNode* lhs, HuffmanTree::HuffNode* rhs) {
	return lhs->weight>rhs->weight;
}

void HuffmanTree::createTree() {
	typedef HuffNode N;
	priority_queue<N*, deque<N*>, bool (*)(N* lhs, N* rhs)> pq(greaterWeight);

	for (int i(0); i<DIFF_CHARS; ++i) {
		char ci(static_cast<char>(i));
		if (theCounts.getCount(ci)>0) {
			HuffNode* newNode(new HuffNode(i, theCounts.getCount(ci), 0, 0, 0));
			theNodes[i]=newNode;
			pq.push(newNode);
		}
	}
	theNodes[END]=new HuffNode(END, 1, 0, 0, 0);
	pq.push(theNodes[END]);

	while (pq.size()>1) {
		HuffNode* n1(pq.top());
		pq.pop();
		HuffNode* n2(pq.top());
		pq.pop();
		HuffNode* result(new HuffNode(INCOMPLETE_CODE, n1->weight+n2->weight, n1, n2, 0));
		n1->parent=n2->parent=result;
		pq.push(result);
	}
	root=pq.top();
}


void Compressor::compress(const string& inFile) {
	string compressedFile(inFile+".huf");
	ifstream in(inFile.c_str(), READ_MODE);
	CharCounter countObj(in);
	HuffmanTree codeTree(countObj);
	ofstream out(compressedFile.c_str(), WRITE_MODE);
	codeTree.writeEncodingTable(out);
	obstream bout(out);
	in.clear();
	in.seekg(0, ios::beg); // Rewind the stream
	char ch;
	while (in.get(ch))
		bout.writeBits(codeTree.getCode(static_cast<unsigned char>(ch)));
	bout.writeBits(codeTree.getCode(EOF));
}

void Compressor::uncompress(const string& compressedFile) {
	string::size_type indexPunt(compressedFile.rfind("."));
	string inFile(compressedFile.substr(0, indexPunt));
	string extension(compressedFile.substr(indexPunt));
	if (extension!=".huf") {
		cerr<<"Not a compressed file"<<endl;
		return;
	}
	inFile+=".uc";  // for debugging, so as to not clobber original
	ifstream in(compressedFile.c_str(), READ_MODE);
	ofstream out(inFile.c_str(), WRITE_MODE);
	HuffmanTree codeTree;
	codeTree.readEncodingTable(in);
	ibstream bin(in);
	deque<int> bits;
	while (true)
	{
		int bit(bin.readBit());
		bits.push_back(bit);
		int decode(codeTree.getChar(bits));
		if (decode==HuffmanTree::INCOMPLETE_CODE)
			continue;
		else if (decode==HuffmanTree::ERROR) {
			cerr<<"Error decoding!"<<endl;
			break;
		}
		else if(decode==HuffmanTree::END)
			break;
		else {
			out.put(static_cast<char>(decode));
			bits.resize(0);
		}
	}
}

int main(int argc, char* argv[]) {
	if(argc < 3) {
		cerr<<"Usage: "<<argv[0]<<" -[cu] files"<<endl;
		return 1;
	}
	string option(argv[1]);
	for (int i(2); i<argc; ++i) {
		string nextFile(argv[i]);
		if (option=="-c")
			Compressor::compress(nextFile);
		else if (option=="-u")
			Compressor::uncompress(nextFile);
		else {
			cerr<<"Usage: "<<argv[0]<<" -[cu] files"<<endl;
			return 1;
		}
	}
	return 0;
}