// <> "stop" functionality needed

//#define WRITE_ZT
#define WRITE_ZT_PLANE 7

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

#define SIGN_CONTEXTS		4

#define TOTMAX			4000
#define INC				30

#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)

typedef struct {
	arithInfo *ari;
	int signs_p0[SIGN_CONTEXTS],signs_pt[SIGN_CONTEXTS];
} signInfo;

void * initSigns(arithInfo *ari);
void freeSigns(void *sign_coder);

void encodeSigns(signInfo *me,int *band,int w,int h,int fullw,int *parent);
void decodeSigns(signInfo *me,int *band,int w,int h,int fullw,int *parent);

void * initSigns(arithInfo *ari)
{
signInfo *d;
int i;

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

	d->ari = ari;

	for(i=0;i<SIGN_CONTEXTS;i++) {
		d->signs_p0[i] = 100;
		d->signs_pt[i] = 200;
	}
return d;
}

void freeSigns(signInfo *d)
{
	if ( d ) {
		free(d);
	}
}

void encodeSigns(signInfo *me,int *band,int width,int height,int fullw,int *parent)
{
int x,y,context; //,V;
int *dp;
arithInfo *ari;
int *signs_p0,*signs_pt;

	ari = me->ari;
	signs_p0 = me->signs_p0;
	signs_pt = me->signs_pt;

	dp = band;

	for(x=0;x<width;x++) {
		if ( *dp == 0 ) { dp++; continue; }
		if ( x == 0 ) context = 0;
		else context = (dp[-1]&1);
		bitEnc((*dp&1),ari,signs_p0[context],signs_pt[context]);
		dp++;
	}
	dp += fullw - width;

	for(y=1;y<height;y++) {
		if ( *dp == 0 ) dp++;
		else {
			context = ((dp[-fullw]&1)<<1);
			bitEnc((*dp&1),ari,signs_p0[context],signs_pt[context]);
			dp++;
		}

		for(x=1;x<width;x++) {
			if ( *dp == 0 ) { dp++; continue; }
			context = (dp[-1]&1) + ((dp[-fullw]&1)<<1);

			bitEnc((*dp&1),ari,signs_p0[context],signs_pt[context]);
			dp++;
		}
		dp += fullw - width;
	}
}

void decodeSigns(signInfo *me,int *band,int width,int height,int fullw,int *parent)
{
int x,y,sign,context;
int *dp;
arithInfo *ari;
int *signs_p0,*signs_pt;

	ari = me->ari;
	signs_p0 = me->signs_p0;
	signs_pt = me->signs_pt;

	dp = band;

	for(x=0;x<width;x++) {
		if ( *dp == 0 ) { dp++; continue; }
		if ( x == 0 ) context = 0;
		else context = (dp[-1]&1);
		bitDec(sign,ari,signs_p0[context],signs_pt[context]);
		*dp++ += sign;
	}
	dp += fullw - width;

	for(y=1;y<height;y++) {

		if ( *dp == 0 ) dp++;
		else {
			context = ((dp[-fullw]&1)<<1);
			bitDec(sign,ari,signs_p0[context],signs_pt[context]);
			*dp++ += sign;
		}

		for(x=1;x<width;x++) {
			if ( *dp == 0 ) { dp++; continue; }
			context = (dp[-1]&1) + ((dp[-fullw]&1)<<1);

			bitDec(sign,ari,signs_p0[context],signs_pt[context]);
			*dp++ += sign;
		}
		dp += fullw - width;
	}
}


void findZeroTrees(image *im,int bits,int levels,int offx,int offy)
{
int *fmband,*toband,*fmp,*top;
int x,y,l,w,h,wt,ht,fullw,p;
int bitmask;

	bitmask = (1<<bits) + FLAG_ISOLATED_ZERO;

	fullw = im->width;
	for(p=0;p<im->planes;p++) {

		for(l=1;l<levels;l++) {	//** don't do very top level, cuz that would use the LL as parents

			w = (im->width) >>l;
			h = (im->height)>>l;
			wt = w>>1;
			ht = h>>1;
			fmband = im->data[p][h *offy] + w *offx;
			toband = im->data[p][ht*offy] + wt*offx;

			fmp = fmband; top = toband;
			for(y=0;y<ht;y++) {
				for(x=0;x<wt;x++) {
					if ( (*top&bitmask) == 0 ) {	// *top might be a zerotree, (if not, nothing to do)
						if ( ((fmp[0]&bitmask) + (fmp[1]&bitmask) + (fmp[fullw]&bitmask) + (fmp[fullw+1]&bitmask)) == 0 ) { // all kids zero
							// if *top is zero, it can stay zero to be a zero-tree
							// change kids of zerotree roots to children
							fmp[0] |= FLAG_CHILD_ZERO;
							fmp[1] |= FLAG_CHILD_ZERO;
							fmp[fullw  ] |= FLAG_CHILD_ZERO;
							fmp[fullw+1] |= FLAG_CHILD_ZERO;
						} else {
							*top |= FLAG_ISOLATED_ZERO; // label as an isolated zero
						}
					}
					top ++;
					fmp += 2;
				}
				top += fullw - wt;
				fmp += fullw + fullw - w;
			}
		}
	}
}


void encodeImageZT(coder * encoder,image *im,int levels)
{
int top_val,top_bitpn,bits;
int *dp,x,y,**rows;
int p,l,sizeX,sizeY;
void *sign_coder;

	if ( ! encoder->encodeBandZT ) errexit("don't call imageZT unless you mean it!");

	top_val = 0;

	for(p=0;p<im->planes;p++) {
		rows = im->data[p];

		/** global pre-pass : remove signs & find top_val **/

		sizeX = (im->width) >> levels;
		sizeY = (im->height) >> levels;
		for(y=0;y<(im->height);y++) {
			dp = rows[y];
			if ( y < sizeY ) { x = sizeX; dp += sizeX; }
			else x = 0;
			for(;x<(im->width);x++) {
				if ( isneg(*dp) ) *dp = 1 - (*dp + *dp);
				else *dp += *dp;
				if ( *dp > top_val ) top_val = *dp;
				dp++;
			}
		}
	}

	for(top_bitpn=0;(1<<(top_bitpn+1))<=top_val;top_bitpn++) ;

	cu_putExpanding_ari(top_bitpn,encoder->arith,16,8);

	for(bits = top_bitpn;bits>=1;bits--) {
		findZeroTrees(im,bits,levels,1,0);
		findZeroTrees(im,bits,levels,0,1);
		findZeroTrees(im,bits,levels,1,1);

#ifdef WRITE_ZT
		if ( bits == WRITE_ZT_PLANE ) {	/** write out a zero tree map **/
			image *zt_im;
			int r,*ptr,bitmask;
			zt_im = copyImage(im);

			for(p=0;p>im->planes;p++) {
				ptr = zt_im->data[p][0][0]; bitmask = 1<<bits;
				for(r=zt_im->plane_size;r--;) {
					if ( *ptr & bitmask ) *ptr++ = 255;
					else { 
						*ptr = ((*ptr>>(CODE_MAX_BPN+1))<<6); ptr++; 
					}
					/** 256 is on
						128	is isolated
						 64 is child
						  0 is root
					****/

					ptr[-1] -= 0x80;
				}
			}

			writeImageFile("zerotree.raw",zt_im);

			freeImage(zt_im);
		}
#endif			/** end zero-tree writing **/

		for (l = levels; l > 0; l--) {
			sizeX = (im->width) >> l;
			sizeY = (im->height) >> l;

			for(p=0;p<im->planes;p++) {
				rows = im->data[p];

				encoder->encodeBandZT(encoder,rows[0] + sizeX,	sizeX, sizeY,bits,im->width,rows); 
				encoder->encodeBandZT(encoder,rows[sizeY],		sizeX, sizeY,bits,im->width,rows); 
				encoder->encodeBandZT(encoder,rows[sizeY]+sizeX,sizeX, sizeY,bits,im->width,rows); 
			}
		}
	}

	sign_coder = initSigns(encoder->arith);			
	for(p=0;p<im->planes;p++) {
		rows = im->data[p];
		for (l = levels; l > 0; l--) {
			sizeX = (im->width) >> l;
			sizeY = (im->height) >> l;
			encodeSigns(sign_coder,rows[0] + sizeX	,	sizeX, sizeY, im->width,rows[0] + (sizeX>>1));
			encodeSigns(sign_coder,rows[sizeY]		,	sizeX, sizeY, im->width,rows[(sizeY>>1)]);
			encodeSigns(sign_coder,rows[sizeY]+sizeX,	sizeX, sizeY, im->width,rows[(sizeY>>1)]+ (sizeX>>1));
		}
	}
	freeSigns(sign_coder);
}

void decodeImageZT(coder * decoder,image *im,int levels)
{
int l,sizeX,sizeY,x,y,*dp;
int **rows;
int top_bitpn,bits,p;
void *sign_coder;

	top_bitpn = cu_getExpanding_ari(decoder->arith,16,8);

	for(bits = top_bitpn;bits>=1;bits--) {
		for (l = levels; l > 0; l--) {
			for(p=0;p<im->planes;p++) {
				rows = im->data[p];

				sizeX = (im->width) >> l;
				sizeY = (im->height) >> l;

				decoder->decodeBandZT(decoder,rows[0] + sizeX,	sizeX, sizeY,bits,im->width,rows); 
				decoder->decodeBandZT(decoder,rows[sizeY],		sizeX, sizeY,bits,im->width,rows);
				decoder->decodeBandZT(decoder,rows[sizeY]+sizeX,sizeX, sizeY,bits,im->width,rows);
			}
		}
	}

	sign_coder = initSigns(decoder->arith);
	for(p=0;p<im->planes;p++) {
		rows = im->data[p];
		for (l = levels; l > 0; l--) {
			sizeX = (im->width) >> l;
			sizeY = (im->height) >> l;

			decodeSigns(sign_coder,rows[0] + sizeX	,	sizeX, sizeY, im->width,rows[0] + (sizeX>>1));
			decodeSigns(sign_coder,rows[sizeY]		,	sizeX, sizeY, im->width,rows[(sizeY>>1)]);
			decodeSigns(sign_coder,rows[sizeY]+sizeX,	sizeX, sizeY, im->width,rows[(sizeY>>1)]+ (sizeX>>1));
		}
	}
	freeSigns(sign_coder);

		/** global post-pass : restore signs 
		*	(<> decodeSigns could do this for us, but that would require him to
		*		do an if() when he builds contexts) 
		**/
	
	for(p=0;p<im->planes;p++) {
		rows = im->data[p];
		sizeX = (im->width) >> levels;
		sizeY = (im->height) >> levels;
		for(y=0;y<im->height;y++) {
			dp = rows[y];
			if ( y < sizeY ) { x = sizeX; dp += sizeX; }
			else x = 0;
			for(;x<im->width;x++) {
				if ( *dp & 1 ) *dp++ = - ((*dp)>>1);
				else *dp++ >>= 1;
			}
		}
	}
}