#ifndef ARICODER_H
#define ARICODER_H

#include "bitops.h"

#include <algorithm>
#include <cstdint>
#include <memory>
#include <vector>

// defines for coder
constexpr uint32_t CODER_USE_BITS = 31; // Must never be above 31.
constexpr uint32_t CODER_LIMIT100 = uint32_t(1 << CODER_USE_BITS);
constexpr uint32_t CODER_LIMIT025 = CODER_LIMIT100 / 4;
constexpr uint32_t CODER_LIMIT050 = (CODER_LIMIT100 / 4) * 2;
constexpr uint32_t CODER_LIMIT075 = (CODER_LIMIT100 / 4) * 3;
constexpr uint32_t CODER_MAXSCALE = CODER_LIMIT025 - 1;
constexpr uint32_t ESCAPE_SYMBOL = CODER_LIMIT025;

// symbol struct, used in arithmetic coding
struct symbol {
	uint32_t low_count;
	uint32_t high_count;
	uint32_t scale;
};

// table struct, used in in statistical models,
// holding all info needed for one context
struct table {
	// counts for each symbol contained in the table
	std::vector<uint16_t> counts;
	// links to higher order contexts
	std::vector<table*> links;
	// accumulated counts
	uint32_t scale = uint32_t(0);

	/* -----------------------------------------------
	Recursively deletes all the tables pointed to in links.
	----------------------------------------------- */
	~table() {
		for (auto& link : links) {
			if (link != nullptr) {
				delete link;
			}
		}
	}

	/* -----------------------------------------------
	Checks if counts exist, creating it if it does not.
	----------------------------------------------- */
	inline void check_counts() {
		// check if counts are available
		if (counts.empty()) {
			// setup counts for current table
			counts.resize(2, uint16_t(1));
			// set scale
			scale = uint32_t(2);
		}
	}

	/* -----------------------------------------------
	Resizes the table by rightshifting each count by 1.
	----------------------------------------------- */
	inline void rescale_table() {
		// Do nothing if counts is not set:
		if (!counts.empty()) {
			// Scale the table by bitshifting each count, be careful not to set any count zero:
			counts[0] = std::max(uint16_t(1), uint16_t(counts[0] >> 1));
			counts[1] = std::max(uint16_t(1), uint16_t(counts[1] >> 1));
			scale = counts[0] + counts[1];
		}
	}

	/* -----------------------------------------------
	Recursively runs rescale_table on this and all linked contexts.
	----------------------------------------------- */
	inline void recursive_flush() {
		for (auto& link : links) {
			if (link != nullptr) {
				link->recursive_flush();
			}
		}
		// rescale specific table
		rescale_table();
	}
};

// special table struct, used in in model_s,
// holding additional info for a speedier 'totalize_table'
struct table_s {
	// counts for each symbol contained in the table
	std::vector<uint16_t> counts;
	// links to higher order contexts
	std::vector<table_s*> links;
	// speedup info
	uint16_t max_count = uint16_t(0);
	uint16_t max_symbol = uint16_t(0);

	/* -----------------------------------------------
	Recursively deletes all the tables pointed to in links.
	----------------------------------------------- */
	~table_s() {
		for (auto& link : links) {
			if (link != nullptr) {
				delete link;
			}
		}
	}

	/* -----------------------------------------------
	Resizes the table by rightshifting each count by 1.
	----------------------------------------------- */
	inline void rescale_table() {
		// Nothing to do if counts has not been set.
		if (counts.empty()) return;

		// now scale the table by bitshifting each count
		int lst_symbol = max_symbol;
		int i;
		for (i = 0; i < lst_symbol; i++) {
			counts[i] >>= 1; // Counts will not become negative since it is an unsigned type.
		}

		// also rescale tables max count
		max_count >>= 1;

		// seek for new last symbol
		for (i = lst_symbol - 1; i >= 0; i--) {
			if (counts[i] > 0) {
				break;
			}
		}
		max_symbol = i + 1;
	}

	/* -----------------------------------------------
	Recursively runs rescale_table on this and all linked contexts.
	----------------------------------------------- */
	inline void recursive_flush() {
		for (auto& link : links) {
			if (link != nullptr) {
				link->recursive_flush();
			}
		}

		// rescale specific table
		rescale_table();
	}
};

class ArithmeticBitWriter {
public:
	template <std::uint8_t bit>
	void write_bit();

	void write_n_zero_bits(std::size_t n);

	void write_n_one_bits(std::size_t n);

	void pad();

	std::vector<std::uint8_t> get_data() const;


private:
	std::vector<std::uint8_t> data_;
	std::uint8_t curr_byte_ = 0;
	std::size_t curr_bit_ = 0;
};


/* -----------------------------------------------
	class for arithmetic coding of data to/from iostream
	----------------------------------------------- */

class ArithmeticEncoder
{
	public:
    ArithmeticEncoder(Writer& writer);
	~ArithmeticEncoder();
	void encode( symbol* s );

    void finalize();
	
	private:
	
	// i/o variables
    
    bool finalized = false;
    Writer& writer_;
    std::unique_ptr<ArithmeticBitWriter> bitwriter_ = std::make_unique<ArithmeticBitWriter>();
	
	// arithmetic coding variables
	unsigned int ccode = 0;
	unsigned int clow = 0;
	unsigned int chigh = CODER_LIMIT100 - 1;
	unsigned int cstep = 0;
	unsigned int nrbits = 0;
};

class ArithmeticDecoder {
	public:
    ArithmeticDecoder(Reader& reader);
	~ArithmeticDecoder() {}
	unsigned int decode_count( symbol* s );
	void decode( symbol* s );
	
	private:
	unsigned char read_bit();
	
	// i/o variables
    Reader& reader_;
	unsigned char bbyte = 0;
	unsigned char cbit = 0;
	
	// arithmetic coding variables
	unsigned int ccode = 0;
	unsigned int clow = 0;
	unsigned int chigh = CODER_LIMIT100 - 1;
	unsigned int cstep = 0;
};


/* -----------------------------------------------
	universal statistical model for arithmetic coding
	----------------------------------------------- */
	
class model_s
{	
	public:
	
	model_s( int max_s, int max_c, int max_o, int c_lim );
	~model_s();
	
	void update_model( int symbol );
	void shift_context( int c );
	void flush_model();
	void exclude_symbols(int c);
	
	int  convert_int_to_symbol( int c, symbol *s );
	void get_symbol_scale( symbol *s );
	int  convert_symbol_to_int(uint32_t count, symbol *s);
	
	private:

	inline void totalize_table(table_s* context);

	const int max_symbol;
	const int max_context;
	const int max_order;
	const int max_count;

	int current_order;
	int sb0_count;
	
	std::vector<uint32_t> totals;
	bool* scoreboard;
	std::vector<table_s*> contexts;
};


/* -----------------------------------------------
	binary statistical model for arithmetic coding
	----------------------------------------------- */
	
class model_b
{	
	public:
	
	model_b( int max_c, int max_o, int c_lim );
	~model_b();
	
	void update_model( int symbol );
	void shift_context( int c );
	void flush_model();
	
	int  convert_int_to_symbol( int c, symbol *s );
	void get_symbol_scale( symbol *s );
	int  convert_symbol_to_int(uint32_t count, symbol *s);	
	
	private:
	
	const int max_context;
	const int max_order;
	const int max_count;
	
	std::vector<table*> contexts;
};

// Base case for shifting an arbitrary number of contexts into the model.
template <typename M>
static void shift_model(M) {}

// Shift an arbitrary number of contexts into the model (at most max_c contexts).
template <typename M, typename C, typename... Cargs>
static void shift_model(M model, C context, Cargs ... contextList) {
	model->shift_context(context);
	shift_model(model, contextList...);
}

/* -----------------------------------------------
	generic model_s encoder function
	----------------------------------------------- */
static inline void encode_ari( ArithmeticEncoder* encoder, model_s* model, int c )
{
	symbol s;
	int esc;
	
	do {		
		esc = model->convert_int_to_symbol( c, &s );
		encoder->encode( &s );
	} while ( esc );
	model->update_model( c );
}

/* -----------------------------------------------
	generic model_s decoder function
	----------------------------------------------- */	
static inline int decode_ari( ArithmeticDecoder* decoder, model_s* model )
{
	symbol s;
	uint32_t count;
	int c;
	
	do{
		model->get_symbol_scale( &s );
		count = decoder->decode_count( &s );
		c = model->convert_symbol_to_int( count, &s );
		decoder->decode( &s );	
	} while ( c == ESCAPE_SYMBOL );
	model->update_model( c );
	
	return c;
}

/* -----------------------------------------------
	generic model_b encoder function
	----------------------------------------------- */	
static inline void encode_ari( ArithmeticEncoder* encoder, model_b* model, int c )
{
	symbol s;
	
	model->convert_int_to_symbol( c, &s );
	encoder->encode( &s );
	model->update_model( c );
}

/* -----------------------------------------------
	generic model_b decoder function
	----------------------------------------------- */	
static inline int decode_ari( ArithmeticDecoder* decoder, model_b* model )
{
	symbol s;
	
	model->get_symbol_scale( &s );
	uint32_t count = decoder->decode_count( &s );
	int c = model->convert_symbol_to_int( count, &s );
	decoder->decode( &s );	
	model->update_model( c );
	
	return c;
}

#endif