// КДЗ-1 по дисциплине Алгоритмы и структуры данных
// Ирхин Федор Денисович, БПИ154(1), 25.11.2016
// Visual Studio 2015, состав проекта (файлы *.cpp и *.h)
// Что сделано, а что нет

#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <algorithm>
#include <queue>
#include <iterator>
#include <stdexcept>
#include <assert.h>
#include <map>
#include <windows.h>
#include <clocale>
#include <codecvt>
#include <locale>

using namespace std;
struct Node
{
	bool isLeaf;
	wchar_t simbolChar;
	int freq;
	Node* left = nullptr;
	Node* right = nullptr;

	//Comparing operators for sort
	bool operator < (const Node& str) const
	{
		return (freq < str.freq);
	}
	bool operator > (const Node& str) const
	{
		return (freq > str.freq);
	}
};




//Header
void tableFilling(Node *root, std::wofstream& fout);
void treeDelete(Node *);
void writeCodedText(string pathFrom, std::wofstream& fout);
std::vector<Node*> nodesFromString(string path);

Node* buildTreeHaff(vector<Node> nodes);
Node* buildTreeSh(vector<Node> nodes);
Node* split(pair<int, Node*> left, pair<int, Node*> right);


std::wstring tableFreq[1024];

std::vector<Node*> nodesFromString(string path)
{
	std::locale locale(std::locale(), new std::codecvt_utf8<wchar_t>);
	std::wifstream fin(path, std::ios::binary);
	fin.imbue(locale);

	//Chechking is file exist
	if (!fin.is_open())
	{
		std::cout << "Input file not found" << std::endl;
		system("pause");
		throw std::logic_error("Input file not found");
	}

	//Vector of different Chars
	std::vector<wchar_t> charSet;
	//Vector of Nodes, which containt char and it's frequency
	std::vector<Node*> nodes;
	wchar_t currentChar;
	int j = 0;

	//Filling nodes vector, counting frequency of every node
	while (!fin.eof())
	{
		fin.get(currentChar);
		if (!fin.eof())
		{
			int i;
			for (i = 0; i < charSet.size() && charSet[i] != currentChar; i++);

			//Add new char
			if (i == charSet.size())
			{
				Node* tmp = new Node();

				charSet.push_back(currentChar);

				tmp->simbolChar = currentChar;
				tmp->freq = 0;
				tmp->isLeaf = true;

				nodes.push_back(tmp);

				nodes[j]->freq++;
				j++;
			}
			//Plus frequency if char is already exist
			else
			{
				Node tmp;

				tmp.simbolChar = currentChar;
				tmp.freq = 0;
				tmp.isLeaf = true;

				std::vector<Node*>::iterator it = nodes.begin();

				int i = 0;
				while (it < nodes.end())
				{
					if ((*it)->simbolChar == currentChar)
					{
						break;
					}
					i++;
					it++;
				}
				nodes[i]->freq++;
			}
		}
	}
	return nodes;
}
void decodeString(Node *root, string encoded_str, string outPath);
void decodeString(Node *root, string encoded_str, string outPath)
{
	wstring res = L"";
	Node* node = root;
	for (int i = 0; i != encoded_str.size(); ++i)
	{
		if (encoded_str[i] == '0') { // left child

			node = node->left;

		}
		else { // rigth child
			node = node->right;
		}
		if (node->isLeaf)
		{
			res += node->simbolChar;
			node = root;
		}
	}

	std::wofstream foutStream(outPath, std::ios_base::binary);
	foutStream << res;
}

void write_tree(wofstream& foutTree, Node* root);
void write_tree(wofstream& foutTree, Node* root)
{
	unsigned char buf = 0;
	unsigned char count = 0;
	bool bit;
	//	std::wofstream foutTree(path);
	if (root->isLeaf)
	{
		bit = 1;
		buf = buf | bit << (7 - count);
		count++;
		if (count > 7)
			count = 0;
		foutTree << '1';
		foutTree << root->simbolChar;
	}
	else
	{
		bit = 0;
		buf = buf | bit << (7 - count);
		count++;
		if (count > 7)
			count = 0;
		foutTree << '0';
		write_tree(foutTree, root->left);
		write_tree(foutTree, root->right);
	}
	//foutTree.close();
}

static string read_string(basic_ifstream<wchar_t>& fin)
{
	string str = "";
	wchar_t currentChar;


	while (!fin.eof())
	{
		fin.get(currentChar);
		if (!fin.eof())
			str += currentChar;
	}
	return str;
}
static Node* read_tree(basic_ifstream<wchar_t>& fin)
{
	wchar_t currentChar;
	//Filling nodes vector, counting frequency of every node
	if (fin.is_open());
	{
		fin.get(currentChar);
		if (currentChar == '0')
		{
			Node* left = read_tree(fin);
			Node* right = read_tree(fin);
			Node* newNode = new Node;
			newNode->left = left;
			newNode->right = right;
			newNode->isLeaf = false;
			return newNode;
		}
		else
		{
			if (currentChar == '1')
			{
				Node* newNode = new Node;

				if (true)
				{
					fin.get(currentChar);
					newNode->simbolChar = currentChar;
					newNode->freq = 0;
					newNode->isLeaf = true;
					return newNode;
				}
				else
				{
					throw std::logic_error("Tree reading went wrong");
				}
			}
		}
	}

}


Node* split(Node** left, Node** right)
{
	if (left == right)
	{
		(*left)->isLeaf = true;
		return *left;
	}
	int allSum = 0;
	for (Node** i = left; i <= right; i++)
		allSum += (*i)->freq;

	int halfSum = (*left)->freq;
	Node** nextRight = left + 1;
	while (abs(halfSum - allSum / 2 + (*nextRight)->freq) < abs(halfSum - allSum / 2))
	{
		halfSum += (*nextRight)->freq;
		nextRight++;
	}

	Node* rootNode = new Node;
	rootNode->left = split(left, nextRight - 1);
	rootNode->right = split(nextRight, right);
	rootNode->isLeaf = false;
	//	rootNode->simbolChar = (*left)->simbolChar;

	return rootNode;

}

bool Compare(const Node* a, const Node* b) { return a->freq < b->freq; }
Node* buildTreeSh(vector<Node*> nodes)
{
	std::sort(nodes.rbegin(), nodes.rend(), Compare);

	return split(&nodes[0], &nodes[nodes.size() - 1]);
}
Node* buildTreeHaff(vector<Node*> nodes)
{

	std::multimap<int, Node*> tree;
	for (int i = 0; i < nodes.size(); i++)
	{
		tree.insert(std::make_pair(nodes[i]->freq, nodes[i]));
	}
	//Writing Info
	//for (int i = 0; i < tree.size() - 1; i++)
	//	fout << "'" << tree[i]->simbolChar << "'" << tree[i]->freq << ", ";
	//fout << "'" << tree[tree.size() - 1]->simbolChar << "'" << tree[tree.size() - 1]->freq << std::endl << std::endl;

	//Building a tree
	while (tree.size() > 1)
	{
		std::map<int, Node*>::iterator it = tree.begin();
		Node* firstNode = it->second;
		tree.erase(it);

		it = tree.begin();
		Node* secondNode = it->second;
		tree.erase(it);

		Node* newNode = new Node();
		newNode->left = firstNode;
		newNode->right = secondNode;
		newNode->freq = firstNode->freq + secondNode->freq;

		tree.insert(std::make_pair(newNode->freq, newNode));
	}

	return tree.begin()->second;
}

//std::ofstream foutHaff("output.haff");
void codeHaff(string inPath);
void codeHaff(string inPath)
{
	string outputPath = inPath;
	outputPath = outputPath.erase(inPath.size() - 4, 4) + ".haff";

	std::vector<Node*> nodes = nodesFromString(inPath);
	Node* treeHaff = buildTreeHaff(nodes);
	std::wofstream foutStream(outputPath, std::ios_base::binary);

	write_tree(foutStream, treeHaff);
	tableFilling(treeHaff, foutStream);

	writeCodedText(inPath, foutStream);

	treeDelete(treeHaff);
	foutStream.close();
}

void codeShan(string inPath);
void codeShan(string inPath)
{
	string outputPath = inPath;
	outputPath = outputPath.erase(inPath.size() - 4, 4) + ".shan";

	std::vector<Node*> nodes = nodesFromString(inPath);
	Node* treeShan = buildTreeSh(nodes);
	std::wofstream foutStream(outputPath, std::ios_base::binary);

	write_tree(foutStream, treeShan);

	foutStream << std::endl;

	tableFilling(treeShan, foutStream);

	writeCodedText(inPath, foutStream);

	treeDelete(treeShan);
	foutStream.close();
}

void decode(string inPath);
void decode(string inPath)
{
	string outputPath = inPath;
	if (inPath.size() < 4)
		throw std::logic_error("Error in file name");
	char type = inPath[inPath.size() - 4];


	if (type == 'h')
		outputPath = outputPath.erase(inPath.size() - 5, 5) + "-unz-h.txt";
	else
	{
		if (type == 's')
			outputPath = outputPath.erase(inPath.size() - 5, 5) + "-unz-s.txt";
		else
		{
			throw std::logic_error("Error in file name");
		}
	}

	//Reading tree Root
	basic_ifstream<wchar_t> fin(inPath, std::ios_base::in | std::ios::binary);
	std::locale locale(std::locale(), new std::codecvt_utf8<wchar_t>);
	fin.imbue(locale);
	//Chechking is file exist
	if (!fin.is_open())
	{
		std::cout << "Input file not found" << std::endl;
		system("pause");
		throw std::logic_error("Input file not found");
	}
	Node* treeRoot = read_tree(fin);

	//
	std::wofstream foutStream("test.txt", std::ios_base::binary);

	write_tree(foutStream, treeRoot);
	foutStream.close();

	//Reading coded string
	string stringToDecode = read_string(fin);

	//Decode
	decodeString(treeRoot, stringToDecode, outputPath);

}
void tableFilling(Node *root, std::wofstream& foutHaff)
{
	static std::wstring s = L"";
	if (!root->isLeaf)
	{
		s += L"1";

		tableFilling(root->right, foutHaff);
		s.erase(s.size() - 1, 1);

		s += L"0";

		tableFilling(root->left, foutHaff);
		//
		s.erase(s.size() - 1, 1);
	}
	else
	{
		//	foutHaff << "'" << root->simbolChar << "'" << " = " << s << std::endl;
		tableFreq[root->simbolChar] = s;
	}
}

void treeDelete(Node *root)
{
	if (!root->isLeaf)
	{
		treeDelete(root->right);
		//tree_del(root->br[1]);

		treeDelete(root->left);
	}
	else
		delete root;
}

void writeCodedText(string pathFrom, wofstream& fout)
{
	basic_ifstream<wchar_t> fin(pathFrom, std::ios_base::binary);
	wchar_t ch;

	while (!fin.eof())
	{
		fin.get(ch);
		if (!fin.eof())
			fout << tableFreq[ch];
	}
}

int main(int argc, char** argv)
{
	//string filename = "testrus.txt";
	//std::locale locale(std::locale(), new std::codecvt_utf8<wchar_t>);
	//std::wifstream stream(filename, std::ios::binary);
	//stream.imbue(locale);
	//wchar_t ch;
	//stream.get(ch);

	//std::locale::global(std::locale("rus"));

	string filePath = string(argv[2]) + ".txt";
	//	if (argv[1] == "-compress_h")
	codeHaff(filePath);
	//	if (argv[1] == "-compress_s")
	codeShan(filePath);

	//	if (argv[1] == "-decode_h")
	//	{
	string filePathDecode = string(argv[2]) + ".haff";
	decode(filePathDecode);
	//	}
	//	if (argv[1] == "-decode_s")
	//	{
	filePathDecode = string(argv[2]) + ".shan";
	decode(filePathDecode);
	//	}

	return 0;
}