compressor.cpp

/**
 * @file compressor.cpp
 * @author Marián Tarageľ (xtarag01)
 * @brief Compression algorithm with static and adaptive scan
 * @date 8.5.2024
 */

#include <cstdio>
#include <cstdint>
#include <vector>
#include <map>
#include "model.hpp"
#include "huffman.hpp"
#include "utils.hpp"
#include "compressor.hpp"
#include "huff_codec.hpp"

void static_scan(std::vector<uint8_t> buffer, int model, const char *ofile)
{
    if (model)
        buffer = rle(diff_model(buffer));

    std::map<uint8_t, unsigned long int> frequencies = count_frequency(buffer);
    std::map<uint8_t, uint8_t> bit_lengths = huff_code_len(frequencies);
    std::map<uint8_t, std::vector<bool>> codes = canonical_huff_codes(bit_lengths);

    FILE *fp = fopen(ofile, "wb");
    if (fp == NULL)
        handle_error("fopen");
    
    // write codebook
    for (std::pair<uint8_t, uint8_t> symbol : bit_lengths) {
        std::vector<bool> length = int2ba(symbol.second, 8);
        write_bits(length, fp);
    }

    // write data
    for (uint8_t symbol : buffer) {
        std::vector<bool> code = codes[symbol];
        write_bits(code, fp);
    }

    // find ending symbol
    std::vector<bool> end_symbol;
    for (std::pair<uint8_t, std::vector<bool>> code : codes) {
        if (code.second.size() >= 8)
            end_symbol = code.second;
    }

    flush_buffer(fp, end_symbol);
    fclose(fp);
}

void adaptive_scan(std::vector<uint8_t> buffer, unsigned long int width, unsigned long int height, int model, const char *ofile)
{
    FILE *fp = fopen(ofile, "wb");

    std::vector<bool> width_bits = int2ba(width, 16);
    write_bits(width_bits, fp);

    for (unsigned long int y = 0; y < height; y+=ADAPTIVE_SCAN_HEIGHT) {
        for (unsigned long int x = 0; x < width; x+=ADAPTIVE_SCAN_WIDTH) {
            
            unsigned long int scan_width = ADAPTIVE_SCAN_WIDTH;
            unsigned long int scan_height = ADAPTIVE_SCAN_HEIGHT;

            // make edge block scan width smaller
            if (x + ADAPTIVE_SCAN_WIDTH > width)
                scan_width = width - x;
            if (y + ADAPTIVE_SCAN_HEIGHT > height)
                scan_height = height - y;

            // scan the block horizontally and vertically
            std::vector<uint8_t> sector_h = scanh(buffer, x, y, width, scan_width, scan_height);
            std::vector<uint8_t> sector_v = scanv(buffer, x, y, width, scan_width, scan_height);

            // apply model
            if (model) {
                sector_h = rle(diff_model(sector_h));
                sector_v = rle(diff_model(sector_v));
            }

            // compress the block
            std::vector<bool> compressed_h = compress_block(sector_h, 0);
            std::vector<bool> compressed_v = compress_block(sector_v, 1);

            if (((sector_h.size() * 8) < compressed_h.size()) && ((sector_h.size() * 8) < compressed_v.size())) { // uncompressed

                write_bits({0}, fp);
                std::vector<bool> size = int2ba(sector_h.size(), 10);
                write_bits(size, fp);
                for (uint8_t pixel : sector_h) {
                    std::vector<bool> bits = int2ba(pixel, 8);
                    write_bits(bits, fp);
                }

            } else if (compressed_h.size() <= compressed_v.size()) { // compressed horizontal
                write_bits({1}, fp);
                write_bits(compressed_h, fp);
            } else { // compressed vertical
                write_bits({1}, fp);
                write_bits(compressed_v, fp);
            }
        }
    }

    flush_buffer(fp, {0, 0, 0, 0, 0, 0, 0, 0});
    fclose(fp);
}

std::vector<bool> compress_block(std::vector<uint8_t> buffer, bool direction)
{
    std::map<uint8_t, unsigned long int> frequencies = count_frequency(buffer);
    std::map<uint8_t, uint8_t> bit_lengths = huff_code_len(frequencies);
    std::map<uint8_t, std::vector<bool>> codes = canonical_huff_codes(bit_lengths);

    for (std::pair<uint8_t, unsigned long int> freq : frequencies) {
        if (!freq.second)
            bit_lengths.erase(freq.first);
    }

    std::vector<uint8_t> lengths;
    std::vector<uint8_t> symbols;
    for (std::pair<uint8_t, uint8_t> bitl : bit_lengths) {
        symbols.push_back(bitl.first);
        lengths.push_back(bitl.second);
    }

    auto p = sort_permutation(lengths);
    symbols = apply_permutation(symbols, p);
    lengths = apply_permutation(lengths, p);

    std::vector<unsigned long int> count_lengths = count_bit_lengths(lengths);

    std::vector<bool> result = int2ba(count_lengths.size(), 8);

    // write codebook
    for (unsigned long int len : count_lengths) {
        std::vector<bool> bits = int2ba(len, 9);
        result.insert(result.end(), bits.begin(), bits.end());
    }

    for (unsigned long int symb : symbols) {
        std::vector<bool> bits = int2ba(symb, 8);
        result.insert(result.end(), bits.begin(), bits.end());
    }

    std::vector<bool> size = int2ba(buffer.size(), 10);
    result.insert(result.end(), size.begin(), size.end());

    // write scan direction
    result.push_back(direction);

    // write data
    for (uint8_t symbol : buffer) {
        std::vector<bool> code = codes[symbol];
        result.insert(result.end(), code.begin(), code.end());
    }

    return result;
}

std::vector<uint8_t> scanh(std::vector<uint8_t> buffer, unsigned long int x, unsigned long int y, unsigned long int width,
                           unsigned long int scan_width, unsigned long int scan_height)
{
    std::vector<uint8_t> sector;
    sector.resize(scan_height * scan_width);

    for (unsigned long int row = 0; row < scan_height; row++) {
        for (unsigned long int column = 0; column < scan_width; column++) {
            unsigned long int index = width * (y + row) + x + column;
            sector[scan_width * row + column] = buffer[index];
        }
    }

    return sector;
}

std::vector<uint8_t> scanv(std::vector<uint8_t> buffer, unsigned long int x, unsigned long int y, unsigned long int width,
                           unsigned long int scan_width, unsigned long int scan_height)
{
    std::vector<uint8_t> sector;
    sector.resize(scan_height * scan_width);

    for (unsigned long int column = 0; column < scan_width; column++) {
        for (unsigned long int row = 0; row < scan_height; row++) {
            unsigned long int index = width * (y + row) + x + column;
            sector[row + scan_height * column] = buffer[index];
        }
    }

    return sector;
}