/*****

BP Binary Fast ZeroTree

*****/

//#define ZT_STATS

#include <stdio.h>
#include <stdlib.h>
#include <crblib/inc.h>
#include <crblib/arithc.h>
#include "coder.h"

extern int tune_param;

#define	VAL_CONTEXTS		16
#define	VAL_CONTEXT_MAX		(VAL_CONTEXTS -1)
#define SHAPE_BASE			VAL_CONTEXTS
#define SHAPE(x)			(SHAPE_BASE<<(x))
#define NUM_SHAPES			2
#define NUM_CONTEXTS		(VAL_CONTEXTS<<NUM_SHAPES)

#define TOTMAX			4000
#define INC				30

#define P0_INIT			8
#define P1_INIT			0

#define bitModel(bit,P0,PT)		do { PT += INC; if (!(bit)) P0 += INC;  if ( PT > TOTMAX ) { PT >>= 1; P0 >>= 1; P0++; PT += 2; } } while(0)
#define bitEnc(bit,ari,P0,PT)	do { arithEncBit(ari,P0,PT,bit);	bitModel(bit,P0,PT); } while(0)
#define bitDec(bit,ari,P0,PT)	do { bit = arithDecBit(ari,P0,PT);	bitModel(bit,P0,PT); } while(0)

void coderBPBFZT_encodeBandZT(coder *me,int *band,int width,int height,int bits,int fullw,int **rows);
void coderBPBFZT_decodeBandZT(coder *me,int *band,int width,int height,int bits,int fullw,int **rows);

typedef struct {
	int stats_p0[NUM_CONTEXTS],stats_pt[NUM_CONTEXTS];
	int zeros_p0[NUM_CONTEXTS],zeros_pt[NUM_CONTEXTS];
} myInfo;

void coderBPBFZT_init(coder *c)
{
myInfo *d;
int i;

	if ( (d = new(myInfo)) == NULL )
		errexit("ozero init failed");

	c->data = d;

	for(i=0;i<NUM_CONTEXTS;i++) {
		d->stats_p0[i] = P0_INIT+1; d->stats_pt[i] = 2+P0_INIT+P1_INIT;
		d->zeros_p0[i] = 1; d->zeros_pt[i] = 2;
	}

}

void coderBPBFZT_free(coder *c)
{
	if ( c->data ) {
		myInfo *d;

		d = c->data;
		free(d);
		c->data = NULL;
	}
}

coder coderBPBFZT = {
		"BP Bin Fast ZT",
		coderBPBFZT_init,
		coderBPBFZT_free,
		NULL,NULL,NULL,NULL,
		coderBPBFZT_encodeBandZT,
		coderBPBFZT_decodeBandZT
	};

/**********

lazy way to pass the state from getStats to fixStats
	and also interacts with the codeBand()

these are re-initialized at each codeBand() call, so this is
	quite re-entrant as long as we aren't multi-threaded
	(that is, no more than one call to codeBand() at a time)

*********/

static int VD,donemask,nextmask;

static int mcontext(int *dp,int x,int y,int fullw)
{
int N,W;
int context;

	VD	= (*dp)&donemask;	// current val already done

	if ( y == 0 ) {
		N = VD;
		if ( x == 0 ) W = VD; else W = (dp[-1]) & nextmask;
	} else if ( x == 0 ) {
		W = VD;
		N  = (dp[-fullw])	& nextmask;
	} else {
		N = (dp[-fullw])		& nextmask;
		W = (dp[-1])			& nextmask;
	}

	context = min(VAL_CONTEXT_MAX, ((VD + N + W)>>2));

	/** shapes help almost 0.1 bpp **/

	if ( N > VD ) context += VAL_CONTEXTS; 
	if ( W > VD ) context += VAL_CONTEXTS+VAL_CONTEXTS;

return context;
}

void coderBPBFZT_encodeBandZT(coder *me,int *band,int width,int height,int bits,int fullw,int **rows)
{
int x,y,bit,context;
int *dp;
arithInfo *ari = me->arith;
int *stats_p0,*stats_pt;
int *zeros_p0,*zeros_pt;
int bitmask;
bool bottom;
#ifdef ZT_STATS
int zt_kids=0,zt_root=0,zt_on=0,zt_isolated=0;
#endif

	stats_p0 = ((myInfo *)me->data)->stats_p0;
	stats_pt = ((myInfo *)me->data)->stats_pt;
	zeros_p0 = ((myInfo *)me->data)->zeros_p0;
	zeros_pt = ((myInfo *)me->data)->zeros_pt;

	if ( width == (fullw >>1) ) bottom = true;
	else bottom = false;

	bitmask = 1<<bits;
	for(x=bitmask,nextmask=0; x<(1<<CODE_MAX_BPN) ;x<<=1) nextmask += x;
	donemask = nextmask - bitmask;

	dp = band;
	for(y=0;y<height;y++) {	
		for(x=0;x<width;x++) {
			if ( dp[x] & FLAG_CHILD_ZERO ) {
				dp[x] -= FLAG_CHILD_ZERO;	// it's a zerotree child, don't code
#ifdef ZT_STATS
zt_kids ++;
#endif
				continue;
			} else {
				context = mcontext(dp+x,x,y,fullw);

				bit = ((dp[x])&bitmask)?1:0;
				bitEnc(bit,ari,stats_p0[context],stats_pt[context]);
				if ( ! bit  ) {
					if ( ! bottom ) { // at bottom theres no need for zt/isolated flags
						// send : zerotree root/isolate zero
						if ( dp[x] & FLAG_ISOLATED_ZERO ) {
							dp[x] -= FLAG_ISOLATED_ZERO;
							bitEnc(0,ari,zeros_p0[context],zeros_pt[context]);			
#ifdef ZT_STATS
zt_isolated ++;
#endif
						} else {
							bitEnc(1,ari,zeros_p0[context],zeros_pt[context]);	
#ifdef ZT_STATS
zt_root ++;
#endif
						}
					}
				} 
#ifdef ZT_STATS
else zt_on ++;
#endif
			}
		}
		dp += fullw;
	}

#ifdef ZT_STATS
		if ( ! bottom ) {
		int zt_tot;
		zt_tot = zt_kids+zt_root+zt_on+zt_isolated;
		if ( zt_tot != width*height ) { dbf(); errputs("error : didn't scan all pels!"); }
		printf("zt_kids=%2.1f,zt_root=%2.1f,zt_on=%2.1f,zt_isolated=%2.1f,coded=%2.1f\n",
				zt_kids*100.0/zt_tot,zt_root*100.0/zt_tot,zt_on*100.0/zt_tot,zt_isolated*100.0/zt_tot,(zt_tot - zt_kids)*100.0/zt_tot);
		} else {
		int zt_tot;
		zt_tot = zt_kids+zt_on;
		zt_isolated = width*height - zt_tot;
		printf("zt_kids=%2.1f,zt_root=%2.1f,zt_on=%2.1f,zt_isolated=%2.1f,coded=%2.1f\n",
				zt_kids*100.0/zt_tot,zt_root*100.0/zt_tot,zt_on*100.0/zt_tot,zt_isolated*100.0/zt_tot,(zt_tot - zt_kids)*100.0/zt_tot);
		}
#endif

}

void coderBPBFZT_decodeBandZT(coder *me,int *band,int width,int height,int bits,int fullw,int **rows)
{
int x,y,bit,context;
int *dp;
arithInfo *ari = me->arith;
int *stats_p0,*stats_pt;
int *zeros_p0,*zeros_pt;
int bitmask;
int offwidth,offheight;
bool bottom;

	stats_p0 = ((myInfo *)me->data)->stats_p0;
	stats_pt = ((myInfo *)me->data)->stats_pt;
	zeros_p0 = ((myInfo *)me->data)->zeros_p0;
	zeros_pt = ((myInfo *)me->data)->zeros_pt;

	if ( width == (fullw >>1) ) bottom = true;
	else bottom = false;

	offheight	= (band - rows[0])/fullw;
	offwidth	= (band - rows[0]) - fullw*offheight;

	bitmask = 1<<bits;
	for(x=bitmask,nextmask=0; x<(1<<CODE_MAX_BPN) ;x<<=1) nextmask += x;
	donemask = nextmask - bitmask;

	dp = band;
	for(y=0;y<height;y++) {
		for(x=0;x<width;x++) {
			if ( dp[x] & FLAG_ALREADY_DECODED ) {	
				dp[x] -= FLAG_ALREADY_DECODED;
				continue;
			} else {
				context = mcontext(dp+x,x,y,fullw);

				bitDec(bit,ari,stats_p0[context],stats_pt[context]);
				if ( bit ) dp[x] += bitmask;
				else if ( ! bottom ) {	// at bottom there's no need for zt/isolated flags
					bitDec(bit,ari,zeros_p0[context],zeros_pt[context]);
					if ( bit ) {
						int bx,by,tx,ty,n;
						bx = (offwidth + x)<<1;
						by = (offheight+ y)<<1;
						n = 2;

						// <> could unroll as a switch on levels

						while( bx < fullw && by < fullw ) {
							for(tx=bx;tx<(bx+n);tx+=2) {
								for(ty=by;ty<(by+n);ty+=2) {
									rows[ty  ][tx  ]	+= FLAG_ALREADY_DECODED;
									rows[ty  ][tx+1]	+= FLAG_ALREADY_DECODED;
									rows[ty+1][tx  ]	+= FLAG_ALREADY_DECODED;
									rows[ty+1][tx+1]	+= FLAG_ALREADY_DECODED;
								}
							}

							bx += bx; by += by; n += n;
						}
					}
				}
			}
		}
		dp += fullw;
	}
}