compression_handler.cpp

#include "compression_handler.h"
#include <fstream>
#include <cstring>
#include "quantization.h"
#include "prediction.h"
#include "entropy.h"
#include <algorithm>

static std::vector<uint8_t> pack_block_modes_3bit_prefixed(const std::vector<uint8_t>& modes)
{
    uint32_t count = (uint32_t)modes.size();

    std::vector<uint8_t> out;
    out.reserve(4 + (modes.size() * 3 + 7) / 8);

    out.push_back((uint8_t)(count & 0xFF));
    out.push_back((uint8_t)((count >> 8) & 0xFF));
    out.push_back((uint8_t)((count >> 16) & 0xFF));
    out.push_back((uint8_t)((count >> 24) & 0xFF));

    uint32_t acc = 0;
    int accBits = 0;

    for (uint8_t m : modes) {
        m &= 7;
        acc |= (uint32_t)m << accBits;
        accBits += 3;

        while (accBits >= 8) {
            out.push_back((uint8_t)(acc & 0xFF));
            acc >>= 8;
            accBits -= 8;
        }
    }

    if (accBits > 0) out.push_back((uint8_t)(acc & 0xFF));
    return out;
}

static std::vector<uint8_t> unpack_block_modes_3bit_prefixed(const std::vector<uint8_t>& packed)
{
    if (packed.size() < 4) return {};

    uint32_t count =
        (uint32_t)packed[0] |
        ((uint32_t)packed[1] << 8) |
        ((uint32_t)packed[2] << 16) |
        ((uint32_t)packed[3] << 24);

    std::vector<uint8_t> out;
    out.reserve(count);

    uint32_t acc = 0;
    int accBits = 0;
    size_t idx = 4;

    while (out.size() < count) {
        while (accBits < 3) {
            if (idx >= packed.size()) return out;
            acc |= (uint32_t)packed[idx++] << accBits;
            accBits += 8;
        }

        out.push_back((uint8_t)(acc & 7u));
        acc >>= 3;
        accBits -= 3;
    }

    return out;
}

// Load all YUV frames from .yuv file
static std::vector<YUVFrame> load_yuv_frames(const CompressionSettings& settings)
{
	// Construct frame list
    std::vector<YUVFrame> frames;
    int w = settings.width;
    int h = settings.height;
    const char* path = settings.inputYuvPath.c_str();
    std::ifstream in(path, std::ios::binary | std::ios::ate);
    if (!in) return frames;

    std::streamoff size = in.tellg();
    in.seekg(0, std::ios::beg);

    size_t yBytes = w * h;
    size_t cBytes = yBytes / 4;
    size_t frameBytes = yBytes + 2 * cBytes;
    int frameCount = static_cast<int>(size / frameBytes);

	// Read all frames
    for (int f = 0; f < frameCount; ++f) {
        std::vector<uint8_t> buf(frameBytes);
        in.read(reinterpret_cast<char*>(buf.data()), frameBytes);
        if (in.gcount() != (std::streamsize)frameBytes) break;

        YUVFrame frame;
        frame.Y.resize(yBytes);
        for (size_t i = 0; i < yBytes; ++i)
            frame.Y[i] = static_cast<int16_t>(buf[i]);
        frame.U.resize(cBytes);
        for (size_t i = 0; i < cBytes; ++i)
            frame.U[i] = static_cast<int16_t>(buf[yBytes + i]);
        frame.V.resize(cBytes);
        for (size_t i = 0; i < cBytes; ++i)
            frame.V[i] = static_cast<int16_t>(buf[yBytes + cBytes + i]);

        frames.push_back(std::move(frame));
    }
    return frames;
}

static YUVFrame convert_to_uint8(const YUVFrame& src) {
    YUVFrame dst;
    dst.Y.resize(src.Y.size());
    dst.U.resize(src.U.size());
    dst.V.resize(src.V.size());
    for (size_t i = 0; i < src.Y.size(); ++i)
        dst.Y[i] = static_cast<uint8_t>(std::clamp<int>(src.Y[i], 0, 255));
    for (size_t i = 0; i < src.U.size(); ++i)
        dst.U[i] = static_cast<uint8_t>(std::clamp<int>(src.U[i], 0, 255));
    for (size_t i = 0; i < src.V.size(); ++i)
        dst.V[i] = static_cast<uint8_t>(std::clamp<int>(src.V[i], 0, 255));
    return dst;
}


// temporales Downsampling (skip every Nth frame)
// frameStep = N: jedes N-te Frame wird NICHT gespeichert (1-basiert: N,2N,3N,...)
static std::vector<YUVFrame> temporal_downsample(const std::vector<YUVFrame>& in, int frameStep)
{
    if (frameStep <= 1) return in;

    std::vector<YUVFrame> out;
    out.reserve(in.size());

    for (size_t i = 0; i < in.size(); ++i) {
        // 1-basiert: i=0 -> Frame 1
        bool isSkipped = (((i + 1) % (size_t)frameStep) == 0);
        if (!isSkipped) out.push_back(in[i]);
    }
    return out;
}

// Upsampling durch Interpolation (Decoder-Seite)
// frameStep = N: jedes N-te Frame fehlt im Stream und wird aus Nachbarn interpoliert
static YUVFrame interpolate_half(const YUVFrame& a, const YUVFrame& b)
{
    YUVFrame out;
    out.Y.resize(a.Y.size());
    out.U.resize(a.U.size());
    out.V.resize(a.V.size());

    for (size_t i = 0; i < a.Y.size(); ++i) out.Y[i] = (int16_t)((a.Y[i] + b.Y[i]) / 2);
    for (size_t i = 0; i < a.U.size(); ++i) out.U[i] = (int16_t)((a.U[i] + b.U[i]) / 2);
    for (size_t i = 0; i < a.V.size(); ++i) out.V[i] = (int16_t)((a.V[i] + b.V[i]) / 2);
    return out;
}

static std::vector<YUVFrame> temporal_upsample_interpolate(
    const std::vector<YUVFrame>& stored,
    int frameStep,
    int originalFrameCount)
{
    if (frameStep <= 1) return stored;
    if (originalFrameCount <= 0) return stored;

    std::vector<YUVFrame> out;
    out.reserve((size_t)originalFrameCount);

    size_t storedIdx = 0;

    for (int t = 0; t < originalFrameCount; ++t) {
        bool isSkipped = (((t + 1) % frameStep) == 0);

        if (!isSkipped) {
            // dieses Frame war gespeichert
            if (storedIdx < stored.size()) {
                out.push_back(stored[storedIdx++]);
            } else {
                // fail-safe: falls Stream zu kurz
                if (!out.empty()) out.push_back(out.back());
                else if (!stored.empty()) out.push_back(stored.back());
                else out.push_back(YUVFrame{});
            }
        } else {
            // dieses Frame fehlte -> interpolieren zwischen prev (out.back) und next (stored[storedIdx])
            if (out.empty()) {
                // sollte praktisch nicht passieren, aber fail-safe
                if (storedIdx < stored.size()) out.push_back(stored[storedIdx]);
                else out.push_back(YUVFrame{});
            } else if (storedIdx < stored.size()) {
                out.push_back(interpolate_half(out.back(), stored[storedIdx]));
            } else {
                // kein "next" vorhanden (Ende): wiederholen
                out.push_back(out.back());
            }
        }
    }

    return out;
}




// Helper to run prediction and block mode extraction
static void run_prediction(
    const CompressionSettings& settings,
    const std::vector<YUVFrame>& inputFrames,
    std::vector<YUVFrame>& predictedFrames,
    std::vector<uint8_t>& blockModesFlat)
{
    if (settings.predictionType == PRED_NONE) {
        predictedFrames.clear();
        predictedFrames.reserve(inputFrames.size());
        for (const auto& frame : inputFrames)
            predictedFrames.push_back(convert_to_uint8(frame));
        blockModesFlat.clear();
    }

    else if (settings.predictionType == PRED_BLOCK) {
        std::vector<BlockModes> allFramesBlockModes;
        predictedFrames = apply_prediction_auto(settings, inputFrames, allFramesBlockModes);
        blockModesFlat = flattenBlockModes(allFramesBlockModes);
    }
    else {
        // Unknown prediction type, treat as no prediction
        predictedFrames = inputFrames;
        blockModesFlat.clear();
    }
}


// Entropy encoding dispatcher
EncodedStream run_entropy(const CompressionSettings& settings,
                          const std::vector<YUVFrame>& frames,
                          const std::vector<uint8_t>& blockModes)
{
    switch (settings.entropyType) {
    case ENTROPY_RICE:
        return entropy_encode_rice(settings, frames, blockModes);

    case ENTROPY_HUFFMAN:
        return entropy_encode_huffman(settings, frames, blockModes);

    case ENTROPY_NONE:
    default:
        return bitpack_and_header(settings, frames, blockModes);
    }
}

// Bitpack quantized frames and construct header
EncodedStream bitpack_and_header(const CompressionSettings& settings, const std::vector<YUVFrame>& quantizedFrames, const std::vector<uint8_t>& blockModes)
{
    EncodedStream stream;
    int w = settings.width, h = settings.height, fps = settings.fps;
    int bitsY = settings.bitsY, bitsC = settings.bitsC;
    int frameCount = static_cast<int>(quantizedFrames.size());

    stream.header.magic[0] = 'V';
    stream.header.magic[1] = 'I';
    stream.header.magic[2] = 'D';
    stream.header.width = w;
    stream.header.height = h;
    stream.header.fps = fps;
    stream.header.bitsY = bitsY;
    stream.header.bitsC = bitsC;
    stream.header.frameCount = frameCount;

    // Frame-Step speichern
    stream.header.frameStep = (settings.frameStep <= 0) ? 1 : settings.frameStep;

	stream.header.predictionType = settings.predictionType;
	stream.header.entropyType = settings.entropyType;
	stream.header.blockModeCount = static_cast<int>(blockModes.size());
    stream.header.riceKY = 0;
    stream.header.riceKC = 0;

    stream.blockModes = blockModes;

    size_t yBytes = w * h;
    size_t cBytes = yBytes / 4;

    // Bitpacking
    std::vector<int16_t> data;
    if (settings.predictionType == PRED_NONE) {
        // 8-bit packing: two uint8_t per int16_t
        std::vector<uint8_t> bytes;
        for (const auto& frame : quantizedFrames) {
            for (auto v : frame.Y) bytes.push_back(static_cast<uint8_t>(std::clamp<int>(v, 0, 255)));
            for (auto v : frame.U) bytes.push_back(static_cast<uint8_t>(std::clamp<int>(v, 0, 255)));
            for (auto v : frame.V) bytes.push_back(static_cast<uint8_t>(std::clamp<int>(v, 0, 255)));
        }
        // Pack two bytes into one int16_t
        for (size_t i = 0; i < bytes.size(); i += 2) {
            int16_t packed = bytes[i];
            if (i + 1 < bytes.size())
                packed |= (static_cast<int16_t>(bytes[i + 1]) << 8);
            data.push_back(packed);
        }
    }
    else {
        // 16-bit packing: one value per int16_t
        for (const auto& frame : quantizedFrames) {
            data.insert(data.end(), frame.Y.begin(), frame.Y.end());
            data.insert(data.end(), frame.U.begin(), frame.U.end());
            data.insert(data.end(), frame.V.begin(), frame.V.end());
        }
    }

    stream.data = std::move(data);
    return stream;
}

// Write the encoded stream to file
bool write_vid_file(const CompressionSettings& settings, const EncodedStream& stream)
{
    std::ofstream out(settings.outputVidPath, std::ios::binary);
    if (!out) return false;

    // Write header
    out.write(reinterpret_cast<const char*>(&stream.header), sizeof(VidHeader));
    if (!out) return false;

    // Write blockModes
	if (!stream.blockModes.empty()) {
		out.write(reinterpret_cast<const char*>(stream.blockModes.data()), stream.blockModes.size());
		if (!out) return false;
	}

    // Write frame data
    out.write(reinterpret_cast<const char*>(stream.data.data()), stream.data.size() * sizeof(int16_t));

    return (bool)out;
}

// Orchestrate the pipeline
bool export_vid(const CompressionSettings& settings)
{
    // 1. Load YUV frames
    auto framesAll = load_yuv_frames(settings);
    if (framesAll.empty()) return false;
    int originalFrameCount = (int)framesAll.size();

    // Temporales Downsampling (Frame-Step)
    int frameStep = (settings.frameStep <= 0) ? 1 : settings.frameStep;
    auto frames = temporal_downsample(framesAll, frameStep);
    if (frames.empty()) return false;

    // 2. Prediction and block mode handling
    std::vector<YUVFrame> predictedFrames;
    std::vector<uint8_t> blockModesFlat;
    run_prediction(settings, frames, predictedFrames, blockModesFlat);

    std::vector<uint8_t> blockModesPacked;
    if (!blockModesFlat.empty()) blockModesPacked = pack_block_modes_3bit_prefixed(blockModesFlat);


    // Quantisierung (nur für PRED_NONE)
    std::vector<YUVFrame> qFrames = predictedFrames;

    // 3. Entropy encode or bitpack
    EncodedStream encoded;
    if (settings.entropyType != ENTROPY_NONE)
        encoded = run_entropy(settings, qFrames, blockModesPacked);
    else
        encoded = bitpack_and_header(settings, qFrames, blockModesPacked);

    //  Frame-Step ins Header schreiben (für ENTROPY_* Fälle)
    encoded.header.frameStep = frameStep;
    encoded.header.originalFrameCount = originalFrameCount;

    // 4. Write to .vid file
    return write_vid_file(settings, encoded);
}


// Reads the header and encoded data from a .vid file
VidHeader read_vid_header(const std::string& vidPath, std::vector<int16_t>& encodedData, std::vector<uint8_t>& blockModes)
{
    std::ifstream in(vidPath, std::ios::binary);
    VidHeader header{};
    if (!in) return header;
    in.read(reinterpret_cast<char*>(&header), sizeof(VidHeader));
    if (!in) return header;

    // Read blockModes
    blockModes.resize(header.blockModeCount);
    in.read(reinterpret_cast<char*>(blockModes.data()), header.blockModeCount);

    // Determine the size of the remaining file (data section)
    std::streamoff dataStart = sizeof(VidHeader) + header.blockModeCount;
    in.seekg(0, std::ios::end);
    std::streamoff fileEnd = in.tellg();
    std::streamoff dataBytes = fileEnd - dataStart;
    size_t numInt16 = static_cast<size_t>(dataBytes) / sizeof(int16_t);

    // Go back to the start of the data section
    in.seekg(dataStart, std::ios::beg);
    encodedData.resize(numInt16);
    in.read(reinterpret_cast<char*>(encodedData.data()), numInt16 * sizeof(int16_t));

    return header;
}

// Unpacks the bitstream into quantized frames
std::vector<YUVFrame> unpack_bitstream(const VidHeader& header, const std::vector<int16_t>& encodedData)
{
    size_t yBytes = header.width * header.height;
    size_t cBytes = yBytes / 4;
    int frameCount = header.frameCount;
    std::vector<YUVFrame> frames(frameCount);

    if (header.predictionType == PRED_NONE) {
        // 8-bit packed: two samples per int16_t
        std::vector<uint8_t> bytes;
        bytes.reserve(encodedData.size() * 2);
        for (size_t i = 0; i < encodedData.size(); ++i) {
            int16_t val = encodedData[i];
            bytes.push_back(static_cast<uint8_t>(val & 0xFF));
            bytes.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
        }
        size_t byteIdx = 0;
        for (int f = 0; f < frameCount; ++f) {
            frames[f].Y.resize(yBytes);
            frames[f].U.resize(cBytes);
            frames[f].V.resize(cBytes);
            for (size_t i = 0; i < yBytes; ++i)
                frames[f].Y[i] = (byteIdx < bytes.size()) ? bytes[byteIdx++] : 0;
            for (size_t i = 0; i < cBytes; ++i)
                frames[f].U[i] = (byteIdx < bytes.size()) ? bytes[byteIdx++] : 0;
            for (size_t i = 0; i < cBytes; ++i)
                frames[f].V[i] = (byteIdx < bytes.size()) ? bytes[byteIdx++] : 0;
        }
    }
    else {
        // 16-bit packed: one sample per int16_t
        size_t dataIdx = 0;
        for (int f = 0; f < frameCount; ++f) {
            frames[f].Y.resize(yBytes);
            frames[f].U.resize(cBytes);
            frames[f].V.resize(cBytes);
            for (size_t i = 0; i < yBytes; ++i)
                frames[f].Y[i] = encodedData[dataIdx++];
            for (size_t i = 0; i < cBytes; ++i)
                frames[f].U[i] = encodedData[dataIdx++];
            for (size_t i = 0; i < cBytes; ++i)
                frames[f].V[i] = encodedData[dataIdx++];
        }
    }
    return frames;
}


std::vector<YUVFrame> run_entropy_decode(const VidHeader& header, const std::vector<int16_t>& encodedData)
{
    switch (header.entropyType) {
    case ENTROPY_RICE:
        return entropy_decode_rice(header, encodedData);

    case ENTROPY_HUFFMAN:
        return entropy_decode_huffman(header, encodedData);

    case ENTROPY_NONE:
    default:
        return unpack_bitstream(header, encodedData);
    }
}

// Helper to run reverse prediction based on predictionType
static std::vector<YUVFrame> run_reverse_prediction(
    const VidHeader& header,
    const std::vector<YUVFrame>& decodedFrames,
    const std::vector<uint8_t> blockModesFlat)
{
    if (header.predictionType == PRED_NONE) {
        // No prediction, return frames as-is
        return decodedFrames;
    }
    else if (header.predictionType == PRED_BLOCK) {
        // 3. Unflatten blockModes
        int widthY = header.width, heightY = header.height;
        int widthC = widthY / 2, heightC = heightY / 2;
        int blockSize = BLOCK_SIZE, blockSizeC = blockSize / 2;
        int numBlocksY = ((widthY + blockSize - 1) / blockSize) * ((heightY + blockSize - 1) / blockSize);
        int numBlocksC = ((widthC + blockSizeC - 1) / blockSizeC) * ((heightC + blockSizeC - 1) / blockSizeC);
        auto allFramesBlockModes = unflattenBlockModes(blockModesFlat, header.frameCount, numBlocksY, numBlocksC);

        // Block-based prediction, use block modes
        std::vector<YUVFrame> yuvFrames;
        for (int f = 0; f < header.frameCount; ++f) {
            std::vector<YUVFrame> singleFrame{ decodedFrames[f] };
            auto recon = reverse_prediction_auto(header, singleFrame, allFramesBlockModes[f]);
            yuvFrames.push_back(recon[0]);
        }
        return yuvFrames;
    }
    else {
        // Unknown prediction type, treat as no prediction
        return decodedFrames;
    }
}


// Writes YUV frames to a file
bool write_yuv_file(const std::string& outYuvPath, const std::vector<YUVFrame>& yuvFrames)
{
    if (yuvFrames.empty()) return false;
    std::ofstream out(outYuvPath, std::ios::binary);
    if (!out) return false;
    for (const auto& frame : yuvFrames) {
        std::vector<uint8_t> ybuf(frame.Y.size());
        for (size_t i = 0; i < frame.Y.size(); ++i)
            ybuf[i] = static_cast<uint8_t>(std::clamp<int>(frame.Y[i], 0, 255));
        out.write(reinterpret_cast<const char*>(ybuf.data()), ybuf.size());

        std::vector<uint8_t> ubuf(frame.U.size());
        for (size_t i = 0; i < frame.U.size(); ++i)
            ubuf[i] = static_cast<uint8_t>(std::clamp<int>(frame.U[i], 0, 255));
        out.write(reinterpret_cast<const char*>(ubuf.data()), ubuf.size());

        std::vector<uint8_t> vbuf(frame.V.size());
        for (size_t i = 0; i < frame.V.size(); ++i)
            vbuf[i] = static_cast<uint8_t>(std::clamp<int>(frame.V[i], 0, 255));
        out.write(reinterpret_cast<const char*>(vbuf.data()), vbuf.size());
    }
    return true;
}


bool decompress_vid(const std::string& vidPath, const std::string& outYuvPath, int& outW, int& outH, int& outFps)
{
	// 1. Validate Data and Header
    if (vidPath.empty() || outYuvPath.empty())
		return false; // ung�ltige Pfade

    std::ifstream in(vidPath, std::ios::binary);
	if (!in) return false; // Datei nicht gefunden

    std::vector<int16_t> encodedData;
    std::vector<uint8_t> blockModesFlat;
    VidHeader header = read_vid_header(vidPath, encodedData, blockModesFlat);
    if (!blockModesFlat.empty()) blockModesFlat = unpack_block_modes_3bit_prefixed(blockModesFlat);

    // Magic pr�fen
    if (std::memcmp(header.magic, "VID", 3) != 0)
        return false; // falsches Format
    if (header.width <= 0 || header.height <= 0 || header.fps <= 0)
		return false; // ung�ltige Parameter

    // Output parameters
    outW = header.width;
    outH = header.height;
    outFps = header.fps;

    // 2. Entropy decode
    auto frames = run_entropy_decode(header, encodedData);

    // 4. Reverse prediction
    auto yuvFrames = run_reverse_prediction(header, frames, blockModesFlat);

    // temporales Upsampling (Frames wieder auf volle Timeline)
    int frameStep = (header.frameStep <= 0) ? 1 : header.frameStep;
    int originalFrameCount = (header.originalFrameCount <= 0) ? (int)yuvFrames.size() : header.originalFrameCount;
    yuvFrames = temporal_upsample_interpolate(yuvFrames, frameStep, originalFrameCount);

    // 5. Write YUV frames
    return write_yuv_file(outYuvPath, yuvFrames);
}