#ifndef _COMPRESSION_ADAPTIVE_ARITHMETIC_DECOMPRESSION_INTEGER_H_
#define _COMPRESSION_ADAPTIVE_ARITHMETIC_DECOMPRESSION_INTEGER_H_

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

#include "ArithmeticModel.h"

#include <string/LinearString.h>

namespace stringology {

namespace compression {

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

	unsigned long long m_value = 0;
	unsigned long long m_low = 0;
	unsigned long long m_high = 0;

	unsigned long long m_range = 0;
	unsigned m_prob_count = 0;
	unsigned m_scaled_value = 0;

	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 >
	void fillVariables ( Callback && callback ) {
		for( ; ; ) {
			if ( m_high < one_half || m_low >= one_half ) {
				//do nothing, both bits are a zero or both bits are one
			} else if ( m_low >= one_fourth && m_high < three_fourths ) {
				m_value -= one_fourth;
				m_low -= one_fourth;
				m_high -= one_fourth;
			} else
				break;
			m_low <<= 1;
			m_high <<= 1;
			m_value <<= 1;

			m_high++;
			m_value += callback ( );

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

		m_range = m_high - m_low + 1;
		m_prob_count = m_model.getCount ( );
		m_scaled_value = ( ( m_value - m_low + 1 ) * m_prob_count - 1 ) / m_range;
	}

public:
	template < class Callback >
	ArithmeticDecoder ( Model model, Callback && callback ) : m_model ( std::move ( model ) ) {
		fillVariables ( std::forward < Callback && > ( callback ) );
	}

	bool isEof ( ) {
		return m_model.isEof ( m_scaled_value );
	}

	template < class Callback >
	SymbolType decode ( Callback && callback ) {
		unsigned prob_low;
		unsigned prob_high;
		SymbolType c = m_model.getChar ( m_scaled_value, prob_low, prob_high );
		m_model.update ( c );

		m_high = m_low + m_range * prob_high / m_prob_count - 1;
		m_low = m_low + m_range * prob_low / m_prob_count;

		fillVariables ( std::forward < Callback && > ( callback ) );

		return c;
	}

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

class AdaptiveIntegerArithmeticDecompression {
public:
	template < class SymbolType >
	static string::LinearString < SymbolType > decompress ( const ext::vector < bool > & source, const ext::set < SymbolType > & alphabet ) {
		size_t index = 0;
		auto input = [ & ] ( ) {
			return index >= source.size ( ) ? 0 : source [ index ++ ] ? 1 : 0;
		};

		ArithmeticDecoder < SymbolType, ArithmeticModel < SymbolType > > decoder ( ArithmeticModel < SymbolType > ( alphabet ), input );
		ext::vector < SymbolType > result;

		while ( ! decoder.isEof ( ) ) {
			result.push_back ( decoder.decode ( input ) );
		}

		return string::LinearString < SymbolType > ( alphabet, result );
	}
};

} /* namespace compression */

} /* namespace stringology */

#endif /* _COMPRESSION_ADAPTIVE_ARITHMETIC_DECOMPRESSION_INTEGER_H_ */
