#define QUANTIZER 0

/*****

on checa 256 @ 0.50 bpp

z = 2 : psnr 36.89
z = 0 : psnr 37.16

so fixing is slightly better, but not much

on checa 256 @ 0.25 bpp

z = 10 : psnr 31.78
z = 4  : psnr 32.54
z = 0  : psnr 33.09

*****/

//#define MTF_ADDED 	// move it to front ; hurts 0.007
//#define MTF_ONE_ADDED 	// identical to the big mtf_added ! ?

#define FIRST_MATCH		//else best : first HELPS 0.02 , big surprise

#define MTF_BIGSTEP		// helps 0.05
#define BIGSTEP_SHIFT	2	/** two seems optimal **/

#define DO_LOG

//#define AMORTIZE_ERROR
#define AMORTIZE_MAX_ERR 	50
#define AMORTIZE_MAX_IDX	4		/** best at -1 **/
#define AMORTIZE_SHIFTDOWN	3
	/*** this never helps, contrary to intuition.
	*	compression is *VERY* sensitive to the exact setting of all of these param!  
	*		lots of wierd local minima & maxima
	*	Changes of even 1 make a huge difference!
	*		and the optimum varies strongly with 'q' and 'z' ;
	****/

/************

an LRU vector quantizer.

this simple coder is quite competitive at high 'q' and low 'l'

	we need high 'q' because our "escape" is pretty poor
	we need low 'l' because we handle the top levels so poorly

	about 0.05 bpp off of the "bitplane" coder, only slightly
		better than the "order 0" coder.

todos :

	0. Amortize_Max todos :
		addVec() currently adds a new one : should we replace the amortized one instead?

	1. MTF_BIGSTEP helped so much, we should try more improvements of the lru walking 

	2. we're operating as a *terrible* quantizer right now; increasing the mainline
		quantizer helps much more than increasing our "quantizer".  We work best lossless.

performance concerns :

	0. scontext with such large alphabets is awfully slow

	1. code literals as delta to nearest match?

	2.	possible inefficiency :
			add a vec A to the database
			code a vec B which differs only slightly
			code B 9999 times -> large total error,
			A stays in database, B never added
		this is not as rare as it might seem; for example the
			common vector "0010" will never get added if
			"0000" is added first.
		how to fix?
			always send delta after coding from a vector?
				(tried : doesn't seem to change the R-D curve)
			detect this special case of a vector in error being
				sent many times and explicitly choose to send it
				as an escape?  (hard to code this logic).
			**** VQ is fixing about (1/3) of matches now, I think this is critical

				tried to handle this with the "Amortize Error" option, which	
				adds a new vector whenever it seems that an old one is being
				coded from often & inexactly.


theory from Vetterli :

	1. periodically update lru-vectors to the midpoints of the source-vectors coded onto them

	2. don't always add a new escaped-out vector ; only do so if the cost of the escape is
		less than the cost incurred by adding a new vector

	3. periodically shrink the dictionary to cut chaffe ;
		not really important if we add often and entropy code correctly

	4. VQ-Vetterli gets 30 dB at 0.5 bpp on Lena !  SPIHT gets 37 dB !!!
			VQ bites ASS!!!

tried, didn't work :

	1. add vectors after fixing (hurt oodles!)

	2. mtf nearest when failing a match (on fenwick's tip) ; didn't hurt, and slow

*************/

#include <stdio.h>
#include <stdlib.h>
#include <crblib/inc.h>
#include <crblib/arithc.h>
#include <crblib/scontext.h>
#include <crblib/intmath.h>

typedef struct {
	int tot_err;
	int A,B,C,D;
} vector;

#ifdef DO_LOG
static int matches=0,escapes=0,fixes=0;
#endif

#ifdef DO_LOG
#define LOG(x)	if(0) ; else { x; }
#else
#define LOG(x)
#endif

extern int tune_param;

#define NUM_VECS		100		/** compression roughly monotonically decreases with more vecs **/
#define VEC_ESCAPE		NUM_VECS

#define VEC_CNTXMAX			4		/** helps a bunch when 'quant' is small, irrelevant when large **/
#define VEC_CONTEXTS		(VEC_CNTXMAX+1)
#define VEC_CONTEXT(cntx)	min(VEC_CNTXMAX,intlog2(cntx))

#define VEC_TOTMAX		5000
#define VEC_INC			15
#define VEC_ALPHABET	(VEC_ESCAPE+1)

#define LIT_ESCAPE		30
#define LIT_TOTMAX		5000
#define LIT_INC			5
#define LIT_ALPHABET	(LIT_ESCAPE+1)

#include "coder.h"

void coderVQ_encodeBand(coder *me,int *band,int w,int h,int fullw,int *parent);
void coderVQ_decodeBand(coder *me,int *band,int w,int h,int fullw,int *parent);

typedef struct {
	scontext **vec_o1,*lit_o0,*delta_o0;
	vector * vec_alloc;
	vector ** lru;
	arithInfo * ari;
	int max_err;
	bool fix_errors;
} myInfo;

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

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

	c->data = d;

	d->ari = c->arith;

	if ( (d->vec_alloc = newarray(vector,NUM_VECS)) == NULL )
		errexit("alloc failed");
	if ( (d->lru = newarray(vector *,NUM_VECS)) == NULL )
		errexit("alloc failed");
	for(i=0;i<NUM_VECS;i++) {
		d->lru[i] = &(d->vec_alloc[i]);
	}

	d->max_err = QUANTIZER * QUANTIZER;
		/** we do the quantizing, not just the coding
		*	err = mse * 4  (four pels per vec)
		*		= (quantizer/2)^2 * 4 = quantizer^2
		***/

	if ( d->max_err == 0 ) {
		d->max_err = 2;
		d->fix_errors = true;
	}

	/*** tuning shows roughly an 0.3 psnr change for each increment of max_err
	**		(= 0.6 mse change , or almost exacly 0.1 rmse change)
	*	we get about an 0.1 bpp improvement for each increment of max_err up to 2,
	*		and then it slows to 0.02 bpp.  Hence our choice.
	*****/

	/** <> load in some stock vectors ? ; not very important
	***	since the first plane we pass is random anyway **/

	if ( (d->vec_o1 = newarray(scontext *,VEC_CONTEXTS)) == NULL )
		errexit("alloc failed");

	for(i=0;i<VEC_CONTEXTS;i++) {
		if ( (d->vec_o1[i] = scontextCreate(c->arith,VEC_ALPHABET,0,VEC_TOTMAX,VEC_INC,true)) == NULL )
			errexit("ozero init failed");
	}

	if ( (d->lit_o0 = scontextCreate(c->arith,LIT_ALPHABET,0,LIT_TOTMAX,LIT_INC,true)) == NULL )
		errexit("ozero init failed");

	if ( (d->delta_o0 = scontextCreate(c->arith,3,0,10000,30,true)) == NULL )
		errexit("ozero init failed");
}

void coderVQ_free(coder *c)
{

#ifdef DO_LOG
if ( matches+escapes > 0 ) {
	printf("matches = %d, fixes = %d, escapes = %d\n",matches,fixes,escapes);
	matches = fixes = escapes = 0;
}
#endif

	if ( c->data ) {
		myInfo *d;
		d = c->data;
		smartfree(d->vec_alloc);
		smartfree(d->lru);
		if ( d->vec_o1 ) { int i;
			for(i=0;i<VEC_CONTEXTS;i++)
				if ( d->vec_o1[i] ) scontextFree(d->vec_o1[i]);
			free(d->vec_o1);
		}
		if ( d->lit_o0 ) scontextFree(d->lit_o0);
		if ( d->delta_o0 ) scontextFree(d->delta_o0);
		free(d);
		c->data = NULL;
	}
}

coder coderVQ = {
		"VQ",
		coderVQ_init,
		coderVQ_free,
		coderVQ_encodeBand,
		coderVQ_decodeBand
	};


static int diffVec(vector *x,vector *y)
{
return
	(x->A - y->A)*(x->A - y->A) +
	(x->B - y->B)*(x->B - y->B) +
	(x->C - y->C)*(x->C - y->C) +
	(x->D - y->D)*(x->D - y->D);
}

static void mtfVec(myInfo *mi,int i)
{
int steps;
vector *swap;

	if ( i == 0 ) return; //do nothing

#ifdef MTF_BIGSTEP
	steps = i>>(BIGSTEP_SHIFT);
	if ( steps < 1 ) steps=1;
#else /** just one step **/
	steps = 1;
#endif
	swap = mi->lru[i];
	while(steps--) {
		mi->lru[i] = mi->lru[i-1];
		i--;
	}
	mi->lru[i] = swap;
}

static void mtfOneVec(myInfo *mi,int i)
{
vector *swap;

	if ( i == 0 ) return; //do nothing

	swap = mi->lru[i];
	mi->lru[i] = mi->lru[i-1];
	mi->lru[i-1] = swap;
}

static void addVec(myInfo *mi,vector *vec)
{
int tail;

	/** add at tail **/

	tail = NUM_VECS -1;
	memcpy(mi->lru[tail],vec,sizeof(vector));
	mi->lru[tail]->tot_err = 0;

#ifdef MTF_ADDED // move it to front
	mtfVec(mi,tail);
#endif
#ifdef MTF_ONE_ADDED // slide it a taddle
	mtfOneVec(mi,tail);
#endif
}

static void encode_esc(scontext *o0,arithInfo *ari,int val)
{
if ( val == 0 ) { scontextEncode(o0,0); return; }
else { 
	int v = abs(val);
	if ( v < LIT_ESCAPE ) scontextEncode(o0,v);
	else {
		scontextEncode(o0,LIT_ESCAPE);
		encode_m1(ari,v - LIT_ESCAPE);
	}
	if ( isneg(val) ) arithBit(ari,1);
	else arithBit(ari,0);
	}
}

static int decode_esc(scontext *o0,arithInfo *ari)
{
int val;
	val = scontextDecode(o0);
	if ( val == 0 ) return 0;
	else if ( val == LIT_ESCAPE ) {
		val += decode_m1(ari);
	}
	if ( arithGetBit(ari) ) val = -val;
return val;
}

static void encodeDelta(myInfo *mi,vector *delta,int err)
{
int pos,sign,sign2;

	switch(err) {
		case 0: 
			break;
		case 1:
			// 8 values : 2 bit pos, 1 bit sign

#if 0
pos = abs(delta->A) + abs(delta->B) + abs(delta->C) + abs(delta->D);
if ( pos > 1 ) errexit("err > 1");
#endif

			if ( delta->A ) { pos=0; sign = delta->A; }
			if ( delta->B ) { pos=1; sign = delta->B; }
			if ( delta->C ) { pos=2; sign = delta->C; }
			if ( delta->D ) { pos=3; sign = delta->D; }

			if ( sign == -1 ) pos += 4;
			arithEncode(mi->ari,pos,pos+1,8);
			break;
		case 2:
			// 6 positions, 4 sign values

			if ( delta->A ) { sign = delta->A; 
				if ( delta->B ) { pos=0; sign2 = delta->B; }
				else if ( delta->C ) { pos=1; sign2 = delta->C; }
				else if ( delta->D ) { pos=2; sign2 = delta->D; }
				else errexit("should not get here");
			} else if ( delta->B ) { sign = delta->B; 
				if ( delta->C ) { pos=3; sign2 = delta->C; }
				else if ( delta->D ) { pos=4; sign2 = delta->D; }
				else errexit("should not get here");
			} else { pos =5; sign = delta->C; sign2 = delta->D; }

			if ( sign == - 1) pos += 6;
			if ( sign2 == - 1) pos += 12;

			arithEncode(mi->ari,pos,pos+1,24);
			break;
	}
}

static void decodeDelta(myInfo * mi,vector *vec,int err) /** add the delta onto vec **/
{
int pos,sign,sign2;

	switch(err) {
		case 0: 
			break;
		case 1:
			pos = arithGet(mi->ari,8); arithDecode(mi->ari,pos,pos+1,8);
			if ( pos&4 ) { pos -=4; sign = -1; }
			else sign = 1;
			switch(pos) {
				case 0: vec->A += sign; break;
				case 1: vec->B += sign; break;
				case 2: vec->C += sign; break;
				case 3: vec->D += sign; break;
			}
			break;
		case 2:
			pos = arithGet(mi->ari,24); arithDecode(mi->ari,pos,pos+1,24);
			if ( pos >= 12) { pos -= 12; sign2 = -1; } else sign2 = 1;
			if ( pos >= 6) { pos -= 6; sign = -1; } else sign = 1;

			switch(pos) {
				case 0: vec->A += sign; vec->B += sign2; break;
				case 1: vec->A += sign; vec->C += sign2; break;
				case 2: vec->A += sign; vec->D += sign2; break;
				case 3: vec->B += sign; vec->C += sign2; break;
				case 4: vec->B += sign; vec->D += sign2; break;
				case 5: vec->C += sign; vec->D += sign2; break;
				default: errexit("should not get here");
			}
			break;
	}
}

void coderVQ_encodeBand(coder *me,int *band,int width,int height,int fullw,int *parent)
{
int x,y,cntx,i,cur_err;
int best_i,best_err;
int *dp,*pp,*dpn;
vector vec,*vs,delta;
myInfo *mi = (myInfo *)(me->data); 
arithInfo *ari = mi->ari;

	dp = band;	pp = parent;
	for(y=0;y<height;y+=2) {
		dpn = dp + fullw;
		if ( coder_timetostop(me) ) { coder_didstop(me,y); return; }
		for(x=0;x<width;x+=2) {
			cntx = abs(pp[x>>1]);
			cntx = VEC_CONTEXT(cntx);

			vec.A = dp[x];	vec.B = dp[x+1];
			vec.C = dpn[x];	vec.D = dpn[x+1];

			best_err = mi->max_err + 1; best_i = NUM_VECS;
			for(i=0;i<NUM_VECS;i++) {
				vs = mi->lru[i];
				cur_err = diffVec(vs,&vec);
				if ( cur_err < best_err ) {
					best_i = i; best_err = cur_err;	
#ifdef FIRST_MATCH
					break;
#endif
				}
			}

#ifdef AMORTIZE_ERROR
			if ( !mi->fix_errors && best_i < NUM_VECS ) {
				if ( mi->lru[best_i]->tot_err > AMORTIZE_MAX_ERR && best_i <= AMORTIZE_MAX_IDX ) {
					mi->lru[best_i]->tot_err >>= AMORTIZE_SHIFTDOWN;
					best_i = NUM_VECS; /** force escape **/
				}
			}
#endif
			if ( best_i == NUM_VECS ) {	/** escape **/
				LOG(escapes++);
				scontextEncode(mi->vec_o1[cntx],VEC_ESCAPE);
				encode_esc(mi->lit_o0,ari,vec.A); encode_esc(mi->lit_o0,ari,vec.B);
				encode_esc(mi->lit_o0,ari,vec.C); encode_esc(mi->lit_o0,ari,vec.D);
				addVec(mi,&vec);
			} else {
				i = best_i; vs = mi->lru[i];
				scontextEncode(mi->vec_o1[cntx],i);
				mtfVec(mi,i);
				LOG(matches++);
				if ( mi->fix_errors ) {
					scontextEncode(mi->delta_o0,best_err);
					if ( best_err > 0 ) {
						LOG(fixes++);
						vs->tot_err += best_err;
						delta.A = vec.A - vs->A;
						delta.B = vec.B - vs->B;
						delta.C = vec.C - vs->C;
						delta.D = vec.D - vs->D;
						encodeDelta(mi,&delta,best_err);
#ifdef AMORTIZE_ERROR
						if ( vs->tot_err > AMORTIZE_MAX_ERR && i <= AMORTIZE_MAX_IDX ) {
							vs->tot_err >>= AMORTIZE_SHIFTDOWN;
							addVec(mi,&vec); 
						}
#endif
					}
				} else {
					if ( best_err != 0 ) {	/** must put in the used values for context coding **/
						dp[x] =  vs->A;
						dp[x+1]= vs->B;
						dpn[x] = vs->C;
						dpn[x+1]=vs->D;
						vs->tot_err += best_err;
						/** decoder doesn't know what this error is, but we can
						** 	use it for amortizing, because we explicitly send an escape
						***/
					}
				}
			}
		}
		pp += fullw;
		dp += fullw + fullw;
	}
}

void coderVQ_decodeBand(coder *me,int *band,int width,int height,int fullw,int *parent)
{
int x,y,cntx,idx,err;
int *dp,*pp,*dpn;
vector *vs,vec;
myInfo *mi = (myInfo *)(me->data); 
arithInfo *ari = me->arith;

	dp = band;	pp = parent;
	for(y=0;y<height;y+=2) {
		dpn = dp + fullw;
		if ( coder_timetostopd(me,y) ) return;
		for(x=0;x<width;x+=2) {	/** x & y are the parent's location *2 **/
			cntx = abs(pp[x>>1]);
			cntx = VEC_CONTEXT(cntx);

			idx = scontextDecode(mi->vec_o1[cntx]);

			if ( idx == VEC_ESCAPE ) {
				dp[x]	= vec.A = decode_esc(mi->lit_o0,ari);
				dp[x+1] = vec.B = decode_esc(mi->lit_o0,ari);
				dpn[x] 	= vec.C = decode_esc(mi->lit_o0,ari);
				dpn[x+1]= vec.D = decode_esc(mi->lit_o0,ari);
				addVec(mi,&vec);
			} else {
				vs = mi->lru[idx];
				mtfVec(mi,idx);
				if ( mi->fix_errors ) {
					err = scontextDecode(mi->delta_o0);
					if ( err ) {
						vs->tot_err += err;
						memcpy(&vec,vs,sizeof(vector));
						decodeDelta(mi,&vec,err);
						dp[x]  = vec.A;	dp[x+1]  = vec.B;
						dpn[x] = vec.C;	dpn[x+1] = vec.D;
#ifdef AMORTIZE_ERROR
						if ( vs->tot_err > AMORTIZE_MAX_ERR && idx <= AMORTIZE_MAX_IDX ) {
							vs->tot_err >>= AMORTIZE_SHIFTDOWN;
							addVec(mi,&vec); 
						}
#endif
					} else {
						dp[x]  = vs->A;	dp[x+1]  = vs->B;
						dpn[x] = vs->C;	dpn[x+1] = vs->D;
					}
				} else {
					dp[x]  = vs->A;	dp[x+1]  = vs->B;
					dpn[x] = vs->C;	dpn[x+1] = vs->D;
				}
			}
		}
		pp += fullw;
		dp += fullw + fullw;
	}
}