bwt_compress_v5.hpp

// BWT Compression v5 - Multi-tree Huffman with Symbol Bitmap
//
// Key optimization over v4: Symbol bitmap (bzip2-style)
// - Only store tree lengths for actually used symbols
// - Saves ~2.5 bits per gap symbol per tree
// - Gap symbols = unused values between 0 and max_sym
//
// Format:
// - 16-bit range header: which of 16 ranges (0-15, 16-31, ..., 240-255) are used
// - Per used range: 16-bit sub-bitmap for specific symbols in that range
// - Then: tree lengths ONLY for symbols marked in bitmap
//
// v5.1: Integrated libsais for O(n) BWT construction (was O(n²) naive sort)

#pragma once
#include <vector>
#include <algorithm>
#include <numeric>
#include <cstdint>
#include <cstring>
#include <cstdio>
#include <queue>
#include <functional>

// libsais - O(n) suffix array / BWT construction
// Source: https://github.com/IlyaGrebnov/libsais
extern "C" {
#include "libsais.h"
}

namespace bwt5 {

// Helper: Check if size causes libsais_unbwt crash on MinGW
// Pattern: crashes when (n mod 4 == 2) AND (n >= 30)
// Some edge cases at idx transitions (54, 108, etc.) are safe but we pad anyway
inline bool is_libsais_problematic_size(size_t n) {
    return (n % 4 == 2) && (n >= 30);
}

// Helper: Pad size to next safe value (add bytes until n mod 4 != 2)
inline size_t get_safe_padded_size(size_t n) {
    while (is_libsais_problematic_size(n)) {
        n++;
    }
    return n;
}

constexpr int MAX_ALPHA = 258;  // 256 + RUNA + RUNB
constexpr int GROUP_SIZE = 44;  // Optimized: smaller groups = better local adaptation
constexpr int MAX_TREES = 6;
constexpr int N_ITERS = 4;      // bzip2-style: 4 iterations sufficient

// =============================================================================
// Symbol Bitmap - bzip2-style compact symbol set representation
// Uses 3-byte header to support symbols 0-271 (ZRLE can produce 256+)
// =============================================================================
struct SymbolBitmap {
    uint32_t range_header = 0;  // Which of 17 ranges are used
    uint16_t sub_bitmaps[17] = {0};  // Per-range: which symbols are used
    std::vector<uint16_t> used_symbols;  // Ordered list of used symbols

    void build(const std::vector<uint16_t>& data) {
        range_header = 0;
        memset(sub_bitmaps, 0, sizeof(sub_bitmaps));
        used_symbols.clear();

        // Mark used symbols
        bool used[MAX_ALPHA] = {false};
        for (auto s : data) {
            if (s < MAX_ALPHA) used[s] = true;
        }

        // Build bitmap (17 ranges for symbols 0-271)
        for (int s = 0; s < MAX_ALPHA; s++) {
            if (used[s]) {
                int range = s / 16;
                int bit = s % 16;
                range_header |= (1u << range);
                sub_bitmaps[range] |= (1 << bit);
                used_symbols.push_back(s);
            }
        }
    }

    void write(std::vector<uint8_t>& out) const {
        // Variable-length header:
        // - Byte 0: ranges 0-7
        // - Byte 1: ranges 8-14 (bits 0-6), extension flag (bit 7)
        // - Byte 2 (if extension): ranges 15-16 (bits 0-1)
        // Extension is needed if ranges 15 or 16 are used
        bool need_extension = (range_header & 0x18000) != 0;  // Range 15 or 16 used?

        out.push_back(range_header & 0xFF);  // Ranges 0-7
        uint8_t byte1 = (range_header >> 8) & 0x7F;  // Ranges 8-14 (only 7 bits!)
        if (need_extension) {
            byte1 |= 0x80;  // Set extension flag
        }
        out.push_back(byte1);

        if (need_extension) {
            // Store ranges 15 and 16 in extension byte
            uint8_t ext = ((range_header >> 15) & 0x03);  // Bits 15-16
            out.push_back(ext);
        }

        for (int r = 0; r < 17; r++) {
            if (range_header & (1u << r)) {
                out.push_back(sub_bitmaps[r] & 0xFF);
                out.push_back((sub_bitmaps[r] >> 8) & 0xFF);
            }
        }
    }

    size_t read(const uint8_t* data, size_t max_size) {
        if (max_size < 2) return 0;
        used_symbols.clear();
        memset(sub_bitmaps, 0, sizeof(sub_bitmaps));

        // Variable-length header:
        // - Byte 0: ranges 0-7
        // - Byte 1: ranges 8-14 (bits 0-6), extension flag (bit 7)
        // - Byte 2 (if extension): ranges 15-16 (bits 0-1)
        uint8_t byte0 = data[0];
        uint8_t byte1 = data[1];
        bool need_extension = (byte1 & 0x80) != 0;
        byte1 &= 0x7F;  // Clear extension flag, keep ranges 8-14

        size_t pos = 2;
        if (need_extension) {
            if (max_size < 3) return 0;
            // Ranges 0-7 | ranges 8-14 | ranges 15-16
            range_header = byte0 | ((uint32_t)byte1 << 8) | ((uint32_t)(data[2] & 0x03) << 15);
            pos = 3;
        } else {
            // Ranges 0-7 | ranges 8-14 (no ranges 15-16)
            range_header = byte0 | ((uint32_t)byte1 << 8);
        }

        for (int r = 0; r < 17; r++) {
            if (range_header & (1u << r)) {
                if (pos + 2 > max_size) return 0;
                sub_bitmaps[r] = data[pos] | (data[pos + 1] << 8);
                pos += 2;
                for (int b = 0; b < 16; b++) {
                    if (sub_bitmaps[r] & (1 << b)) {
                        used_symbols.push_back(r * 16 + b);
                    }
                }
            }
        }
        return pos;
    }

    int to_compact(uint16_t sym) const {
        for (size_t i = 0; i < used_symbols.size(); i++) {
            if (used_symbols[i] == sym) return i;
        }
        return -1;
    }

    uint16_t from_compact(int idx) const {
        return (idx >= 0 && idx < (int)used_symbols.size()) ? used_symbols[idx] : 0;
    }
};

// =============================================================================
// BWT Transform - using libsais for O(n) construction
// =============================================================================
inline std::vector<uint8_t> bwt_encode(const uint8_t* data, size_t n, uint32_t& primary_idx) {
    if (n == 0) return {};
    if (n > INT32_MAX) return {};  // libsais uses int32_t for size

    std::vector<uint8_t> bwt(n);
    // Note: Using malloc for A to avoid heap corruption issues with std::vector on MinGW
    int32_t* A = (int32_t*)malloc(n * sizeof(int32_t));
    if (!A) return {};

    // libsais_bwt returns the primary index (position of original string end)
    int32_t idx = libsais_bwt(data, bwt.data(), A, (int32_t)n, 0, nullptr);
    free(A);

    if (idx < 0) {
        // Error - fall back to empty (caller should handle)
        return {};
    }

    primary_idx = (uint32_t)idx;
    return bwt;
}

inline std::vector<uint8_t> bwt_decode(const uint8_t* bwt, size_t n, uint32_t primary_idx, bool debug = false) {
    if (n == 0) return {};
    if (primary_idx >= n) {
        if (debug) printf("ERROR: primary_idx %u >= n %zu\n", primary_idx, n);
        return {};
    }

    if (debug) { printf("bwt_decode: n=%zu, primary_idx=%u\n", n, primary_idx); fflush(stdout); }

    // Use libsais_unbwt for O(n) inverse BWT
    // Note: Using malloc for A to avoid heap corruption issues with std::vector on MinGW
    // Note: Caller must ensure n is a safe size (use is_libsais_problematic_size to check)
    std::vector<uint8_t> output(n);
    int32_t* A = (int32_t*)malloc(n * sizeof(int32_t));
    if (!A) {
        if (debug) printf("ERROR: malloc failed for A array\n");
        return {};
    }

    int32_t result = libsais_unbwt(bwt, output.data(), A, (int32_t)n, nullptr, (int32_t)primary_idx);
    free(A);

    if (result != 0) {
        if (debug) printf("ERROR: libsais_unbwt returned %d\n", result);
        return {};
    }

    return output;
}

// =============================================================================
// MTF Transform (same as v4)
// =============================================================================
inline std::vector<uint8_t> mtf_encode(const uint8_t* data, size_t n) {
    std::vector<uint8_t> alphabet(256);
    std::iota(alphabet.begin(), alphabet.end(), 0);
    std::vector<uint8_t> output(n);
    for (size_t i = 0; i < n; i++) {
        uint8_t c = data[i];
        uint8_t pos = 0;
        while (alphabet[pos] != c) pos++;
        output[i] = pos;
        for (int j = pos; j > 0; j--) alphabet[j] = alphabet[j-1];
        alphabet[0] = c;
    }
    return output;
}

inline std::vector<uint8_t> mtf_decode(const uint8_t* data, size_t n) {
    std::vector<uint8_t> alphabet(256);
    std::iota(alphabet.begin(), alphabet.end(), 0);
    std::vector<uint8_t> output(n);
    for (size_t i = 0; i < n; i++) {
        uint8_t pos = data[i];
        uint8_t c = alphabet[pos];
        output[i] = c;
        for (int j = pos; j > 0; j--) alphabet[j] = alphabet[j-1];
        alphabet[0] = c;
    }
    return output;
}

// =============================================================================
// Zero-RLE (same as v4)
// =============================================================================
inline std::vector<uint16_t> zrle_encode(const uint8_t* data, size_t n) {
    std::vector<uint16_t> output;
    size_t i = 0;
    while (i < n) {
        if (data[i] == 0) {
            size_t run = 0;
            while (i < n && data[i] == 0) { run++; i++; }
            while (run > 0) {
                if (run & 1) {
                    output.push_back(0);
                    run = (run - 1) / 2;
                } else {
                    output.push_back(1);
                    run = (run - 2) / 2;
                }
            }
        } else {
            output.push_back(data[i] + 1);
            i++;
        }
    }
    return output;
}

inline std::vector<uint8_t> zrle_decode(const uint16_t* data, size_t n) {
    std::vector<uint8_t> output;
    size_t i = 0;
    while (i < n) {
        if (data[i] <= 1) {
            size_t run = 0;
            size_t power = 1;
            while (i < n && data[i] <= 1) {
                run += (data[i] + 1) * power;
                power *= 2;
                i++;
            }
            for (size_t j = 0; j < run; j++) output.push_back(0);
        } else {
            output.push_back(data[i] - 1);
            i++;
        }
    }
    return output;
}

// =============================================================================
// Pre-BWT Run-Length Encoding (bzip2-style)
// =============================================================================
// Encodes runs of 4+ identical bytes as: byte byte byte byte (count-4)
// where count-4 is 0-255, meaning runs of 4-259 bytes
inline std::vector<uint8_t> pre_rle_encode(const uint8_t* data, size_t n) {
    if (n == 0) return {};

    std::vector<uint8_t> output;
    output.reserve(n);

    size_t i = 0;
    while (i < n) {
        uint8_t byte = data[i];
        size_t run = 1;
        while (i + run < n && data[i + run] == byte && run < 259) {
            run++;
        }

        if (run >= 4) {
            // Encode as: byte byte byte byte (run-4)
            output.push_back(byte);
            output.push_back(byte);
            output.push_back(byte);
            output.push_back(byte);
            output.push_back(run - 4);
            i += run;
        } else {
            // Output bytes literally
            for (size_t j = 0; j < run; j++) {
                output.push_back(byte);
            }
            i += run;
        }
    }

    return output;
}

inline std::vector<uint8_t> pre_rle_decode(const uint8_t* data, size_t n) {
    if (n == 0) return {};

    std::vector<uint8_t> output;
    output.reserve(n * 2);

    size_t i = 0;
    while (i < n) {
        uint8_t byte = data[i];

        // Check for run encoding: 4 identical bytes followed by count
        if (i + 4 < n &&
            data[i+1] == byte &&
            data[i+2] == byte &&
            data[i+3] == byte) {
            size_t run = 4 + data[i+4];
            for (size_t j = 0; j < run; j++) {
                output.push_back(byte);
            }
            i += 5;
        } else {
            output.push_back(byte);
            i++;
        }
    }

    return output;
}

// =============================================================================
// Huffman Tree (same as v4)
// =============================================================================
struct HuffTree {
    int code_lengths[MAX_ALPHA] = {0};
    uint32_t codes[MAX_ALPHA] = {0};
    int max_length = 0;

    void build_from_freqs(const uint32_t* freqs, int alpha_size) {
        // bzip2-style: scale frequencies and rebuild if maxLen exceeded
        std::vector<uint32_t> scaled_freqs(freqs, freqs + alpha_size);

        for (int attempt = 0; attempt < 10; attempt++) {
            struct Node { int symbol; uint32_t freq; int left, right; };
            std::vector<Node> nodes;
            std::priority_queue<std::pair<uint32_t, int>, std::vector<std::pair<uint32_t, int>>, std::greater<>> pq;

            for (int i = 0; i < alpha_size; i++) {
                if (scaled_freqs[i] > 0) {
                    int idx = nodes.size();
                    nodes.push_back({i, scaled_freqs[i], -1, -1});
                    pq.push({scaled_freqs[i], idx});
                }
            }

            if (pq.empty()) return;
            if (pq.size() == 1) {
                code_lengths[nodes[pq.top().second].symbol] = 1;
                codes[nodes[pq.top().second].symbol] = 0;
                max_length = 1;
                break;
            }

            while (pq.size() > 1) {
                auto [f1, i1] = pq.top(); pq.pop();
                auto [f2, i2] = pq.top(); pq.pop();
                int idx = nodes.size();
                nodes.push_back({-1, f1 + f2, i1, i2});
                pq.push({f1 + f2, idx});
            }

            max_length = 0;
            memset(code_lengths, 0, sizeof(code_lengths));

            std::function<void(int, int)> assign_lengths = [&](int idx, int depth) {
                if (nodes[idx].symbol >= 0) {
                    code_lengths[nodes[idx].symbol] = depth;
                    max_length = std::max(max_length, depth);
                } else {
                    assign_lengths(nodes[idx].left, depth + 1);
                    assign_lengths(nodes[idx].right, depth + 1);
                }
            };
            assign_lengths(pq.top().second, 0);

            // bzip2: if maxLen > 17, scale frequencies and rebuild
            if (max_length <= 17) break;

            // Scale all frequencies: freq = 1 + freq/2 (ensures min freq = 1)
            for (int i = 0; i < alpha_size; i++) {
                if (scaled_freqs[i] > 0) {
                    scaled_freqs[i] = 1 + scaled_freqs[i] / 2;
                }
            }
        }

        // Final clamp just in case
        for (int i = 0; i < alpha_size; i++) {
            if (code_lengths[i] > 17) code_lengths[i] = 17;
        }
        max_length = std::min(max_length, 17);

        std::vector<std::pair<int, int>> syms;
        for (int i = 0; i < alpha_size; i++) {
            if (code_lengths[i] > 0) syms.push_back({code_lengths[i], i});
        }
        std::sort(syms.begin(), syms.end());

        uint32_t code = 0;
        int prev_len = 0;
        for (auto [len, sym] : syms) {
            code <<= (len - prev_len);
            codes[sym] = code++;
            prev_len = len;
        }
    }

    void build_from_lengths(const int* lengths, int alpha_size) {
        if (alpha_size <= 0 || alpha_size > 258) return;  // Bounds check (max 256 + 2 special)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wstringop-overflow"
        memcpy(code_lengths, lengths, static_cast<size_t>(alpha_size) * sizeof(int));
#pragma GCC diagnostic pop
        max_length = 0;
        for (int i = 0; i < alpha_size; i++) {
            max_length = std::max(max_length, code_lengths[i]);
        }

        std::vector<std::pair<int, int>> syms;
        for (int i = 0; i < alpha_size; i++) {
            if (code_lengths[i] > 0) syms.push_back({code_lengths[i], i});
        }
        std::sort(syms.begin(), syms.end());

        uint32_t code = 0;
        int prev_len = 0;
        for (auto [len, sym] : syms) {
            code <<= (len - prev_len);
            codes[sym] = code++;
            prev_len = len;
        }
    }
};

// =============================================================================
// Multi-Tree Encoder with Bitmap (key change from v4)
// =============================================================================
class MultiTreeEncoder {
    std::vector<uint8_t>& output;
    uint32_t buffer = 0;
    int bits_in_buffer = 0;

    void write_bits(uint32_t value, int n_bits) {
        buffer |= (value << bits_in_buffer);
        bits_in_buffer += n_bits;
        while (bits_in_buffer >= 8) {
            output.push_back(buffer & 0xFF);
            buffer >>= 8;
            bits_in_buffer -= 8;
        }
    }

    // Unary encoding: 0=0, 1=10, 2=110, 3=1110, etc.
    void write_unary(int value) {
        for (int i = 0; i < value; i++) {
            write_bits(1, 1);
        }
        write_bits(0, 1);
    }

public:
    MultiTreeEncoder(std::vector<uint8_t>& out) : output(out) {}

    void encode(const std::vector<uint16_t>& data, int alpha_size) {
        if (data.empty()) return;

        SymbolBitmap bitmap;
        bitmap.build(data);
        int compact_size = bitmap.used_symbols.size();

        bitmap.write(output);

        // Optimized tree count based on data size
        int n_trees;
        size_t n = data.size();
        if (n < 200) n_trees = 2;
        else if (n < 600) n_trees = 3;
        else if (n < 1200) n_trees = 4;
        else n_trees = 6;  // 6 trees optimal for all sizes >= 1200

        int n_groups = (n + GROUP_SIZE - 1) / GROUP_SIZE;
        std::vector<std::vector<uint32_t>> tree_freqs(n_trees, std::vector<uint32_t>(compact_size, 0));
        std::vector<int> group_tree(n_groups, 0);

        // Step 1: Compute global symbol frequencies
        std::vector<uint32_t> sym_freq(compact_size, 0);
        for (size_t i = 0; i < n; i++) {
            int idx = bitmap.to_compact(data[i]);
            if (idx >= 0) sym_freq[idx]++;
        }

        // Step 2: bzip2-style initial seeding - divide SYMBOLS by cumulative frequency
        // Each tree gets a contiguous range of symbols with roughly equal total frequency
        std::vector<int> sym_tree_owner(compact_size, n_trees - 1);  // Default to last tree
        {
            int nPart = n_trees;
            uint64_t remF = n;  // Total symbol count
            int gs = 0;

            while (nPart > 0 && gs < compact_size) {
                uint64_t tFreq = remF / nPart;  // Target frequency for this partition
                int ge = gs - 1;
                uint64_t aFreq = 0;

                // Accumulate symbols until we reach target frequency
                while (aFreq < tFreq && ge < compact_size - 1) {
                    ge++;
                    aFreq += sym_freq[ge];
                }

                // Assign symbols [gs, ge] to this tree
                for (int v = gs; v <= ge; v++) {
                    sym_tree_owner[v] = nPart - 1;
                }

                nPart--;
                gs = ge + 1;
                remF -= aFreq;
            }
        }

        // Step 3: Initialize tree costs based on symbol ownership
        // BZ_LESSER_ICOST = 0, BZ_GREATER_ICOST = 15
        std::vector<std::vector<int>> initial_costs(n_trees, std::vector<int>(compact_size, 15));
        for (int s = 0; s < compact_size; s++) {
            initial_costs[sym_tree_owner[s]][s] = 0;
        }

        // Step 4: Initial group assignment based on lowest cost tree
        for (int g = 0; g < n_groups; g++) {
            size_t start = g * GROUP_SIZE;
            size_t end = std::min(start + GROUP_SIZE, n);

            int best_tree = 0;
            int best_cost = INT_MAX;

            for (int t = 0; t < n_trees; t++) {
                int cost = 0;
                for (size_t i = start; i < end; i++) {
                    int idx = bitmap.to_compact(data[i]);
                    if (idx >= 0) cost += initial_costs[t][idx];
                }
                if (cost < best_cost) {
                    best_cost = cost;
                    best_tree = t;
                }
            }
            group_tree[g] = best_tree;
        }

        for (int iter = 0; iter < N_ITERS; iter++) {
            for (auto& tf : tree_freqs) std::fill(tf.begin(), tf.end(), 0);

            for (int g = 0; g < n_groups; g++) {
                int t = group_tree[g];
                size_t start = g * GROUP_SIZE;
                size_t end = std::min(start + GROUP_SIZE, n);
                for (size_t i = start; i < end; i++) {
                    int compact_idx = bitmap.to_compact(data[i]);
                    if (compact_idx >= 0) tree_freqs[t][compact_idx]++;
                }
            }

            for (auto& tf : tree_freqs) {
                for (auto& f : tf) if (f == 0) f = 1;
            }

            std::vector<HuffTree> trees(n_trees);
            for (int t = 0; t < n_trees; t++) {
                trees[t].build_from_freqs(tree_freqs[t].data(), compact_size);
            }

            for (int g = 0; g < n_groups; g++) {
                size_t start = g * GROUP_SIZE;
                size_t end = std::min(start + GROUP_SIZE, n);

                int best_tree = 0;
                size_t best_cost = SIZE_MAX;

                for (int t = 0; t < n_trees; t++) {
                    size_t cost = 0;
                    for (size_t i = start; i < end; i++) {
                        int compact_idx = bitmap.to_compact(data[i]);
                        if (compact_idx >= 0) {
                            cost += trees[t].code_lengths[compact_idx];
                        }
                    }
                    if (cost < best_cost) {
                        best_cost = cost;
                        best_tree = t;
                    }
                }
                group_tree[g] = best_tree;
            }
        }

        for (auto& tf : tree_freqs) std::fill(tf.begin(), tf.end(), 0);
        for (int g = 0; g < n_groups; g++) {
            int t = group_tree[g];
            size_t start = g * GROUP_SIZE;
            size_t end = std::min(start + GROUP_SIZE, n);
            for (size_t i = start; i < end; i++) {
                int compact_idx = bitmap.to_compact(data[i]);
                if (compact_idx >= 0) tree_freqs[t][compact_idx]++;
            }
        }
        for (auto& tf : tree_freqs) {
            for (auto& f : tf) if (f == 0) f = 1;
        }

        std::vector<HuffTree> trees(n_trees);
        for (int t = 0; t < n_trees; t++) {
            trees[t].build_from_freqs(tree_freqs[t].data(), compact_size);
        }

#ifdef DEBUG_V5
        printf("Encoder: n_trees=%d, n_groups=%d, compact_size=%d\n", n_trees, n_groups, compact_size);
        for (int t = 0; t < n_trees; t++) {
            printf("Tree %d: max_len=%d, lengths: ", t, trees[t].max_length);
            for (int s = 0; s < compact_size; s++) printf("%d ", trees[t].code_lengths[s]);
            printf("\n");
        }
#endif

        write_bits(n_trees - 1, 3);
        write_bits(n_groups, 16);

        // Selector MTF + unary encoding (bzip2-style)
        // 97% of consecutive groups use same tree → MTF[0] dominates → ~1 bit avg
        std::vector<int> selector_mtf(n_trees);
        std::iota(selector_mtf.begin(), selector_mtf.end(), 0);

        for (int g = 0; g < n_groups; g++) {
            int tree = group_tree[g];

            // Find position in MTF list
            int pos = 0;
            while (selector_mtf[pos] != tree) pos++;

            // Write unary-encoded position
            write_unary(pos);

            // Move to front
            for (int j = pos; j > 0; j--) {
                selector_mtf[j] = selector_mtf[j - 1];
            }
            selector_mtf[0] = tree;
        }

        for (int t = 0; t < n_trees; t++) {
            int curr_len = trees[t].code_lengths[0];
            write_bits(curr_len, 5);

            for (int s = 0; s < compact_size; s++) {
                int len = trees[t].code_lengths[s];
                while (curr_len < len) { write_bits(1, 2); curr_len++; }
                while (curr_len > len) { write_bits(3, 2); curr_len--; }
                write_bits(0, 1);
            }
        }

        int curr_group = -1;
        int curr_tree = 0;
        for (size_t i = 0; i < n; i++) {
            int g = i / GROUP_SIZE;
            if (g != curr_group) {
                curr_group = g;
                curr_tree = group_tree[g];
            }
            int compact_idx = bitmap.to_compact(data[i]);
            if (compact_idx >= 0) {
                // Reverse code bits for LSB-first writing
                uint32_t code = trees[curr_tree].codes[compact_idx];
                int len = trees[curr_tree].code_lengths[compact_idx];
                uint32_t rev_code = 0;
                for (int b = 0; b < len; b++) {
                    rev_code = (rev_code << 1) | ((code >> b) & 1);
                }
#ifdef DEBUG_V5
                if (i < 5) {
                    printf("  encode[%zu]: full_sym=%d, compact=%d, code=%u->%u, len=%d\n",
                           i, data[i], compact_idx, code, rev_code, len);
                }
#endif
                write_bits(rev_code, len);
            }
        }

        if (bits_in_buffer > 0) {
            output.push_back(buffer & 0xFF);
        }
    }
};

// =============================================================================
// Multi-Tree Decoder with Bitmap
// =============================================================================
class MultiTreeDecoder {
    const uint8_t* data;
    size_t size;
    size_t byte_pos = 0;
    uint32_t buffer = 0;
    int bits_in_buffer = 0;

    void ensure_bits(int n) {
        while (bits_in_buffer < n && byte_pos < size) {
            buffer |= ((uint32_t)data[byte_pos++] << bits_in_buffer);
            bits_in_buffer += 8;
        }
    }

    uint32_t read_bits(int n) {
        ensure_bits(n);
        uint32_t val = buffer & ((1u << n) - 1);
        buffer >>= n;
        bits_in_buffer -= n;
        return val;
    }

    // Unary decoding: count 1s until 0
    int read_unary() {
        int count = 0;
        while (read_bits(1) == 1) {
            count++;
        }
        return count;
    }

public:
    MultiTreeDecoder(const uint8_t* d, size_t n) : data(d), size(n) {}

    std::vector<uint16_t> decode(size_t expected_count, bool debug = false) {
        SymbolBitmap bitmap;
        size_t bitmap_size = bitmap.read(data, size);
        if (bitmap_size == 0) return {};

        byte_pos = bitmap_size;
        int compact_size = bitmap.used_symbols.size();

        int n_trees = read_bits(3) + 1;
        int n_groups = read_bits(16);

        if (debug) printf("Decoder: n_trees=%d, n_groups=%d, compact_size=%d\n", n_trees, n_groups, compact_size);

        // Decode selectors with MTF + unary
        std::vector<int> selector_mtf(n_trees);
        std::iota(selector_mtf.begin(), selector_mtf.end(), 0);

        std::vector<int> group_tree(n_groups);
        for (int g = 0; g < n_groups; g++) {
            int pos = read_unary();
            int tree = selector_mtf[pos];
            group_tree[g] = tree;

            // Move to front
            for (int j = pos; j > 0; j--) {
                selector_mtf[j] = selector_mtf[j - 1];
            }
            selector_mtf[0] = tree;
        }

        std::vector<HuffTree> trees(n_trees);
        for (int t = 0; t < n_trees; t++) {
            int curr_len = read_bits(5);
            std::vector<int> lengths(compact_size);

            for (int s = 0; s < compact_size; s++) {
                while (true) {
                    uint32_t bit = read_bits(1);
                    if (bit == 0) break;
                    bit = read_bits(1);
                    if (bit == 0) curr_len++;
                    else curr_len--;
                }
                lengths[s] = curr_len;
            }
            trees[t].build_from_lengths(lengths.data(), compact_size);
            if (debug) {
                printf("Tree %d: max_len=%d, lengths: ", t, trees[t].max_length);
                for (int s = 0; s < compact_size; s++) printf("%d ", trees[t].code_lengths[s]);
                printf("\n");
            }
        }

        struct DecodeEntry { int symbol; int length; };
        std::vector<std::vector<DecodeEntry>> decode_tables(n_trees);

        for (int t = 0; t < n_trees; t++) {
            int max_len = trees[t].max_length;
            if (max_len == 0) max_len = 1;
            decode_tables[t].resize(1 << max_len);

            if (debug) {
                printf("Building decode table for tree %d, max_len=%d\n", t, max_len);
                for (int s = 0; s < compact_size; s++) {
                    printf("  sym %d: len=%d code=%u\n", s, trees[t].code_lengths[s], trees[t].codes[s]);
                }
            }

            for (int s = 0; s < compact_size; s++) {
                int len = trees[t].code_lengths[s];
                if (len > 0) {
                    uint32_t code = trees[t].codes[s];
                    uint32_t rev_code = 0;
                    for (int b = 0; b < len; b++) {
                        rev_code = (rev_code << 1) | ((code >> b) & 1);
                    }

                    int fill = 1 << (max_len - len);
                    for (int f = 0; f < fill; f++) {
                        uint32_t idx = rev_code | (f << len);
                        if (debug && s < 3) {
                            printf("  table[%u] = {%d, %d} (code=%u->rev=%u, f=%d)\n", idx, s, len, code, rev_code, f);
                        }
                        decode_tables[t][idx] = {s, len};
                    }
                }
            }
        }

        std::vector<uint16_t> output;
        output.reserve(expected_count);

        int curr_group = -1;
        int curr_tree = 0;

        for (size_t i = 0; i < expected_count; i++) {
            int g = i / GROUP_SIZE;
            if (g != curr_group) {
                curr_group = g;
                curr_tree = group_tree[g];
            }

            int max_len = trees[curr_tree].max_length;
            if (max_len == 0) max_len = 1;
            ensure_bits(max_len);

            uint32_t peek = buffer & ((1u << max_len) - 1);
            auto& entry = decode_tables[curr_tree][peek];

            if (debug && i < 5) {
                printf("  decode[%zu]: peek=%u, sym=%d, len=%d, full_sym=%d\n",
                       i, peek, entry.symbol, entry.length, bitmap.from_compact(entry.symbol));
            }

            output.push_back(bitmap.from_compact(entry.symbol));
            buffer >>= entry.length;
            bits_in_buffer -= entry.length;
        }

        return output;
    }

    // Decode until EOB symbol (highest symbol in alphabet)
    std::vector<uint16_t> decode_until_eob(bool debug = false) {
        if (debug) { printf("decode_until_eob called: size=%zu\n", size); fflush(stdout); }
        SymbolBitmap bitmap;
        size_t bitmap_size = bitmap.read(data, size);
        if (debug) { printf("bitmap_size=%zu\n", bitmap_size); fflush(stdout); }
        if (bitmap_size == 0) return {};

        byte_pos = bitmap_size;
        int compact_size = bitmap.used_symbols.size();
        if (debug) { printf("compact_size=%d\n", compact_size); fflush(stdout); }

        // EOB is the highest symbol
        if (compact_size == 0) return {};
        uint16_t eob_sym = bitmap.used_symbols.back();
        if (debug) { printf("eob_sym=%d\n", eob_sym); fflush(stdout); }

        int n_trees = read_bits(3) + 1;
        int n_groups = read_bits(16);

        if (debug) { printf("Decoder EOB: n_trees=%d, n_groups=%d, compact_size=%d, eob=%d\n",
                          n_trees, n_groups, compact_size, eob_sym); fflush(stdout); }

        // Decode selectors
        if (debug) { printf("Decoding %d selectors...\n", n_groups); fflush(stdout); }
        std::vector<int> selector_mtf(n_trees);
        std::iota(selector_mtf.begin(), selector_mtf.end(), 0);

        std::vector<int> group_tree(n_groups);
        for (int g = 0; g < n_groups; g++) {
            int pos = read_unary();
            if (pos >= n_trees) {
                if (debug) { printf("ERROR: selector pos %d >= n_trees %d at group %d\n", pos, n_trees, g); fflush(stdout); }
                return {};
            }
            int tree = selector_mtf[pos];
            group_tree[g] = tree;
            for (int j = pos; j > 0; j--) selector_mtf[j] = selector_mtf[j - 1];
            selector_mtf[0] = tree;
        }
        if (debug) { printf("Selectors decoded\n"); fflush(stdout); }

        // Build trees
        if (debug) { printf("Building %d trees...\n", n_trees); fflush(stdout); }
        std::vector<HuffTree> trees(n_trees);
        for (int t = 0; t < n_trees; t++) {
            int curr_len = read_bits(5);
            std::vector<int> lengths(compact_size);
            for (int s = 0; s < compact_size; s++) {
                while (true) {
                    uint32_t bit = read_bits(1);
                    if (bit == 0) break;
                    bit = read_bits(1);
                    if (bit == 0) curr_len++;
                    else curr_len--;
                }
                lengths[s] = curr_len;
            }
            trees[t].build_from_lengths(lengths.data(), compact_size);
            if (debug) { printf("Tree %d: max_len=%d\n", t, trees[t].max_length); fflush(stdout); }
        }
        if (debug) { printf("Trees built\n"); fflush(stdout); }

        // Build decode tables
        if (debug) { printf("Building decode tables...\n"); fflush(stdout); }
        struct DecodeEntry { int symbol; int length; };
        std::vector<std::vector<DecodeEntry>> decode_tables(n_trees);

        for (int t = 0; t < n_trees; t++) {
            int max_len = trees[t].max_length;
            if (max_len == 0) max_len = 1;
            if (max_len > 20) {
                if (debug) { printf("ERROR: tree %d max_len %d too large!\n", t, max_len); fflush(stdout); }
                return {};
            }
            decode_tables[t].resize(1 << max_len);

            for (int s = 0; s < compact_size; s++) {
                int len = trees[t].code_lengths[s];
                if (len > 0) {
                    uint32_t code = trees[t].codes[s];
                    uint32_t rev_code = 0;
                    for (int b = 0; b < len; b++) {
                        rev_code = (rev_code << 1) | ((code >> b) & 1);
                    }
                    int fill = 1 << (max_len - len);
                    for (int f = 0; f < fill; f++) {
                        uint32_t idx = rev_code | (f << len);
                        decode_tables[t][idx] = {s, len};
                    }
                }
            }
            if (debug) { printf("Decode table %d: size=%zu\n", t, decode_tables[t].size()); fflush(stdout); }
        }
        if (debug) { printf("Decode tables built\n"); fflush(stdout); }

        // Decode until EOB
        if (debug) { printf("Starting decode loop (expected ~%d symbols)...\n", n_groups * GROUP_SIZE); fflush(stdout); }
        std::vector<uint16_t> output;
        output.reserve(n_groups * GROUP_SIZE);

        int curr_group = -1;
        int curr_tree = 0;
        size_t i = 0;

        while (true) {
            int g = i / GROUP_SIZE;
            if (g != curr_group && g < n_groups) {
                curr_group = g;
                curr_tree = group_tree[g];
            }

            int max_len = trees[curr_tree].max_length;
            if (max_len == 0) max_len = 1;
            ensure_bits(max_len);

            uint32_t peek = buffer & ((1u << max_len) - 1);
            auto& entry = decode_tables[curr_tree][peek];

            // Debug: check for invalid entry
            if (entry.length == 0) {
                if (debug) { printf("ERROR at i=%zu: entry.length=0, peek=%u, tree=%d\n", i, peek, curr_tree); fflush(stdout); }
                return {};
            }

            uint16_t sym = bitmap.from_compact(entry.symbol);

            if (sym == eob_sym) {
                if (debug) { printf("EOB found at i=%zu\n", i); fflush(stdout); }
                break;
            }

            output.push_back(sym);
            buffer >>= entry.length;
            bits_in_buffer -= entry.length;
            i++;

            // Debug: periodic progress
            if (debug && (i % 5000 == 0)) { printf("Decoded %zu symbols...\n", i); fflush(stdout); }

            if (i > 100000000) break;  // Safety limit
        }

        return output;
    }
};

// =============================================================================
// Main Compress/Decompress
// =============================================================================
// Varint encoding: 7 bits per byte, high bit = continuation
inline void write_varint(std::vector<uint8_t>& out, uint32_t val) {
    while (val >= 128) {
        out.push_back((val & 0x7F) | 0x80);
        val >>= 7;
    }
    out.push_back(val & 0x7F);
}

inline uint32_t read_varint(const uint8_t* data, size_t max_size, size_t& pos) {
    uint32_t val = 0;
    int shift = 0;
    while (pos < max_size && shift < 35) {
        uint8_t b = data[pos++];
        val |= (uint32_t)(b & 0x7F) << shift;
        if (!(b & 0x80)) break;
        shift += 7;
    }
    return val;
}

inline std::vector<uint8_t> compress(const uint8_t* data, size_t n) {
    if (n == 0) return {};

    // Step 1: Pre-BWT RLE (bzip2-style)
    auto rle1 = pre_rle_encode(data, n);
    uint32_t original_rle1_size = rle1.size();

    // Step 1.5: Pad to safe size if needed (avoids libsais_unbwt crash on MinGW)
    if (is_libsais_problematic_size(rle1.size())) {
        size_t safe_size = get_safe_padded_size(rle1.size());
        rle1.resize(safe_size, 0);  // Pad with zeros
    }

    // Step 2: BWT on (possibly padded) RLE1 output
    uint32_t primary_idx;
    auto bwt = bwt_encode(rle1.data(), rle1.size(), primary_idx);

    // Step 3: MTF
    auto mtf = mtf_encode(bwt.data(), bwt.size());

    // Step 4: Zero-RLE (RUNA/RUNB)
    auto rle = zrle_encode(mtf.data(), mtf.size());

    if (rle.empty()) return {};

    int max_sym = 0;
    for (auto s : rle) if (s > max_sym) max_sym = s;

    // Add EOB symbol (eliminates need to store rle_count)
    uint16_t eob_sym = max_sym + 1;
    rle.push_back(eob_sym);
    int alpha_size = eob_sym + 1;

    std::vector<uint8_t> output;

    // Header: 'BB' for v11 (compact header with EOB)
    output.push_back('B');
    output.push_back('B');

    // Varint: RLE1 size (original, before padding - typically 3 bytes for 256KB)
    write_varint(output, original_rle1_size);

    // Varint: Primary index (typically 3 bytes)
    write_varint(output, primary_idx);

    MultiTreeEncoder encoder(output);
    encoder.encode(rle, alpha_size);

    return output;
}

inline std::vector<uint8_t> decompress(const uint8_t* data, size_t n, bool debug = false) {
    if (debug) { printf("bwt5::decompress called: n=%zu, header=%c%c\n", n, data[0], data[1]); fflush(stdout); }
    if (n < 4) return {};

    // Check header format
    if (data[0] == 'B' && data[1] == 'B') {
        // New compact format with EOB
        if (debug) { printf("Matched BB header\n"); fflush(stdout); }
        size_t pos = 2;
        uint32_t rle1_size = read_varint(data, n, pos);
        uint32_t primary_idx = read_varint(data, n, pos);

        if (debug) { printf("Compact format: rle1_size=%u, primary_idx=%u, data_start=%zu\n",
                          rle1_size, primary_idx, pos); fflush(stdout); }

        MultiTreeDecoder decoder(data + pos, n - pos);
        auto rle = decoder.decode_until_eob(debug);

        if (debug) {
            printf("Decoded RLE (%zu symbols, excluding EOB): ", rle.size());
            for (size_t i = 0; i < std::min(rle.size(), (size_t)20); i++)
                printf("%d ", rle[i]);
            if (rle.size() > 20) printf("...");
            printf("\n"); fflush(stdout);
        }

        // ZRLE decode
        if (debug) { printf("Starting ZRLE decode...\n"); fflush(stdout); }
        auto mtf_data = zrle_decode(rle.data(), rle.size());
        if (debug) { printf("ZRLE decoded: %zu bytes\n", mtf_data.size()); fflush(stdout); }

        // MTF decode
        if (debug) { printf("Starting MTF decode...\n"); fflush(stdout); }
        auto bwt_data = mtf_decode(mtf_data.data(), mtf_data.size());
        if (debug) { printf("MTF decoded: %zu bytes\n", bwt_data.size()); fflush(stdout); }

        // BWT decode
        if (debug) { printf("Starting BWT decode (primary_idx=%u)...\n", primary_idx); fflush(stdout); }
        auto rle1_data = bwt_decode(bwt_data.data(), bwt_data.size(), primary_idx, debug);
        if (debug) { printf("BWT decoded: %zu bytes (expected %u)\n", rle1_data.size(), rle1_size); fflush(stdout); }

        // Truncate to original size (remove padding added during compression)
        if (rle1_data.size() > rle1_size) {
            if (debug) { printf("Truncating from %zu to %u bytes (removing padding)\n", rle1_data.size(), rle1_size); fflush(stdout); }
            rle1_data.resize(rle1_size);
        }

        // Pre-RLE decode
        if (debug) { printf("Starting Pre-RLE decode...\n"); fflush(stdout); }
        auto output = pre_rle_decode(rle1_data.data(), rle1_data.size());
        if (debug) { printf("Pre-RLE decoded: %zu bytes\n", output.size()); fflush(stdout); }

        return output;
    }
    else if (data[0] == 'B' && data[1] == 'A') {
        // Old format (for backwards compatibility)
        if (n < 20) return {};

        uint32_t orig_size = data[2] | (data[3] << 8) | (data[4] << 16) | (data[5] << 24);
        uint32_t rle1_size = data[6] | (data[7] << 8) | (data[8] << 16) | (data[9] << 24);
        uint32_t primary_idx = data[10] | (data[11] << 8) | (data[12] << 16) | (data[13] << 24);
        uint32_t rle_count = data[14] | (data[15] << 8) | (data[16] << 16) | (data[17] << 24);

        size_t encoded_start = 20;
        MultiTreeDecoder decoder(data + encoded_start, n - encoded_start);
        auto rle = decoder.decode(rle_count, debug);

        if (rle.size() != rle_count) return {};

        auto mtf_data = zrle_decode(rle.data(), rle.size());
        auto bwt_data = mtf_decode(mtf_data.data(), mtf_data.size());
        auto rle1_data = bwt_decode(bwt_data.data(), bwt_data.size(), primary_idx);
        auto output = pre_rle_decode(rle1_data.data(), rle1_data.size());

        return output;
    }

    return {};
}

} // namespace bwt5