#include "aricoder.h"

#include "bitops.h"

#include <algorithm>
#include <cstdlib>
#include <functional>
#include <limits>

template <std::uint8_t bit>
void ArithmeticBitWriter::write_bit()
{
    // add bit at last position
    curr_byte_ = (curr_byte_ << 1) | bit;
    // increment bit position
    curr_bit_++;

    // write bit if done
    if (curr_bit_ == 8)
    {
        data_.emplace_back(curr_byte_);
        curr_bit_ = 0;
    }
}

void ArithmeticBitWriter::write_n_zero_bits(std::size_t n)
{
    if (n + curr_bit_ >= 8)
    {
        auto remainingBits = 8 - curr_bit_;
        n -= remainingBits;
        curr_byte_ <<= remainingBits;
        data_.emplace_back(curr_byte_);
        curr_bit_ = 0;
    }

    while (n >= 8)
    {
        data_.emplace_back(0);
        n -= 8;
    }

    curr_byte_ <<= n;
    curr_bit_ += n;
}

void ArithmeticBitWriter::write_n_one_bits(std::size_t n)
{
    constexpr std::uint8_t all_ones = std::numeric_limits<std::uint8_t>::max();
    if (n + curr_bit_ >= 8)
    {
        auto remainingBits = 8 - curr_bit_;
        n -= remainingBits;
        curr_byte_ <<= remainingBits;
        curr_byte_ |= all_ones >> (8 - remainingBits);
        data_.emplace_back(curr_byte_);
        curr_bit_ = 0;
    }

    while (n >= 8)
    {
        data_.emplace_back(all_ones);
        n -= 8;
    }

    curr_byte_ = (curr_byte_ << n) | (all_ones >> (8 - n));
    curr_bit_ += n;
}

void ArithmeticBitWriter::pad()
{
    while (curr_bit_ > 0)
    {
        write_bit<0>();
    }
}

std::vector<std::uint8_t> ArithmeticBitWriter::get_data() const
{
    return data_;
}

ArithmeticDecoder::ArithmeticDecoder(Reader& reader) : reader_(reader)
{
    // code buffer has to be filled before starting decoding
    for (std::uint32_t i = 0; i < CODER_USE_BITS; i++)
    {
        ccode = (ccode << 1) | read_bit();
    }
}

ArithmeticEncoder::ArithmeticEncoder(Writer& writer) : writer_(writer) {}

ArithmeticEncoder::~ArithmeticEncoder()
{
    if (!finalized)
    {
        this->finalize();
    }
}

void ArithmeticEncoder::finalize()
{
    if (finalized)
    {
        return;
    }

    // due to clow < CODER_LIMIT050, and chigh >= CODER_LIMIT050
    // there are only two possible cases
    if (clow < CODER_LIMIT025)
    {
        bitwriter_->write_bit<0>();
        bitwriter_->write_bit<1>();
        bitwriter_->write_n_one_bits(nrbits);
        nrbits = 0;
    }
    else
    {
        // case b.), clow >= CODER_LIMIT025
        bitwriter_->write_bit<1>();
    }
    // done, zeroes are auto-read by the decoder

    bitwriter_->pad(); // Pad code with zeroes.
    writer_.write(bitwriter_->get_data());

    finalized = true;
}

/* -----------------------------------------------
    arithmetic encoder function
    ----------------------------------------------- */

void ArithmeticEncoder::encode(symbol* s)
{
    // Make local copies of clow_ and chigh_ for cache performance:
    uint32_t clow_local = clow;
    uint32_t chigh_local = chigh;
    // update steps, low count, high count
    cstep = (chigh_local - clow_local + 1) / s->scale;
    chigh_local = clow_local + (cstep * s->high_count) - 1;
    clow_local = clow_local + (cstep * s->low_count);

    // e3 scaling is performed for speed and to avoid underflows
    // if both, low and high are either in the lower half or in the higher half
    // one bit can be safely shifted out
    while (clow_local >= CODER_LIMIT050  || chigh_local < CODER_LIMIT050)
    {
        if (chigh_local < CODER_LIMIT050)       // this means both, high and low are below, and 0 can be safely shifted out
        {
            // write 0 bit
            bitwriter_->write_bit<0>();
            // shift out remaing e3 bits
            bitwriter_->write_n_one_bits(nrbits);
            nrbits = 0;
        }
        else   // if the first wasn't the case, it's clow >= CODER_LIMIT050
        {
            // write 1 bit
            bitwriter_->write_bit<1>();
            clow_local &= CODER_LIMIT050 - 1;
            chigh_local &= CODER_LIMIT050 - 1;
            // shift out remaing e3 bits
            bitwriter_->write_n_zero_bits(nrbits);
            nrbits = 0;
        }
        clow_local <<= 1;
        chigh_local <<= 1;
        chigh_local++;
    }

    // e3 scaling, to make sure that theres enough space between low and high
    while ((clow_local >= CODER_LIMIT025) && (chigh_local < CODER_LIMIT075))
    {
        nrbits++;
        clow_local &= CODER_LIMIT025 - 1;
        chigh_local ^= CODER_LIMIT025 + CODER_LIMIT050;
        // clow  -= CODER_LIMIT025;
        // chigh -= CODER_LIMIT025;
        clow_local <<= 1;
        chigh_local <<= 1;
        chigh_local++;
    }

    clow = clow_local;
    chigh = chigh_local;
}


/* -----------------------------------------------
    arithmetic decoder get count function
    ----------------------------------------------- */

unsigned int ArithmeticDecoder::decode_count(symbol* s)
{
    // update cstep, which is needed to remove the symbol from the stream later
    cstep = ((chigh - clow) + 1) / s->scale;

    // return counts, needed to decode the symbol from the statistical model
    return (ccode - clow) / cstep;
}

/* -----------------------------------------------
    arithmetic decoder function
    ----------------------------------------------- */

void ArithmeticDecoder::decode(symbol* s)
{
    // no actual decoding takes place, as this has to happen in the statistical model
    // the symbol has to be removed from the stream, though

    // alread have steps updated from decoder_count
    // update low count and high count
    uint32_t ccode_local = ccode;
    uint32_t clow_local = clow;
    uint32_t chigh_local = clow_local + (cstep * s->high_count) - 1;
    clow_local = clow_local + (cstep * s->low_count);

    // e3 scaling is performed for speed and to avoid underflows
    // if both, low and high are either in the lower half or in the higher half
    // one bit can be safely shifted out
    while ((clow_local >= CODER_LIMIT050) || (chigh_local < CODER_LIMIT050))
    {
        if (clow_local >= CODER_LIMIT050)
        {
            clow_local &= CODER_LIMIT050 - 1;
            chigh_local &= CODER_LIMIT050 - 1;
            ccode_local &= CODER_LIMIT050 - 1;
        } // if the first wasn't the case, it's chigh < CODER_LIMIT050
        clow_local <<= 1;
        chigh_local <<= 1;
        chigh_local++;
        ccode_local <<= 1;
        ccode_local |= read_bit();
    }

    // e3 scaling, to make sure that theres enough space between low and high
    while ((clow_local >= CODER_LIMIT025) && (chigh_local < CODER_LIMIT075))
    {
        clow_local &= CODER_LIMIT025 - 1;
        chigh_local ^= CODER_LIMIT025 + CODER_LIMIT050;
        // clow  -= CODER_LIMIT025;
        // chigh -= CODER_LIMIT025;
        ccode_local -= CODER_LIMIT025;
        clow_local <<= 1;
        chigh_local <<= 1;
        chigh_local++;
        ccode_local <<= 1;
        ccode_local |= read_bit();
    }
    chigh = chigh_local;
    clow = clow_local;
    ccode = ccode_local;
}

/* -----------------------------------------------
    bit reader function
    ----------------------------------------------- */

unsigned char ArithmeticDecoder::read_bit()
{
    // read in new byte if needed
    if (cbit == 0)
    {
        if (!reader_.read_byte(&bbyte))  // read next byte if available
        {
            bbyte = 0;    // if no more data is left in the stream
        }
        cbit = 8;
    }

    // decrement current bit position
    cbit--;
    // return bit at cbit position
    return BITN(bbyte, cbit);
}


/* -----------------------------------------------
    universal statistical model for arithmetic coding

    boundaries of this model:
    max_s (maximum symbol) -> 1 <= max_s <= 1024 (???)
    max_c (maximum context) -> 1 <= max_c <= 1024 (???)
    max_o (maximum order) -> -1 <= max_o <= 4
    c_lim (maximum count) -> 2 <= c_lim <= 4096 (???)
    WARNING: this can be memory intensive, so don't overdo it
    max_s == 256; max_c == 256; max_o == 4 would be way too much
    ----------------------------------------------- */

model_s::model_s(int max_s, int max_c, int max_o, int c_lim) :
    // Copy settings into the model:
    max_symbol(max_s),
    max_context(max_c),
    max_order(max_o + 1),
    max_count(c_lim),

    current_order(max_o + 1),
    sb0_count(max_s),

    totals(max_s + 2),
    scoreboard(new bool[max_s]),
    contexts(max_o + 3)
{
    std::fill(scoreboard, scoreboard + max_symbol, false);

    // set up null table
    table_s* null_table = new table_s;
    null_table->counts = std::vector<uint16_t>(max_symbol, uint16_t(1));  // Set all probabilities to 1.

    // set up internal counts
    null_table->max_count = 1;
    null_table->max_symbol = max_symbol;

    // set up start table
    table_s* start_table = new table_s;
    start_table->links = std::vector<table_s*>(max_context);

    // integrate tables into contexts
    contexts[ 0 ] = null_table;
    contexts[ 1 ] = start_table;

    // build initial 'normal' tables
    for (int i = 2; i <= max_order; i++)
    {
        // set up current order table
        contexts[i] = new table_s;
        // build forward links
        if (i < max_order)
        {
            contexts[i]->links = std::vector<table_s*>(max_context);
        }
        contexts[ i - 1 ]->links[ 0 ] = contexts[ i ];
    }
}


/* -----------------------------------------------
    model class destructor - recursive cleanup of memory is done here
    ----------------------------------------------- */

model_s::~model_s()
{
    // clean up each 'normal' table
    delete contexts[1];

    // clean up null table
    delete contexts[0];

    // free everything else
    delete[] scoreboard;
}


/* -----------------------------------------------
    Updates statistics for a specific symbol / resets to highest order.
    Use -1 if you just want to reset without updating statistics.
    ----------------------------------------------- */
void model_s::update_model(int symbol)
{
    // only contexts, that were actually used to encode
    // the symbol get its count updated
    if (symbol >= 0)
    {
        for (int local_order = (current_order < 1) ? 1 : current_order;
                local_order <= max_order; local_order++)
        {
            table_s* context = contexts[ local_order ];
            auto& count = context->counts[symbol];
            // update count for specific symbol & scale
            count++;
            // store side information for totalize_table
            context->max_count = std::max(count, context->max_count);
            context->max_symbol = std::max(uint16_t(symbol + 1), context->max_symbol);
            // if count for that symbol have gone above the maximum count
            // the table has to be resized (scale factor 2)
            if (count == max_count)
            {
                context->rescale_table();
            }
        }
    }

    // reset scoreboard and current order
    current_order = max_order;
    std::fill(scoreboard, scoreboard + max_symbol, false);
    sb0_count = max_symbol;
}


/* -----------------------------------------------
    shift in one context (max no of contexts is max_c)
    ----------------------------------------------- */

void model_s::shift_context(int c)
{
    // shifting is not possible if max_order is below 1
    // or context index is negative
    if ((max_order < 2) || (c < 0))
    {
        return;
    }

    // shift each orders' context
    for (int i = max_order; i > 1; i--)
    {
        // this is the new current order context
        table_s* context = contexts[ i - 1 ]->links[ c ];

        // check if context exists, build if needed
        if (context == nullptr)
        {
            // reserve memory for next table_s
            context = new table_s;
            // finished here if this is a max order context
            if (i < max_order)
            {
                // build links to higher order tables otherwise
                context->links.resize(max_context);
            }
            // put context to its right place
            contexts[ i - 1 ]->links[ c ] = context;
        }

        // switch context
        contexts[ i ] = context;
    }
}


/* -----------------------------------------------
    Flushes the entire model by calling rescale_table on all contexts.
    ----------------------------------------------- */

void model_s::flush_model()
{
    contexts[1]->recursive_flush();
}


/* -----------------------------------------------
    Excludes every symbol above c.
    ----------------------------------------------- */

void model_s::exclude_symbols(int c)
{
    // exclusions are back to normal after update_model is used

    for (c = c + 1; c < max_symbol; c++)
    {
        if (!scoreboard[ c ])
        {
            scoreboard[ c ] = true;
            sb0_count--;
        }
    }
}


/* -----------------------------------------------
    converts an int to a symbol, needed only when encoding
    ----------------------------------------------- */

int model_s::convert_int_to_symbol(int c, symbol* s)
{
    // search the symbol c in the current context table_s,
    // return scale, low- and high counts

    // totalize table for the current context
    table_s* context = contexts[ current_order ];
    totalize_table(context);

    // finding the scale is easy
    s->scale = totals[ 0 ];

    // check if that symbol exists in the current table. send escape otherwise
    if (context->counts[ c ] > 0)
    {
        // return high and low count for the current symbol
        s->low_count  = totals[ c + 2 ];
        s->high_count = totals[ c + 1 ];
        return 0;
    }

    // return high and low count for the escape symbol
    s->low_count  = totals[ 1 ];
    s->high_count = totals[ 0 ];
    current_order--;
    return 1;
}


/* -----------------------------------------------
    returns the current context scale needed only when decoding
    ----------------------------------------------- */

void model_s::get_symbol_scale(symbol* s)
{
    // getting the scale is easy: totalize the table_s, use accumulated count -> done
    totalize_table(contexts[ current_order ]);
    s->scale = totals[ 0 ];
}


/* -----------------------------------------------
    converts a count to an int, called after get_symbol_scale
    ----------------------------------------------- */

int model_s::convert_symbol_to_int(uint32_t count, symbol* s)
{
    // seek the symbol that matches the count,
    // also, set low- and high count for the symbol - it has to be removed from the stream


    // go through the totals table, search the symbol that matches the count
    int c;
    for (c = 1; c < int(totals.size()); c++)
    {
        if (count >= totals[c])
        {
            break;
        }
    }
    // set up the current symbol
    s->low_count = totals[c]; // It is guaranteed that there exists such a symbol.
    s->high_count = totals[c - 1]; // This is guaranteed to not go out of bounds since the search started at index 1 of totals.
    // send escape if escape symbol encountered
    if (c == 1)
    {
        current_order--;
        return ESCAPE_SYMBOL;
    }

    // return symbol value
    return c - 2 ; // Since c is not one and is a positive number, this will be nonnegative.
}


/* -----------------------------------------------
    totals are calculated by accumulating counts in the current table_s
    ----------------------------------------------- */

void model_s::totalize_table(table_s* context)
{
    // update exclusion is used, so this has to be done each time
    // escape probability calculation also takes place here

    // accumulated counts must never exceed CODER_MAXSCALE
    // as CODER_MAXSCALE is big enough, though, (2^29), this shouldn't happen and is not checked

    const auto& counts = context->counts;

    // check counts
    if (!counts.empty())    // if counts are already set
    {
        // locally store current fill/symbol count
        int local_symb = sb0_count;

        // set the last symbol of the totals to zero
        int i = context->max_symbol - 1;
        totals[i + 2] = 0;

        // (re)set current total
        uint32_t curr_total = 0;

        // go reverse though the whole counts table and accumulate counts
        // leave space at the beginning of the table for the escape symbol
        for (; i >= 0; i--)
        {
            // only count probability if the current symbol is not 'scoreboard - excluded'
            if (!scoreboard[i])
            {
                uint16_t curr_count = counts[i];
                if (curr_count > 0)
                {
                    // add counts for the current symbol
                    curr_total += curr_count;
                    // exclude symbol from scoreboard
                    scoreboard[i] = true;
                    sb0_count--;
                }
            }
            totals[i + 1] = curr_total;
        }
        // here the escape calculation needs to take place
        uint32_t esc_prob;
        if (local_symb == sb0_count)
        {
            esc_prob = 1;
        }
        else if (sb0_count == 0)
        {
            esc_prob = 0;
        }
        else
        {
            // esc_prob = 1;
            esc_prob  =  sb0_count * (local_symb - sb0_count);
            esc_prob /= (local_symb * context->max_count);
            esc_prob++;
        }
        // include escape probability in totals table
        totals[ 0 ] = totals[ 1 ] + esc_prob;
    }
    else     // if counts are not already set
    {
        // setup counts for current table
        context->counts.resize(max_symbol);
        // set totals table -> only escape probability included
        totals[ 0 ] = 1;
        totals[ 1 ] = 0;
    }
}

/* -----------------------------------------------
    special version of model_s for binary coding

    boundaries of this model:
    ... (maximum symbol) -> 2 (0 or 1 )
    max_c (maximum context) -> 1 <= max_c <= 1024 (???)
    max_o (maximum order) -> -1 <= max_o <= 4
    ----------------------------------------------- */

model_b::model_b(int max_c, int max_o, int c_lim) :
    // Copy settings into the model:
    max_context(max_c),
    max_order(max_o + 1),
    max_count(c_lim),

    contexts(max_o + 3)
{
    // set up null table
    table* null_table = new table;
    null_table->counts = std::vector<uint16_t>(2, uint16_t(1));
    null_table->scale = uint32_t(2);

    // set up start table
    table* start_table = new table;
    start_table->links = std::vector<table*>(max_context);

    // integrate tables into contexts
    contexts[ 0 ] = null_table;
    contexts[ 1 ] = start_table;

    // build initial 'normal' tables
    for (int i = 2; i <= max_order; i++)
    {
        // set up current order table
        contexts[i] = new table;
        // build forward links
        if (i < max_order)
        {
            contexts[i]->links = std::vector<table*>(max_context);
        }
        contexts[ i - 1 ]->links[ 0 ] = contexts[ i ];
    }
}


/* -----------------------------------------------
    model class destructor - recursive cleanup of memory is done here
    ----------------------------------------------- */

model_b::~model_b()
{
    // clean up each 'normal' table
    delete contexts[1];

    // clean up null table
    delete contexts[0];
}


/* -----------------------------------------------
    updates statistics for a specific symbol / resets to highest order
    ----------------------------------------------- */

void model_b::update_model(int symbol)
{
    // use -1 if you just want to reset without updating statistics

    table* context = contexts[ max_order ];

    // only contexts, that were actually used to encode
    // the symbol get their counts updated
    if ((symbol >= 0) && (max_order >= 0))
    {
        // update count for specific symbol & scale
        context->counts[ symbol ]++;
        context->scale++;
        // if counts for that symbol have gone above the maximum count
        // the table has to be resized (scale factor 2)
        if (context->counts[ symbol ] >= max_count)
        {
            context->rescale_table();
        }
    }
}


/* -----------------------------------------------
    shift in one context (max no of contexts is max_c)
    ----------------------------------------------- */

void model_b::shift_context(int c)
{
    // shifting is not possible if max_order is below 1
    // or context index is negative
    if ((max_order < 2) || (c < 0))
    {
        return;
    }

    // shift each orders' context
    for (int i = max_order; i > 1; i--)
    {
        // this is the new current order context
        table* context = contexts[ i - 1 ]->links[ c ];

        // check if context exists, build if needed
        if (context == nullptr)
        {
            // reserve memory for next table
            context = new table;
            // finished here if this is a max order context
            if (i < max_order)
            {
                // build links to higher order tables otherwise
                context->links.resize(max_context);
            }
            // put context to its right place
            contexts[ i - 1 ]->links[ c ] = context;
        }

        // switch context
        contexts[ i ] = context;
    }
}


/* -----------------------------------------------
    Flushes the entire model by calling rescale_table on all contexts.
    ----------------------------------------------- */

void model_b::flush_model()
{
    contexts[1]->recursive_flush();
}


/* -----------------------------------------------
    converts an int to a symbol, needed only when encoding
    ----------------------------------------------- */

int model_b::convert_int_to_symbol(int c, symbol* s)
{
    table* context = contexts[ max_order ];

    // check if counts are available
    context->check_counts();

    // finding the scale is easy
    s->scale = context->scale;

    // return high and low count for current symbol
    if (c == 0)     // if 0 is to be encoded
    {
        s->low_count  = uint32_t(0);
        s->high_count = context->counts[ 0 ];
    }
    else   // if 1 is to be encoded
    {
        s->low_count  = context->counts[ 0 ];
        s->high_count = context->scale;
    }

    return 1;
}


/* -----------------------------------------------
    returns the current context scale needed only when decoding
    ----------------------------------------------- */

void model_b::get_symbol_scale(symbol* s)
{
    table* context = contexts[ max_order ];

    // check if counts are available
    context->check_counts();

    // getting the scale is easy
    s->scale = context->scale;
}


/* -----------------------------------------------
    converts a count to an int, called after get_symbol_scale
    ----------------------------------------------- */

int model_b::convert_symbol_to_int(uint32_t count, symbol* s)
{
    table* context = contexts[ max_order ];
    auto counts0 = context->counts[ 0 ];

    // set up the current symbol
    if (count < counts0)
    {
        s->low_count  = uint32_t(0);
        s->high_count = counts0;
        return 0;
    }
    else
    {
        s->low_count  = counts0;
        s->high_count = s->scale;
        return 1;
    }
}