#include <iostream>

#include "TestFramework.h"

#include "HuffmanStream.h"
#include "BufferBitStream.h"

TEST_CASE(TestHuffmanCodeLengthTable, "compression")
{
    HuffmanCodeLengthTable table;

    std::vector<std::pair<unsigned, unsigned char> > mappings {{144, 8}, {112, 9}, {24, 7}, {8 ,8}};
    std::vector<unsigned char> code_length_sequence;
    for(const auto& entry : mappings)
    {
        for(unsigned idx=0;idx<entry.first;idx++)
        {
            code_length_sequence.push_back(entry.second);
        }
    }
    table.setInputLengthSequence(code_length_sequence, false);
    table.buildCompressedLengthSequence();

    auto compressed_sequence = table.getCompressedLengthSequence();
    //for (auto entry : compressed_sequence)
    //{
        // std::cout << "Count " << entry.first << " extra bits " << entry.second << std::endl;
    //}

    auto compressed_lengths = table.getCompressedLengthCounts();
    for(unsigned idx = 0; idx<compressed_lengths.size(); idx++)
    {
        //std::cout << "Slot " << idx << " length " << compressed_lengths[idx] << std::endl;
    }
};


TEST_CASE(TestLiteralsTable, "compression")
{
    std::vector<unsigned char> lengths = {7,4,4,7,5,5,7,7,6,6,7,6,6,6,8,6,6,8,
            6,6,7,6,8,7,7,7,7,7,7,6,6,7,7,6,6,7,7,8,8,7,7,7,6,6,7,7,7,7,6,7,7,7,
            7,7,7,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,11,4,4,5,6,7,7,9,8,8,9,9,8,8,7,
            7,8,7,6,8,9,9,11,10,8,7,7,8,8,7,3,0,9,9,9,9,9,8,6,7,7,7,7,9,5,7,4,7,4,4,4,3,4,4,4,4,4,5,5,6};

    HuffmanCodeLengthTable table;
    table.setInputLengthSequence(lengths, false);

    table.buildCompressedLengthSequence();

    auto compressed_sequence = table.getCompressedLengthSequence();
    //for (auto entry : compressed_sequence)
    //{
        // std::cout << "Code " << entry.first << " extra bits " << entry.second << std::endl;
    //}

    auto compressed_lengths = table.getCompressedLengthCounts();
    for(unsigned idx = 0; idx<compressed_lengths.size(); idx++)
    {
        // std::cout << "Slot " << idx << " length " << compressed_lengths[idx] << std::endl;
    }

    HuffmanCodeLengthTable codingTable;
    std::vector<unsigned char> coding_lengths{4, 0, 6, 7, 3, 2, 4, 2, 7, 4, 6, 3, 0, 6, 0, 0, 0, 0, 0};
    codingTable.setInputLengthSequence(coding_lengths, true);
    codingTable.buildPrefixCodes();

    BufferBitStream out_stream;
    out_stream.writeNBits(1, 1);
    out_stream.writeNBits(2, 2);

    out_stream.writeNBits(29, 5);
    out_stream.writeNBits(29, 5);
    out_stream.writeNBits(10, 4);

    /*
    std::vector<unsigned char> permuted(19, 0);
    static constexpr unsigned DEFLATE_PERMUTATION[19]{16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15};
    unsigned count = 0;
    for (auto length : coding_lengths)
    {
        permuted[DEFLATE_PERMUTATION[count]] = length;
        count++;
    }

    unsigned skip_count = 0;
    for(unsigned idx=0; idx<permuted.size();idx++)
    {
        if (permuted[permuted.size() - 1 - idx] == 0)
        {
            skip_count++;
        }
    }
    std::cout << "Got skip count " << skip_count << std::endl;

    for(unsigned idx=0; idx<permuted.size() - skip_count;idx++)
    {
        out_stream.writeNBits(permuted[idx], 3);
    }

    for(unsigned idx=0; idx<14;idx++)
    {
        out_stream.writeNBits(coding_lengths[idx], 3);
    }

    for (const auto& entry : compressed_sequence)
    {
        auto prefix_code =  *codingTable.getCodeForSymbol(entry.first);
        out_stream.writeNBits(prefix_code.getData(), prefix_code.getLength());

        std::cout << "Stream count " << out_stream.getBuffer().size() << " for entry " << entry.first << std::endl;

        if (entry.first == 16)
        {
            out_stream.writeNBits(entry.second, 2);
        }
        else if (entry.first == 17)
        {
            out_stream.writeNBits(entry.second, 3);
        }
        else if (entry.first == 18)
        {
            out_stream.writeNBits(entry.second, 7);
        }
    }
    out_stream.resetOffsets();

    std::cout << "Output is: " << std::endl;
    auto dump = out_stream.logNextNBytes(out_stream.getBuffer().size());
    std::cout << dump << std::endl;
    */
};


//int main()
//{
    //testHuffmanCodeLengthTable();

    // testLiteralsTable();
    //HuffmanStream stream(nullptr, nullptr);

    //stream.setCodeLengthAlphabetLengths({3, 3, 3, 3, 3, 2, 4, 4});

    //stream.setCodeLengthAlphabetLengths({2, 2, 3, 3, 3, 3});

    //stream.buildCodeLengthMapping();

    //std::cout << "*******" << std::endl;
    //stream.setCodeLengthAlphabetLengths({4, 0, 6, 7, 3, 2, 4, 2, 7, 4, 6, 3, 0, 6});

    //stream.buildCodeLengthMapping();

    //const auto mapping = stream.getCodeLengthMapping();

    //stream.generateFixedCodeMapping();

//    testCodeConversions();


//    return 0;
//}