#include <iostream>
#include <string>
#include <queue>
#include <map>
#include <fstream>

using namespace std;

struct Node {
    char symbol;
    int weight;
    Node *left, *right;

    Node(char symbol, int weight, Node *left, Node *right) {
        this->symbol = symbol;
        this->weight = weight;
        this->left = left;
        this->right = right;
    }
};

struct compareByWeight {
    bool operator()(Node &l, Node &r) {
        // highest priority => lowest weight
        return l.weight > r.weight;
    }
};

map<char, int> countWeights(string &rawText) {
    // Count weights
    map<char, int> weights;
    for (auto &symbol: rawText)
        weights[symbol]++;

    return weights;
}

void writeCodes(Node *root, const string &str, map<char, string> &huffmanCode) {
    if (root == nullptr)
        return;

    // found a leaf node
    if (!root->left && !root->right) {
        huffmanCode[root->symbol] = str;
    }

    writeCodes(root->left, str + "0", huffmanCode);
    writeCodes(root->right, str + "1", huffmanCode);
}

void decode(Node *root, int &index, string encodedText, string &decodedText) {
    if (root == nullptr)
        return;

    // found a leaf node
    if (!root->left && !root->right) {
        decodedText += root->symbol;
        return;
    }

    index++;

    if (encodedText[index] == '0')
        decode(root->left, index, encodedText, decodedText);
    else
        decode(root->right, index, encodedText, decodedText);
}

Node *buildHuffman(map<char, int> &weights) {
    // Tree
    priority_queue<Node, vector<Node>, compareByWeight> nodes;

    // Create a leaf node for each character and add it to nodes.
    for (auto weight: weights)
        nodes.push(*new Node(weight.first, weight.second, nullptr, nullptr));

    while (nodes.size() > 1) {
        // Remove the two nodes with lowest weights from the tree
        Node *left = new Node(nodes.top());
        nodes.pop();
        Node *right = new Node(nodes.top());
        nodes.pop();

        // Merge two deleted nodes in one
        int sum = left->weight + right->weight;
        nodes.push(*new Node('\0', sum, left, right));
    }

    Node *root = new Node(nodes.top());

    return root;
}

map<char, int> decodeWeights(string &codes) {

    int l = 0, weight, counter = 0;
    char symbol;
    map<char, int> weights;
    for (int i = 0; i < codes.length(); i++)
        if (codes[i] == ' ') {
            counter++;
            if (counter % 2 != 0)
                symbol = char(stoi(codes.substr(l, l - i)));
            else {
                weight = stoi(codes.substr(l, l - i));
                weights[symbol] = weight;
            }
            l = i + 1;
        }

    return weights;
}

void decodeText() {
    ifstream encodedFile("../result.txt");
    string codes, encodedText, decodedText;
    getline(encodedFile, codes);
    getline(encodedFile, encodedText);

    map<char, int> weights = decodeWeights(codes);
    Node *root = buildHuffman(weights);

    ofstream decodedFile("../test.txt");
    int index = -1;
    cout << "\nDecoded text:\n";
    while (index + 2 < encodedText.size())
        decode(root, index, encodedText, decodedText);

    cout << decodedText;
    decodedFile << decodedText;
}

void encodeText() {
    // Input
    char test;
    ifstream decodedFile("../test.txt");
    string line, decodedText;

    while (decodedFile.get(test))
        decodedText += test;

    // Calculate
    map<char, int> weights = countWeights(decodedText);
    Node *root = buildHuffman(weights);
    map<char, string> huffmanCode;
    writeCodes(root, "", huffmanCode);

    // Output
    ofstream encodedFile("../result.txt");

    cout << "Symbols weights and codes:\n";
    for (auto &code: huffmanCode) {
        cout << code.first << " " << code.second << " " << weights[code.first] << '\n';
        encodedFile << (int) code.first << ' ' << weights[code.first] << ' ';
    }

    cout << "\nOriginal text:\n" << decodedText << '\n';

    string encodedText;
    for (auto &symbol: decodedText)
        encodedText += huffmanCode[symbol];
    cout << "\nEncoded text:\n" << encodedText << endl;
    encodedFile << endl << encodedText;
}

int main() {

    encodeText();
//    decodeText();

    return 0;
}