/*
* file: encoding.cpp
* ---------------------
* author: Nahva Techlu and Jeremy Bao
* date: Nov. 27
* This is the file that implements the Huffman algorithm.
* It includes functions that build encoding trees, build
* encoding maps, encode data, decode data and clean a tree.
* Finally, the compress and decompress function wraps them
* up to conduct the entire process of compressing and
* decompressing.
*/

#include "encoding.h"
#include <istream>
#include "pqueue.h"
#include "filelib.h"
#include "bitstream.h"
using namespace std;

//Self-defined functions
void recursiveMap(Map<int, string>& encodingMap, HuffmanNode* encodingTree, string code);
void recursiveDecode(ibitstream& input, HuffmanNode* current, HuffmanNode* encodingTree, ostream &output);

//This is a function that builds a frequency table
//for a specific file. It will output a map.
Map<int, int> buildFrequencyTable(istream& input) {
    Map<int, int> freqTable;
    while (true){
        int temp = input.get();
        if (temp == -1){
            break;
        }
        freqTable[temp]++;
    }
    freqTable[PSEUDO_EOF]++;
    return freqTable;
}

//This is a function that build an encoding tree
//based on a frequency table.
HuffmanNode* buildEncodingTree(const Map<int, int>& freqTable) {
    PriorityQueue<HuffmanNode* > hQueue;
    //The enqueuing process
    for (int eachKey : freqTable){
        HuffmanNode* node = new HuffmanNode(eachKey, freqTable[eachKey], NULL, NULL);
        hQueue.enqueue(node, freqTable[eachKey]);
    }

    //Tree building process
    while (hQueue.size() > 1){
        HuffmanNode* temp1 = hQueue.dequeue();
        HuffmanNode* temp2 = hQueue.dequeue();
        HuffmanNode* sum = new HuffmanNode(NOT_A_CHAR, temp1->count + temp2->count, temp1, temp2);
        hQueue.enqueue(sum, temp1->count + temp2->count);
    }
    return hQueue.dequeue();
}

//This is a function that builds an encoding map based
//on an encoding tree. It will output this map, which will
//be used to encode the data.
Map<int, string> buildEncodingMap(HuffmanNode* encodingTree) {
    Map<int, string> encodingMap;
    if (encodingTree == NULL){
        return encodingMap;
    }
    recursiveMap(encodingMap, encodingTree, "");
    return encodingMap;
}

//This is the actual implementation of building the encoding map.
//It uses recursive algorithm to achieve it.
void recursiveMap(Map<int, string>& encodingMap, HuffmanNode* encodingTree, string code){
    if ((encodingTree->zero == NULL) &&(encodingTree->one == NULL)){
        encodingMap[encodingTree->character] = code;
    }
    if (encodingTree->zero != NULL){
        recursiveMap(encodingMap, encodingTree->zero, code + "0"); //Adding 0 or 1 based on the encoding tree
    }
    if (encodingTree->one != NULL){
        recursiveMap(encodingMap, encodingTree->one, code + "1");
    }
}

//This is a function that encodes a file to compressed form
//based on the encoding map. It will output the compressed
//form of the file.
void encodeData(istream& input, const Map<int, string>& encodingMap, obitstream& output) {
    while (true){
        int temp = input.get();
        if (temp == -1){
            break;
        }
        string toWrite = encodingMap[temp];
        for (int i = 0; i < toWrite.length(); i++){
            output.writeBit(toWrite[i] - 48);
        }
    }

    string toWrite = encodingMap[256];
    for (int i = 0; i < toWrite.length(); i++){
        output.writeBit(toWrite[i] - 48);
    }
}

//This is a function that decodes the compressed form based on the
//encoding tree. It uses recursive algorithm to achieve that.
void decodeData(ibitstream& input, HuffmanNode* encodingTree, ostream& output) {
    string toWrite = "";
    recursiveDecode(input, encodingTree, encodingTree, output);
}

//This is the actual implementation of decoding, and uses recursive
//algorithm to achieve that.
void recursiveDecode(ibitstream& input, HuffmanNode* current, HuffmanNode* encodingTree, ostream& output){
    //Base cases
    if (current->character == PSEUDO_EOF){
        return;
    } else if (current->character != NOT_A_CHAR){
        output.put(current->character + 0);
        current = encodingTree;
    }

    //Recursive cases
    int temp = input.readBit();
    if (temp == 0){
        recursiveDecode(input, current->zero, encodingTree, output);
    } else if (temp == 1){
        recursiveDecode(input, current->one, encodingTree, output);
    }
}

//This is the function that wraps up the compression process. It uses several
//previous functions to achieve that.
void compress(istream& input, obitstream& output) {
    //Build frequency table
    Map<int, int> frequencyTable = buildFrequencyTable(input);
    output<<frequencyTable;

    //Build encoding tree
    HuffmanNode* encodingTree = buildEncodingTree(frequencyTable);

    //Build encoding map
    Map<int, string> encodingMap = buildEncodingMap(encodingTree);

    //Encode data
    rewindStream(input);
    encodeData(input, encodingMap, output);

    //Clear tree
    freeTree(encodingTree);
}

//This is the function that wraps up the decompression process. It uses several
//previous functions to achieve that.
void decompress(ibitstream& input, ostream& output) {
    // Build frequency table
    Map<int, int> frequencyTable;
    input>>frequencyTable;

    // Build encoding tree
    HuffmanNode* encodingTree = buildEncodingTree(frequencyTable);

    //Decode data
    decodeData(input, encodingTree, output);

    //Clear tree
    freeTree(encodingTree);
}

//This is the function that clears up a tree that has been built
//before. It prevents the program from leaking memory.
void freeTree(HuffmanNode* node) {
    //Base case
    if (node == NULL){
        return;
    }

    //Recursive case
    if (node->zero != NULL){
        freeTree(node->zero);
    }
    if (node->one != NULL){
        freeTree(node->one);
    }
    delete node;
}
