MIL-STD-188-110C/include/encoder/Interleaver.h

#ifndef INTERLEAVER_H
#define INTERLEAVER_H

#include <algorithm>
#include <cstdint>
#include <stdexcept>
#include <vector>

#include "bitstream/bitstream.h"

/**
 * @class Interleaver
 * @brief A class to interleave input bitstreams based on the specified baud rate and interleave settings.
 * 
 * The Interleaver processes input data in chunks, rearranging the bit order as required 
 * by the MIL-STD-188-110 standard for different baud rates and interleave settings.
 */
class Interleaver {
public:
    /**
     * @brief Constructs an Interleaver with specific baud rate, interleave settings, and operation mode.
     * @param baud_rate The baud rate of the transmission.
     * @param interleave_setting The interleave setting (0, short, or long).
     * @param is_frequency_hopping Boolean flag indicating if frequency-hopping mode is enabled.
     */
    Interleaver(size_t baud_rate, size_t interleave_setting, bool is_frequency_hopping)
        : baud_rate(baud_rate), interleave_setting(interleave_setting), is_frequency_hopping(is_frequency_hopping) {
        setMatrixDimensions();
        matrix.resize(rows * columns, 0);
    }

    /**
     * @brief Interleaves the entire input BitStream and returns the interleaved result.
     * @param input_data The input BitStream to be interleaved.
     * @return A new BitStream containing the interleaved data.
     */
    std::vector<uint8_t> interleaveStream(bitstream::fixed_bit_reader& input_data) {
        std::vector<uint8_t> interleaved_buffer;

        bitstream::growing_bit_writer<std::vector<uint8_t>> interleaved_data(interleaved_buffer);

        size_t chunk_size = rows * columns;
        size_t input_index = 0;

        std::vector<uint8_t> chunk_data((chunk_size + 7) / 8, 0);

        while (input_data.get_remaining_bits() >= chunk_size) {
            std::fill(chunk_data.begin(), chunk_data.end(), 0);

            if (!input_data.serialize_bytes(chunk_data.data(), chunk_size)) {
                throw std::runtime_error("Failed to serialize chunk from input data");
            }

            bitstream::fixed_bit_reader chunk_reader(chunk_data.data(), chunk_size);

            interleaveChunk(interleaved_data, chunk_reader);
        }

        // Apply puncturing for 2400 bps in frequency-hopping mode (Rate 2/3)
        if (baud_rate == 2400 && is_frequency_hopping) {
            std::vector<uint8_t> punctured_buffer;

            bitstream::growing_bit_writer<std::vector<uint8_t>> punctured_writer(punctured_buffer);
            bitstream::fixed_bit_reader interleaved_reader(interleaved_buffer.data(), interleaved_buffer.size() * 8);

            applyPuncturing(punctured_writer, interleaved_reader);

            interleaved_buffer = punctured_buffer;
        }

        bitstream::fixed_bit_reader final_reader(interleaved_buffer.data(), interleaved_buffer.size() * 8);

        return groupSymbols(final_reader);
    }

    /**
     * @brief Retrieves the number of bits required to fully flush the interleaver matrix.
     * @return The number of bits needed for a complete flush.
     */
    size_t getFlushBits() const {
        return (interleave_setting == 0) ? 0 : (rows * columns);
    }

private:
    size_t baud_rate;
    size_t interleave_setting;
    bool is_frequency_hopping;
    size_t rows;
    size_t columns;
    std::vector<uint8_t> matrix; ///< 1D vector representing the interleaver matrix.

    static constexpr size_t ROW_INCREMENT_DEFAULT = 9;
    static constexpr size_t COLUMN_DECREMENT_DEFAULT = 17;

    /**
     * @brief Groups the bits into symbols based on the baud rate (e.g., 2 bits per symbol at 1200 bps).
     * @param input_data The input BitStream to be grouped into symbols.
     * @return A vector of grouped symbols.
     */
    std::vector<uint8_t> groupSymbols(bitstream::fixed_bit_reader& input_data) {
        std::vector<uint8_t> grouped_data;
        size_t bits_per_symbol = (baud_rate == 2400) ? 3 : (baud_rate == 1200 || (baud_rate == 75 && !is_frequency_hopping)) ? 2 : 1;
        
        while (input_data.get_remaining_bits() >= bits_per_symbol) {
            uint32_t symbol = 0;

            input_data.serialize_bits(symbol, bits_per_symbol);

            grouped_data.push_back(static_cast<uint8_t>(symbol));
        }

        return grouped_data;
    }

    /**
     * @brief Interleaves a chunk of the input BitStream and returns the result.
     * @param input_data The input BitStream chunk.
     * @return A BitStream representing the interleaved chunk.
     */
    void interleaveChunk(bitstream::growing_bit_writer<std::vector<uint8_t>>& interleaved_writer, bitstream::fixed_bit_reader& input_data) {
        loadChunk(input_data);
        return fetchChunk(interleaved_writer);
    }

    /**
     * @brief Loads bits from the input BitStream into the interleaver matrix.
     * @param data The input BitStream to load.
     */
    void loadChunk(bitstream::fixed_bit_reader& data) {
        size_t row = 0;
        size_t col = 0;
        size_t index = 0;

        size_t row_increment = (baud_rate == 75 && interleave_setting == 2) ? 7 : 9;

        // Load bits into the matrix
        std::fill(matrix.begin(), matrix.end(), 0);  // Clear previous data

        while (data.get_remaining_bits() > 0 && col < columns) {
            size_t matrix_idx = row * columns + col;
            if (matrix_idx >= matrix.size()) {
                throw std::out_of_range("Matrix index out of bounds in loadChunk()");
            }

            uint32_t bit = 0;
            if (!data.serialize_bits(bit, 1)) {
                throw std::runtime_error("Failed to read bit from chunk_reader in loadChunk()");
            }
            matrix[matrix_idx] = static_cast<uint8_t>(bit);

            row = (row + row_increment) % rows;

            if (row == 0) {
                col++;
            }
        }
    }

    /**
     * @brief Fetches bits from the interleaver matrix in the interleaved order.
     * @return A BitStream containing the fetched interleaved data.
     */
    void fetchChunk(bitstream::growing_bit_writer<std::vector<uint8_t>>& interleaved_writer) {
        size_t row = 0;
        size_t col = 0;

        size_t column_decrement = (baud_rate == 75 && interleave_setting == 2) ? 7 : 17;

        for (size_t i = 0; i < rows * columns; i++) {
            size_t matrix_idx = row * columns + col;
            if (matrix_idx >= matrix.size()) {
                throw std::out_of_range("Matrix index out of bounds in fetchChunk()");
            }

            uint32_t bit = static_cast<uint32_t>(matrix[matrix_idx]);
            if (!interleaved_writer.serialize_bits(bit, 1)) {
                throw std::runtime_error("Failed to write bit to interleaved_writer in fetchChunk()");
            }

            row++;

            if (row == rows) {
                row = 0;
                col = (col + 1) % columns;
            } else {
                col = (col + columns - column_decrement) % columns;
            }
        }
    }



    /**
     * @brief Sets the matrix dimensions based on baud rate and interleave setting.
     */
    void setMatrixDimensions() {
        if (baud_rate == 2400) {
            rows = 40;
            columns = (interleave_setting == 2) ? 576 : 72;
        } else if (baud_rate == 1200) {
            rows = 40;
            columns = (interleave_setting == 2) ? 288 : 36;
        } else if (baud_rate == 600) {
            rows = 40;
            columns = (interleave_setting == 2) ? 144 : 18;
        } else if (baud_rate == 300 || baud_rate == 150) {
            rows = 40;
            columns = (interleave_setting == 2) ? 144 : 18;
        } else if (baud_rate == 75) {
            if (is_frequency_hopping) {
                rows = 40;
                columns = 18;
            } else {
                rows = (interleave_setting == 2) ? 20 : 10;
                columns = (interleave_setting == 2) ? 36 : 9;
            }
        } else {
            throw std::invalid_argument("Invalid baud rate for setMatrixDimensions");
        }
    }

    /**
     * @brief Applies puncturing to the interleaved data (used for 2400 bps in frequency-hopping mode).
     * @param interleaved_data The interleaved data to be punctured.
     * @return A BitStream containing punctured data.
     */
    void applyPuncturing(bitstream::growing_bit_writer<std::vector<uint8_t>>& punctured_data, bitstream::fixed_bit_reader& interleaved_data) {
        size_t bit_index = 0;

        while (interleaved_data.get_remaining_bits() > 0) {
            uint32_t bit = 0;
            interleaved_data.serialize_bits(bit, 1);

            if ((bit_index % 4) != 3) {
                punctured_data.serialize_bits(bit, 1);
            }
            bit_index++;
        }
    }
};

#endif