#include "common.h"
#include "jive.h"

int jiveL_quantize_with_dither(lua_State *L) {

	/* stack is:
		 1: surface
		 2: quantize Rmask
		 3: quantize Gmask
		 4: quantize Bmask
		 5: dodither - boolean to enable/disable dithering

		 Implementation of Quantizer with Floyd-Steinberg dithering - see http://en.wikipedia.org/wiki/Floyd-Steinberg_dithering
		 
		 Note this is very naive at present with minimal work to make it go fast!  I'm not even sure its right...

		 Assumptions:
		 - image stored in 4 bytes per pixel format
		 - Amask covers most significant bit so signed ints can be used for RGB components
	*/

	Uint32 x, y, w, h, Rmask, Gmask, Bmask, Amask, qRmask, qGmask, qBmask, Roff, Goff, Boff, fRmask, fGmask, fBmask;
	Uint32 pixel, new, Acomp, *start;
	Sint32 Rerr, Gerr, Berr, fRerr, fGerr, fBerr, Rcomp, Gcomp, Bcomp;
	int dodither;

	JiveSurface *srf = tolua_tousertype(L, 1, 0);
	SDL_Surface *sdl = srf->sdl;

	// masks to quantize to
	qRmask = (Uint32) luaL_checkint(L, 2);
	qGmask = (Uint32) luaL_checkint(L, 3);
	qBmask = (Uint32) luaL_checkint(L, 4);
	dodither = lua_toboolean(L, 5);
	
	w = sdl->w;
	h = sdl->h;
	Rmask = sdl->format->Rmask;
	Gmask = sdl->format->Gmask;
	Bmask = sdl->format->Bmask;
	Amask = sdl->format->Amask;

	// offsets to add to each component before quantizing (so errors are signed)
	Roff = (qRmask >> 1) & ~qRmask & Rmask;
	Goff = (qGmask >> 1) & ~qGmask & Gmask;
	Boff = (qBmask >> 1) & ~qBmask & Bmask;
	
	// masks to apply to fractional error components to prevent them bleading into less significant components

	// FIXME: calculate these!!!!
	fRmask = 0xffff0000;
	fGmask = 0xffffff00;
	fBmask = 0xffffffff;

	//printf("image w = %d h = %d\n", w, h);
	//printf("orig Rmask = %08x qRmask = %08x Roff = %08x\n", Rmask, qRmask, Roff);
	//printf("orig Gmask = %08x qGmask = %08x Goff = %08x\n", Gmask, qGmask, Goff);
	//printf("orig Bmask = %08x qBmask = %08x Boff = %08x\n", Bmask, qBmask, Boff);

	start = sdl->pixels;

	Uint32 t0 = SDL_GetTicks();

	if (!dodither) {

		for (y = 0; y < h; y++) {
			for (x = 0; x < w; x++) {
				// quantize current pixel
				pixel = *(Uint32 *)(start + (w*y+x));
				pixel &= (Amask | qRmask | qGmask | qBmask);
				*(Uint32 *)(start + (w*y+x)) = pixel;
			}
		}

	} else {

		for (y = 0; y < h; y++) {
			for (x = 0; x < w; x++) {
				// quantize current pixel
				pixel = *(Uint32 *)(start + (w*y+x));
				
				Rcomp = pixel & Rmask; Gcomp = pixel & Gmask; Bcomp = pixel & Bmask; Acomp = pixel & Amask;
				
				// apply offsets but clamp at max value
				Rcomp += Roff;
				Gcomp += Goff;
				Bcomp += Boff;
				if (Rcomp > Rmask) Rcomp = Rmask;
				if (Gcomp > Gmask) Gcomp = Gmask;
				if (Bcomp > Bmask) Bcomp = Bmask;
				
				// do the quantizing
				Rcomp = Rcomp & qRmask;
				Gcomp = Gcomp & qGmask;
				Bcomp = Bcomp & qBmask;
				
				// reassemble new pixel
				new = Rcomp | Gcomp | Bcomp | Acomp;
				*(Uint32 *)(start + (w*y+x)) = new;
				
				// calculate error
				Rerr = (pixel & Rmask) - (new & Rmask);
				Gerr = (pixel & Gmask) - (new & Gmask);
				Berr = (pixel & Bmask) - (new & Bmask);
				
				// propagate 7/16 * of error to next pixel to right
				if (x < w-1) {
					pixel = *(Uint32 *)(start + (w*y+x+1));
					fRerr = (Rerr * 7 / 16) & fRmask; fGerr = (Gerr * 7 / 16) & fGmask; fBerr = (Berr * 7 / 16) & fBmask;
					Rcomp = pixel & Rmask; Gcomp = pixel & Gmask; Bcomp = pixel & Bmask; Acomp = pixel & Amask;
					Rcomp = Rcomp + fRerr;  Gcomp = Gcomp + fGerr;  Bcomp = Bcomp + fBerr;
					if (Rcomp < 0) Rcomp = 0; if (Rcomp > Rmask) Rcomp = Rmask;
					if (Gcomp < 0) Gcomp = 0; if (Gcomp > Gmask) Gcomp = Gmask;
					if (Bcomp < 0) Bcomp = 0; if (Bcomp > Bmask) Bcomp = Bmask;
					*(Uint32 *)(start + (w*y+x+1)) = Rcomp | Gcomp | Bcomp | Acomp;
				}				
				// propagate 3/16 * of error to pixel on next row to left
				if (y < h-1 && x > 0) {
					pixel = *(Uint32 *)(start + (w*(y+1)+x-1));
					fRerr = (Rerr * 3 / 16) & fRmask; fGerr = (Gerr * 3 / 16) & fGmask; fBerr = (Berr * 3 / 16) & fBmask;
					Rcomp = pixel & Rmask; Gcomp = pixel & Gmask; Bcomp = pixel & Bmask; Acomp = pixel & Amask;
					Rcomp = Rcomp + fRerr;  Gcomp = Gcomp + fGerr;  Bcomp = Bcomp + fBerr;
					if (Rcomp < 0) Rcomp = 0; if (Rcomp > Rmask) Rcomp = Rmask;
					if (Gcomp < 0) Gcomp = 0; if (Gcomp > Gmask) Gcomp = Gmask;
					if (Bcomp < 0) Bcomp = 0; if (Bcomp > Bmask) Bcomp = Bmask;
					*(Uint32 *)(start + (w*(y+1)+x-1)) = Rcomp | Gcomp | Bcomp | Acomp;
				}				
				// propagate 5/16 * of error to pixel on next row
				if (y < h-1) {
					pixel = *(Uint32 *)(start + (w*(y+1)+x));
					fRerr = (Rerr * 5 / 16) & fRmask; fGerr = (Gerr * 5 / 16) & fGmask; fBerr = (Berr * 5 / 16) & fBmask;
					Rcomp = pixel & Rmask; Gcomp = pixel & Gmask; Bcomp = pixel & Bmask; Acomp = pixel & Amask;
					Rcomp = Rcomp + fRerr;  Gcomp = Gcomp + fGerr;  Bcomp = Bcomp + fBerr;
					if (Rcomp < 0) Rcomp = 0; if (Rcomp > Rmask) Rcomp = Rmask;
					if (Gcomp < 0) Gcomp = 0; if (Gcomp > Gmask) Gcomp = Gmask;
					if (Bcomp < 0) Bcomp = 0; if (Bcomp > Bmask) Bcomp = Bmask;
					*(Uint32 *)(start + (w*(y+1)+x)) = Rcomp | Gcomp | Bcomp | Acomp;
				}				
				// propagate 1/16 * of error to pixel on next row to right
				if (y < h-1 && x < w-1) {
					pixel = *(Uint32 *)(start + (w*(y+1)+x+1));
					fRerr = (Rerr * 1 / 16) & fRmask; fGerr = (Gerr * 1 / 16) & fGmask; fBerr = (Berr * 1 / 16) & fBmask;
					Rcomp = pixel & Rmask; Gcomp = pixel & Gmask; Bcomp = pixel & Bmask; Acomp = pixel & Amask;
					Rcomp = Rcomp + fRerr;  Gcomp = Gcomp + fGerr;  Bcomp = Bcomp + fBerr;
					if (Rcomp < 0) Rcomp = 0; if (Rcomp > Rmask) Rcomp = Rmask;
					if (Gcomp < 0) Gcomp = 0; if (Gcomp > Gmask) Gcomp = Gmask;
					if (Bcomp < 0) Bcomp = 0; if (Bcomp > Bmask) Bcomp = Bmask;
					*(Uint32 *)(start + (w*(y+1)+x+1)) = Rcomp | Gcomp | Bcomp | Acomp;
				}
			}
		}
	}

	Uint32 t1 = SDL_GetTicks();

	printf("quantizeWithDither took %dms\n", t1-t0);

	return 0;
}


static const struct luaL_Reg ditherlib [] = {
	{ "quantizeWithDither", jiveL_quantize_with_dither },
 	{ NULL, NULL }
};


int luaopen_jive_slim_dither (lua_State *L) {
	luaL_register(L, "dither", ditherlib);
	return 1;
}