#include <string>
#include <queue>
#include <unordered_map>
#include <map>
#include <iostream>
#include <set>
#include <bitset>
#include <sstream>

#include "bitstream.h"
#include "huffman_table.h"
#include "MashZip.h"

using namespace std;

// // A Tree node
// struct Node
// {
//     char ch;
//     int freq;
//     Node *left, *right;
// };

// // Function to allocate a new tree node
// Node* getNode(char ch, int freq, Node* left, Node* right)
// {
//     Node* node = new Node();

//     node->ch = ch;
//     node->freq = freq;
//     node->left = left;
//     node->right = right;

//     return node;
// }

// // Comparison object to be used to order the heap
// struct comp
// {
//     bool operator()(Node* l, Node* r)
//     {
//         // highest priority item has lowest frequency
//         return l->freq > r->freq;
//     }
// };

// // traverse the Huffman Tree and store Huffman Codes
// // in a map.
// void encode(Node* root, int len,
//             map<int, std::set<char> > &huffmanLengths)
// {
//     if (root == nullptr)
//         return;

//     // found a leaf node
//     if (!root->left && !root->right) {
//         // huffmanCode[root->ch] = str;
//         huffmanLengths[len].insert(root->ch);
//     }

//     encode(root->left, len + 1, huffmanLengths);
//     encode(root->right, len + 1, huffmanLengths);
// }

// void initializeTable(const map<int, std::set<char> > &huffmanLengths,
//                     unordered_map<char, std::pair<int, short> > &table)
// {
//     int nextbl = 0;
//     short code = 0;

//     for (auto lenCodePairIt = huffmanLengths.begin(); lenCodePairIt != huffmanLengths.end(); lenCodePairIt++)
//     {
//         auto lenCodePair = *lenCodePairIt;

//         for (auto it = lenCodePair.second.begin(); it != lenCodePair.second.end(); it++)
//         {
//             table[*it].first = lenCodePair.first; // save current bit length for code
//             table[*it].second = code;

//             // code := (code + 1) << ((bit length of the next symbol) − (current bit length))
//             // code++;

//             // if last symbol of length and has symbols with greater length
//             if (next(it) == lenCodePair.second.end() && next(lenCodePairIt) != huffmanLengths.end()) {
//                 nextbl = (*next(lenCodePairIt)).first;
//             } else {
//                 nextbl = lenCodePair.first;
//             }

//             code = (code + 1) << (nextbl - lenCodePair.first);
//         }

//         // code <<= 1;
//     }
// }

// // traverse the Huffman Tree and decode the encoded string
// void decode(Node* root, int &index, string str)
// {
//     if (root == nullptr) {
//         return;
//     }

//     // found a leaf node
//     if (!root->left && !root->right)
//     {
//         cout << root->ch;
//         return;
//     }

//     index++;

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

// // Builds Huffman Tree and decode given input text
// void buildHuffmanTree(basic_istream<char> &is)
// {
//     // count frequency of appearance of each character
//     // and store it in a map
//     unordered_map<char, int> freq;
//     char ch;
//     while (is.get(ch)) {
//         freq[ch]++;
//     }

//     for (auto freqPair : freq) {
//         cout << "Freq " << freqPair.first << " " << freqPair.second << endl;
//     }

//     // Create a priority queue to store live nodes of
//     // Huffman tree;
//     priority_queue<Node*, vector<Node*>, comp> pq;

//     // Create a leaf node for each character and add it
//     // to the priority queue.
//     for (auto pair: freq) {
//         pq.push(getNode(pair.first, pair.second, nullptr, nullptr));
//     }

//     // do till there is more than one node in the queue
//     while (pq.size() != 1)
//     {
//         // Remove the two nodes of highest priority
//         // (lowest frequency) from the queue
//         Node *left = pq.top(); pq.pop();
//         Node *right = pq.top(); pq.pop();

//         // Create a new internal node with these two nodes
//         // as children and with frequency equal to the sum
//         // of the two nodes' frequencies. Add the new node
//         // to the priority queue.
//         int sum = left->freq + right->freq;
//         pq.push(getNode('\0', sum, left, right));
//     }

//     // root stores pointer to root of Huffman Tree
//     Node* root = pq.top();

//     map<int, std::set<char> > huffmanLengths;
//     encode(root, 0, huffmanLengths);

    
//     unordered_map<char, std::pair<int, short> > huffmanCode;
//     initializeTable(huffmanLengths, huffmanCode);

//     cout << "Huffman Codes are :\n" << '\n';
//     for (auto pair: huffmanCode) {
//         cout << pair.first << " " << pair.second.first << " " << bitset<16>(pair.second.second) << '\n';
//     }

//     // cout << "\nOriginal string was :\n" << text << '\n';

//     // // print encoded string
//     // string str = "";
//     // for (char ch: text) {
//     //     str += huffmanCode[ch];
//     // }

//     // cout << "\nEncoded string is :\n" << str << '\n';

//     // // traverse the Huffman Tree again and this time
//     // // decode the encoded string
//     // int index = -1;
//     // cout << "\nDecoded string is: \n";
//     // while (index < (int)str.size() - 2) {
//     //     decode(root, index, str);
//     // }
// }

//// Huffman coding algorithm
int main(int argc, char **argv)
{
// //    buildHuffmanTree(cin);
//     // 10111010 11110101
//     stringstream ss("\xba\xf5");
//     ibitstream ibs(ss);

//     // 10111010 11110101
//     //       ^^
//     int a = ibs.getbits(2);
//     cout << bitset<2>(a) << endl; // 01

//     // 10111010 11110101
//     //   ^^^^
//     int b = ibs.getbits(4);
//     cout << bitset<4>(b) << endl; // 0111


//     // 10111010 11110101
//     // ^^        ^^^^^^^
//     b = ibs.getbits(9);
//     cout << bitset<9>(b) << endl; // 011010111

//     // 10111010 11110101
//     //          ^ + overflow
//     try {
//         b = ibs.getbits(10);
//         cout << b << endl;
//     } catch (std::runtime_error &e) {
//         cout << "Got runtime error: " << e.what() << endl;
//     }

//     ostringstream so("bits: ");
//     obitstream obs(so);

//     obs.writebits(0xf, 3);
//     cout << "After 3 bits: " << so.str() << endl;
//     // cache here: 00000111


//     // de = 11011110
//     obs.writebits(0xde, 7);
//     // here: Flush: 11101111
//     cout << "After 7 bits: " << so.str() << endl;
//     obs.flush();
//     // here: Flush: 00000001
//     cout << "After flush: " << so.str() << endl;


//     obs.writebits(0xbeef, 16);
//     cout << "After 0xbeef: " << so.str() << endl;
//     obs.flush();
//     // here: Flush: 0xEFBE - reversed!
//     cout << "After flush: " << so.str() << endl;

    
    // string s = ;

    stringstream ss1("Some long text!!!!\x01\x02\x03\x04\n123123456789 privet, masha!");
    stringstream ss2("Some long text!!!!\x01\x02\x03\x04\n123123456789 privet, masha!");
    stringstream so1;
    stringstream so2;

    MashZip mz;

    mz.mashzip_file(ss1, ss2, so1);

    std::cout << "After mashzip: " << so1.str();

    mz.unmashzip_stream(so1, so2);

    std::cout << "After unmashzip: " << so2.str();

    // HuffmanTable ht(ss1);

    // char *header = ht.to_header();
    // // for (size_t i = 0; i < 128; i++)
    // // {
    // //     cout << "Code for " << 2 * i << " and " << 2 * i + 1 << ": " << bitset<8>(header[i]) << endl; 
    // // }
    
    // stringstream test;

    // // test << "MASH"; // magic

    // // for (size_t i = 0; i < 128; i++) {
    // //     test << header[i];
    // // }

    // obitstream some_stream(test);

    // for (char c : s) {
    //     ht.write_symbol(some_stream, c);
    // }
    // some_stream.flush();

    // std::cout << test.str() << endl;

    // ibitstream some_instream(test);

    // while(true) {
    //     char sym = (char) ht.decode_one_symbol(some_instream);
    //     if (sym < 0)
    //     {
    //         std::cout << "Read all" << std::endl;
    //         break;
    //     }
    //     std::cerr << "Sym: " << sym << std::endl;
    // }

    return 0;
}