/*
Written 2005 By Brendan G Bohannon

I release this under the public domain. 
However, I do requst mention in any derived source if applicable.

This is intended, where possible, to be used as an example of an encoder.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

//#define WHUFF_NOLONGLONG	//'long long' type not present
//#define WHUFF_BITIO		//a few extra functions, may cost speed

#define WHUFF_LOWER	0x00000000
#define WHUFF_UPPER	0xFFFFFFFF
#define WHUFF_NEUTRAL	0x8000
#define WHUFF_WBITS	16
#define WHUFF_WMASK	((1<<WHUFF_WBITS)-1)

#define WARITH2_CTXBITS		16
#define WARITH2_BITS		8

//#define INLINE	inline
#define INLINE

typedef unsigned char byte;
typedef unsigned short ushort;
typedef unsigned int uint;
typedef unsigned long long ull;


//globals

unsigned short warith2_divtab[65536];	//division table
unsigned int warith2_divtab2[65536];	//division table (32 bit)
int warith2_ctx;
unsigned char *warith2_model;	//pairs of counts

int warith2_ctxbits=WARITH2_CTXBITS;
int warith2_ctxmask=(1<<WARITH2_CTXBITS)-1;


//general

int log2up(int v)
{
	int c;
	c=0;
	while(v>1)
	{
		c++;
		v++;
		v>>=1;
	}
	return(c);
}

unsigned int PDRLELZ_DataAdler32(void *buf, int sz, unsigned int lcrc)
{
	byte *s;
	int i, c, s1, s2;

	s=buf;
	s1=lcrc&0xFFFF;
	s2=(lcrc>>16)&0xFFFF;
	for(i=0; i<sz; i++)
	{
		s1=(s1+(s[i]))%65521;
		s2=(s1+s2)%65521;
	}
	return((s2<<16)+s1);
}

byte *huff_winbuf;	//window buffer
int huff_winofs;	//bit offset in byte
int huff_outbits;	//dummy output bits

unsigned int whuff_rmin;		//window lower range
unsigned int whuff_rmax;		//window upper range
unsigned int whuff_rval;		//window decode value


//bit io

INLINE int HUFF_InputByte()
{
	return(*huff_winbuf++);
}

INLINE void HUFF_OutputByte(int i)
{
	*huff_winbuf++=i;
}


//arithmetic coder core

void WHUFF_OutputBit(int i, unsigned int w)
{
	unsigned int r, r2, v;
	int j;

	r=whuff_rmax-whuff_rmin;
#ifdef WHUFF_NOLONGLONG
	v=whuff_rmin+(r>>WHUFF_WBITS)*w;
	v+=((r&WHUFF_WMASK)*w)>>WHUFF_WBITS;
#else
	v=whuff_rmin+((((ull)r)*((ull)(w<<(32-WHUFF_WBITS))))>>32);
#endif

	if(i)whuff_rmin=v+1;
	else whuff_rmax=v;

	//while bits
	while((whuff_rmin&0xFF000000)==(whuff_rmax&0xFF000000))
	{
		HUFF_OutputByte(whuff_rmin>>24);
		whuff_rmin<<=8;
		whuff_rmax<<=8;
		whuff_rmax|=255;
	}
}

int WHUFF_InputBit(unsigned int w)
{
	unsigned int r, r2, v, i;

	r=whuff_rmax-whuff_rmin;
#ifdef WHUFF_NOLONGLONG
	v=whuff_rmin+(r>>WHUFF_WBITS)*w;
	v+=((r&WHUFF_WMASK)*w)>>WHUFF_WBITS;
#else
	v=whuff_rmin+((((ull)r)*((ull)(w<<(32-WHUFF_WBITS))))>>32);
#endif

	i=(whuff_rval>v);
	if(i)whuff_rmin=v+1;
	else whuff_rmax=v;

	//while bits
	while((whuff_rmin&0xFF000000)==(whuff_rmax&0xFF000000))
	{
		whuff_rmin<<=8;
		whuff_rmax<<=8;
		whuff_rmax|=255;

		whuff_rval<<=8;
		whuff_rval|=HUFF_InputByte();
	}

	return(i);
}

int WHUFF_InputModelBit(unsigned char *model, int ctx)
{
	unsigned int r, v, i, w;

//	w=warith2_divtab2[((unsigned short *)model)[ctx]];
	r=whuff_rmax-whuff_rmin;
//	v=whuff_rmin+(((ull)r*(ull)w)>>32);

#ifdef WHUFF_NOLONGLONG
	w=warith2_divtab[((unsigned short *)model)[ctx]];
	v=whuff_rmin+(r>>WHUFF_WBITS)*w;
	v+=((r&WHUFF_WMASK)*w)>>WHUFF_WBITS;
#else
	w=warith2_divtab2[((unsigned short *)model)[ctx]];
	v=whuff_rmin+(((ull)r*(ull)w)>>32);
#endif

	i=(whuff_rval>v);
	if(i)
	{
		whuff_rmin=v+1;
		model[(ctx<<1)|1]++;

		if(!model[(ctx<<1)|1])
		{
			model[(ctx<<1)|0]>>=1;
			model[(ctx<<1)|1]=0x80;
		}
	}else
	{
		whuff_rmax=v;
		model[ctx<<1]++;

		if(!model[ctx<<1])
		{
			model[(ctx<<1)|0]=0x80;
			model[(ctx<<1)|1]>>=1;
		}
	}

	//while bits
	while(!((whuff_rmin^whuff_rmax)&0xFF000000))
	{
		whuff_rmin<<=8;
		whuff_rmax<<=8;
		whuff_rmax|=255;

		whuff_rval<<=8;
		whuff_rval|=HUFF_InputByte();
	}

	return(i);
}

void WHUFF_FlushWBits()
{
	while((whuff_rmin!=WHUFF_LOWER) || (whuff_rmax!=WHUFF_UPPER))
	{
		HUFF_OutputByte(whuff_rmin>>24);

		whuff_rmin<<=8;
		whuff_rmax<<=8;
		whuff_rmax|=255;
	}
}


int WHUFF_SetupEncode(byte *out)
{
	huff_outbits=0;
	huff_winbuf=out;
	huff_winofs=0;

	whuff_rmin=WHUFF_LOWER;
	whuff_rmax=WHUFF_UPPER;
}

int WHUFF_SetupDecode(byte *in)
{
	int i;

	huff_outbits=0;
	huff_winbuf=in;
	huff_winofs=0;

	whuff_rmin=WHUFF_LOWER;
	whuff_rmax=WHUFF_UPPER;

	whuff_rval=WHUFF_LOWER;
	for(i=0; i<4; i++)
		whuff_rval=(whuff_rval<<8)|HUFF_InputByte();
}


//arithmatic coder model/init

//arithmetic frontend/modeler
#if 1
INLINE int WArith2_Predict(unsigned char *model, int ctx)
{
	return(warith2_divtab[((unsigned short *)model)[ctx]]);
}
#endif

int WArith2_Update(unsigned char *model, int ctx, int v)
{
	if(v)
	{
		model[(ctx<<1)|1]++;
		if(!model[(ctx<<1)|1])
		{
			model[(ctx<<1)|0]>>=1;
			model[(ctx<<1)|1]=0x80;
		}
	}else
	{
		model[(ctx<<1)]++;
		if(!model[ctx<<1])
		{
			model[(ctx<<1)|0]=0x80;
			model[(ctx<<1)|1]>>=1;
		}
	}
}

int WArith2_EncodeSymbol(int v)
{
	int i, j, k;

	i=WARITH2_BITS;
	while(i--)
	{
		j=(v>>i)&1;

		k=WArith2_Predict(warith2_model, warith2_ctx);
		WArith2_Update(warith2_model, warith2_ctx, j);

		WHUFF_OutputBit(j, k);
		warith2_ctx=((warith2_ctx<<1)|j)&warith2_ctxmask;
	}
}

int WArith2_DecodeSymbol()
{
	int i, j;

	i=WARITH2_BITS;
	while(i--)
	{
		j=WHUFF_InputModelBit(warith2_model, warith2_ctx);
		warith2_ctx=((warith2_ctx<<1)|j)&warith2_ctxmask;
	}

	return(warith2_ctx&((1<<WARITH2_BITS)-1));
}

int WArith2_SetupEncode(byte *out)
{
	int i, j, k;

	for(i=0; i<256; i++)
		for(j=0; j<256; j++)
	{
		k=((i+1)<<16)/(i+j+2);
		warith2_divtab[i+(j<<8)]=k;
		warith2_divtab2[i+(j<<8)]=k<<16;
	}

	j=1<<warith2_ctxbits;
	warith2_model=malloc(j*2*sizeof(unsigned char));
	for(i=0; i<j*2; i++)warith2_model[i]=0;
	warith2_ctx=0;

	WHUFF_SetupEncode(out);
}

int WArith2_SetupDecode(byte *in)
{
	int i, j, k;

	for(i=0; i<256; i++)
		for(j=0; j<256; j++)
	{
		k=((i+1)<<16)/(i+j+2);
		warith2_divtab[i+(j<<8)]=k;
		warith2_divtab2[i+(j<<8)]=k<<16;
	}

	j=1<<warith2_ctxbits;
	warith2_model=malloc(j*2*sizeof(unsigned char));
	for(i=0; i<j*2; i++)warith2_model[i]=0;
	warith2_ctx=0;

	WHUFF_SetupDecode(in);
}


//BPAQ0 encoder/decoder

int BPAQ0_EncodeData(byte *in, byte *out, int sz)
{
	int i, j, k, l, ll;
	byte *cs, *ce;

	if(!sz)return(0);

	out[0]='B';
	out[1]='P';
	out[2]='A';
	out[3]='Q';

	out[4]=0;			//version
	out[5]=16;			//header size
	out[6]=WARITH2_BITS;		//data bits
	out[7]=warith2_ctxbits;		//context bits

	out[8]=sz&0xFF;
	out[9]=(sz>>8)&0xFF;
	out[10]=(sz>>16)&0xFF;
	out[11]=(sz>>24)&0xFF;

	i=PDRLELZ_DataAdler32(in, sz, 1);
	out[12]=i&0xFF;
	out[13]=(i>>8)&0xFF;
	out[14]=(i>>16)&0xFF;
	out[15]=(i>>24)&0xFF;

	WArith2_SetupEncode(out+16);


	cs=in;
	ce=in+sz;
	while(cs<ce)
	{
		j=*cs++;
		WArith2_EncodeSymbol(j);
	}

	WHUFF_FlushWBits();
	i=huff_winbuf-out;

	return(i);
}

int BPAQ0_DecodeData(byte *in, byte *out)
{
	byte *cs, *ce;
	int i, j, sz, crc, crc2;

	if((in[0]!='B') || (in[1]!='P') || (in[2]!='A') || (in[3]!='Q'))
		return(-1);

	if(in[4]!=0)return(-2);
	if(in[6]!=8)return(-3);
	warith2_ctxbits=in[7];

	sz=in[8]+(in[9]<<8)+(in[10]<<16)+(in[11]<<24);
	crc=in[12]+(in[13]<<8)+(in[14]<<16)+(in[15]<<24);

	WArith2_SetupDecode(in+in[5]);


	cs=out;
	ce=out+sz;
	while(cs<ce)
	{
		j=WArith2_DecodeSymbol();
		*cs++=j;
	}

	crc2=PDRLELZ_DataAdler32(out, sz, 1);
	if(crc!=crc2)return(-4);

	return(sz);
}


int fail(char *msg)
{
	printf(msg);
	exit(-1);
}

byte *readfile(char *name, int *rsz)
{
	FILE *fd;
	byte *buf;
	int i;

	fd=fopen(name, "rb");
	if(!fd)fail("can't open input\n");

	fseek(fd, 0, 2);
	i=ftell(fd);
	fseek(fd, 0, 0);

	buf=malloc(i);
	fread(buf, 1, i, fd);
	fclose(fd);

	*rsz=i;
	return(buf);
}

int writefile(char *name, byte *buf, int sz)
{
	FILE *fd;

	fd=fopen(name, "wb");
	if(!fd)fail("can't open output\n");

	fwrite(buf, 1, sz, fd);
	fclose(fd);

	return(0);
}

int main(int argc, char *argv[])
{
	byte *ibuf, *obuf;
	int isz, osz, t;
	float ft;

	if((argc<3) || (strcmp(argv[1], "-c") && strcmp(argv[1], "-d")))
	{
		printf("usage: %s mode infile outfile\n");
		printf("where mode is -c or -d, for compress/decompress\n");
		return(-1);
	}

	if(!strcmp(argv[1], "-c"))
	{
		ibuf=readfile(argv[2], &isz);
		obuf=malloc(2*isz);

		t=clock();
		osz=BPAQ0_EncodeData(ibuf, obuf, isz);

		t=clock()-t;
		ft=t/(float)CLOCKS_PER_SEC;
		printf("isz %d osz %d ratio %d%% dt %fs rate %fkB/s\n",
			isz, osz, 100*osz/isz, ft, isz/(ft*1024));

		writefile(argv[3], obuf, osz);
	}

	if(!strcmp(argv[1], "-d"))
	{
		ibuf=readfile(argv[2], &isz);

		osz=ibuf[8]+(ibuf[9]<<8)+(ibuf[10]<<16)+(ibuf[11]<<24);

		obuf=malloc(2*osz);

		t=clock();
		osz=BPAQ0_DecodeData(ibuf, obuf);

		t=clock()-t;
		ft=t/(float)CLOCKS_PER_SEC;
		printf("isz %d osz %d ratio %d%% dt %fs rate %fkB/s\n",
			isz, osz, 100*osz/isz, ft, isz/(ft*1024));

		if(osz<0)switch(osz)
		{
		case -1:
			fail("bad magic\n");
			break;
		case -2:
			fail("bad version\n");
			break;
		case -3:
			fail("bad data bits\n");
			break;
		case -4:
			fail("bad checksum\n");
			break;
		default:
			fail("some error\n");
			break;
		}

		writefile(argv[3], obuf, osz);
	}
}