#ifndef _COMPRESSION_ADAPTIVE_ARITHMETIC_COMPRESSION_INTEGER_H_
#define _COMPRESSION_ADAPTIVE_ARITHMETIC_COMPRESSION_INTEGER_H_

/**
 * Implementation based on https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
 */

#include "ArithmeticModel.h"

#include <alib/vector>

#include <string/LinearString.h>

namespace stringology {

namespace compression {

template < class SymbolType, class Model >
class ArithmeticCoder {
	Model m_model;

	unsigned m_pending_bits = 0;
	unsigned long long m_low = 0;
	unsigned long long m_high = max_code;

	static const unsigned valid_bits = sizeof ( unsigned long long ) * 8 / 2;
	static const unsigned long long max_code = ~0ull >> valid_bits;
	static const unsigned long long one_half = ( max_code >> 1 ) + 1;
	static const unsigned long long one_fourth = ( max_code >> 2 ) + 1;
	static const unsigned long long three_fourths = one_half + one_fourth;

	template < class Callback >
	inline static void put_bit_plus_pending ( Callback && callback, bool bit, unsigned & m_pending_bits) {
		callback ( bit );
		while ( m_pending_bits > 0 ) {
			callback ( ! bit );
			-- m_pending_bits;
		}
	}

	template < class Callback >
	void encode ( unsigned prob_low, unsigned prob_high, unsigned prob_count, Callback callback ) {
		unsigned long long range = m_high - m_low + 1;
		m_high = m_low + range * prob_high / prob_count - 1;
		m_low = m_low + range * prob_low / prob_count;
		for ( ; ; ) {
			if ( m_high < one_half || m_low >= one_half ) {
				put_bit_plus_pending ( std::forward < Callback && > ( callback ), m_low >= one_half, m_pending_bits );
			} else if ( m_low >= one_fourth && m_high < three_fourths ) {
				m_pending_bits++;
				m_low -= one_fourth;
				m_high -= one_fourth;
			} else
				break;
			m_high <<= 1;
			m_high++;
			m_low <<= 1;

			m_low &= max_code;
			m_high &= max_code;
		}
	}

public:
	ArithmeticCoder ( Model model ) : m_model ( std::move ( model ) ) {
	}

	template < class Callback >
	void encode ( const SymbolType & symbol, Callback && callback ) {
		unsigned prob_low;
		unsigned prob_high;
		unsigned prob_count = m_model.getCount ( );

		m_model.getProbability ( symbol, prob_low, prob_high );
		m_model.update ( symbol );

		encode ( prob_low, prob_high, prob_count, std::forward < Callback && > ( callback ) );
	}

	template < class Callback >
	void finalize ( Callback && callback ) {
		unsigned prob_low;
		unsigned prob_high;
		unsigned prob_count = m_model.getCount ( );

		m_model.getProbabilityEof ( prob_low, prob_high );

		encode ( prob_low, prob_high, prob_count, std::forward < Callback && > ( callback ) );

		m_pending_bits++;
		put_bit_plus_pending ( std::forward < Callback && > ( callback ), m_low >= one_fourth, m_pending_bits);
	}

	Model & getModel ( ) {
		return m_model;
	}
};

class AdaptiveIntegerArithmeticCompression {
public:
	template < class SymbolType >
	static ext::vector < bool > compress ( const string::LinearString < SymbolType > & source ) {
		ext::vector < bool > result;
		ArithmeticCoder < SymbolType, ArithmeticModel < SymbolType > > arithmeticCoder ( ArithmeticModel < SymbolType > ( source.getAlphabet ( ) ) );

		auto output = [&] ( bool bit ) {
			result.push_back ( bit );
		};

		for ( size_t index = 0; index < source.getContent ( ).size ( ); ++ index ) {
			arithmeticCoder.encode ( source.getContent ( ) [ index ], output );
		}
		arithmeticCoder.finalize ( output );

		return result;
	}

};

} /* namespace compression */

} /* namespace stringology */

#endif /* _COMPRESSION_ADAPTIVE_ARITHMETIC_COMPRESSION_INTEGER_H_ */
