#include "arith.h"
#include "files.h"


/*
	Declare variables use for arithmetic compression
*/

static int under_flow_count = 0;
static unsigned int low = 0;
static unsigned int range = MAX_RANGE;
static unsigned int value = 0;

static unsigned int code_value = 0;
static unsigned int code_bit_count = 0;
static unsigned codesize = 0;


static int shift_count (unsigned);

static void put_bit (int);
static void put_bit_string (unsigned, int);
static void clear_under_flow (unsigned);

static unsigned get_bit_string (int);
static unsigned get_bit (void);



/*
	Initialize arithmetic coder at start of any compress/expand procedure
*/

void InitCoder (void)

{
	low = 0;
	range = MAX_RANGE;
}



/*
	Terminate arithmetic code operation
	Write bit string whose value is within active code range
	Force final byte to be output to file
*/

void CloseCoder (void)

{
	int i;

	put_bit (1);
	for (i = 0; i < 7; i++) put_bit (0);
	CloseOutputFile ();
}



/*
	Save state of arithmetic coder
	Allows state to be restored at a later time
*/

void SaveCoderState (struct coder_state *p)

{
	p -> low = low;
	p -> range = range;
	p -> uflow = under_flow_count;
	p -> bits = code_bit_count;
	p -> fpos = codesize;
}


/*
	Restore arithmetic coder state from saved value
	Repositin output file
	Set all internal values to their original state
*/

void RestoreCoderState (struct coder_state *p)

{
	int n;
	int cv;
	static unsigned char reset_mask [] =

	{
		0x00,  0x80,  0xC0,  0xE0,  0xF0,  0xF8,  0xFC,  0xFE,  0xFF,
	};

	n = codesize - p -> fpos;
	cv = ResetOutputPointer (n);

	low = p -> low;
	range = p -> range;

	under_flow_count = p -> uflow;
	code_bit_count = p -> bits;
	code_value = cv & reset_mask [code_bit_count];
	codesize -= n;
}


/*
	Estimate length of output code string
	Uses previously saved coder state
	Returns difference between saved position and current position
*/

int CodeLength (struct coder_state *csptr)

{
	int len;

	len = codesize - csptr -> fpos;
	len *= 8;
	len += under_flow_count - csptr -> uflow;
	len += code_bit_count - csptr -> bits;

	return len;
}


/*
	Arithmetic coder
	Encode symbol using input frequencies
	Output as many bits as possible to output file
	Update coder values for next symbol
*/

void EncodeArith (unsigned base, unsigned freq, unsigned cmax)

{
	unsigned long t;
	unsigned x;

	int n1, n2;

	t  = (long) range * (long) base;
	t += (long) base;
	low  += (unsigned) (t / cmax);
	x     = (unsigned) (t % cmax);

	t = (long) range * (long) freq;

	t += (long) (freq + x);
	t -= cmax;
	range = (unsigned) (t / cmax);

	n1 = shift_count ((low + range) ^ low);
	if (n1 && under_flow_count)
	{
		clear_under_flow ((low & 0x8000) != 0);
		low ^= 0x8000;
	}

	if (n1)
	{
		put_bit_string (low, n1);
		low <<= n1;
		range = ((range + 1) << n1) - 1;
	}

	if (range < MIN_RANGE)
	{
		n2 = shift_count (range - 1) - 1;

		under_flow_count += n2;
		low <<= n2;
		low &= 0x7FFF;
		range = ((range + 1) << n2) - 1;
	}
}


/*
	Send bit string based on underflow count
	Uses bit value followed by its complement

	generates: 0111... or 1000...
*/
	
static void clear_under_flow (unsigned bit)

{
	put_bit (bit);
	bit ^= 0x01;
	while (-- under_flow_count) put_bit (bit);
}



/*
	Send bit string
*/

static void put_bit_string (unsigned x, int count)

{
	while (count)
	{
		put_bit ((x & 0x8000) != 0);
		x <<= 1;
		count --;
	}
}



/*
	Write single bit to output file using internal buffer
*/

static void put_bit (int bit)

{
	static unsigned char mask [] =
	{
		0x80,	0x40,	0x20,	0x10,	0x08,	0x04,	0x02,	0x01,
	};

	if (bit) code_value |= mask [code_bit_count];
	if (++ code_bit_count == 8)
	{
		WriteOutputFile (code_value);
		code_value = 0;
		code_bit_count = 0;
		codesize ++;
	}
}



/*
	Determine number of leading zero bits on unsigned value
*/

static int shift_count (unsigned n)

{
	int i;

	i = 0;
	while ((n & 0x8000) == 0 && i < 16)
	{
		i ++;
		n <<= 1;
	}

	return i;
}



/*
	Read bit string from input stream
	Length is limited to word size
*/

static unsigned get_bit_string (int n)

{
	unsigned x = 0;

	while (n)
	{
		x <<= 1;
		x += get_bit ();
		n --;
	}

	return x;
}



/*
	Read a single bit from input stream
	Uses a one byte buffer for intermediate values
*/

static unsigned get_bit (void)

{
	int n;

	if (code_bit_count == 0)
	{
		code_bit_count = 8;
		n = ReadInputFile ();
		code_value = n >= 0 ? n : 0;
	}

	code_bit_count --;
	n = code_value & 0x80;
	code_value <<= 1;

	return n != 0;
}



/*
	Initialize arithmetic decode procedure
*/

void StartDecode (void)

{
	int i;

	value = 0;
	for (i = 0; i < 16; i ++)
	{
		value <<= 1;
		value |= get_bit ();
	}
}


/*
	Return value of next input code
	Input consists of total frequency for active symbol set
	Note that decoder must be updated with actual frequencies used
*/

int DecodeArith (unsigned cmax)

{
	unsigned long t;

	t = (long) (value - low + 1) * (long) cmax;
	t -= 1;
	t /= (long) range + 1;

	return (unsigned) t;
}


/*
	Update arithmetic decoder using actual symbol frequencies
	Read additional bits from input based on symbol values
*/

void UpdateDecoder (unsigned base, unsigned freq, unsigned cmax)

{
	unsigned long t;
	unsigned x, y;

	int n1, n2;

	t  = (long) range * (long) base;
	t += (long) base;
	x  = (unsigned) (t / cmax);
	y  = (unsigned) (t % cmax);

	low += x;

	t = (long) range * (long) freq;
	t += (long) (freq + y);
	t -= cmax;
	range = (unsigned) (t / cmax);

	n1 = shift_count ((low + range) ^ low);
	if (n1)
	{
		low <<= n1;
		range = ((range + 1) << n1) - 1;
		value <<= n1;
		value |= get_bit_string (n1);
	}

	if (range < MIN_RANGE)
	{
		n2 = shift_count (range - 1) - 1;
		value -= low;
		low <<= n2;
		low &= 0x7FFF;
		range = ((range + 1) << n2) - 1;
		value <<= n2;
		value |= get_bit_string (n2);
		value += low;
	}
}