#include <v_lib.h>

/* ==================================================================== */

#define	SCAN_POWER	9
#define	SCAN_WIDTH	((u32_t)1 << SCAN_POWER)

u32_t	image_width;	/* currently must be equal to SCAN_WIDTH */
u32_t	image_height;	/* currently must not be greater than SCAN_WIDTH */
vec_t	gray_pixels[SCAN_WIDTH][SCAN_WIDTH];
vec_t	*pix_rows[SCAN_WIDTH];
matrix_t *mask = NULL;

matrix_t pixmat =
{
	SCAN_WIDTH, SCAN_WIDTH,
	pix_rows,
	(vec_t*)gray_pixels
};

void init_pixmat()
{
	int i;
	pixmat.cell[0] = pixmat.data;
	for(i=1;i<SCAN_WIDTH;i++)
		pixmat.cell[i] = pixmat.cell[i-1] + SCAN_WIDTH;
}
		
vec_t	freq_ccg[SCAN_WIDTH];
s32_t	quant = 128;
vec_t	qfloat;

vboolean space_hsv = false;
vboolean desaturate = false;
vboolean compensate = false;
vboolean quantize = false;
vec_t	pix_scale, pix_offset;

void calc_scale_and_offset(void)
{
	vec_t min = (vec_t)MAXFLOAT, max = -(vec_t)MAXFLOAT;
	u32_t count, line_num;
	for(line_num=0;line_num<image_height;line_num++)
	{
		vec_t *src = gray_pixels[line_num];
		for(count=0;count < SCAN_WIDTH;count++)
		{
			if(src[count] < min)
				min = src[count];
			if(src[count] > max)
				max = src[count];
		}
	}
	if(IsZero(max-min))
	{
		conlPuts(LOG_WARNING,
			"WARNING: IDFT result with zero contrast, luma compensation disabled\n");
		pix_offset = (vec_t)0.5;
		pix_scale = 0;
		return;
	}
	pix_scale = (vec_t)1.0 / (max - min);
	pix_offset =  - min * pix_scale;
}

void scanlines_to_array(u8_t *pixels,u32_t channel)
{
	u32_t count, line_num;
	u32_t stride = SCAN_WIDTH * 4;
	u8_t *ptr, *endptr;
	vec_t *dst, scale = (vec_t)(1.0 / 255.0);
	for(line_num=0;line_num<image_height;line_num++)
	{
		ptr = pixels + stride * line_num;
		endptr = ptr + stride;
		dst = gray_pixels[line_num];
		for(count=0;ptr < endptr;ptr += 4,count++)
			dst[count] = scale * ptr[channel];
	}
}

u8_t clamp(vec_t pixval)
{
	if(quantize)
	{
		if(pixval < qfloat)
			return 0;
		else
			return 255;
	}
	else
	{
		if(pixval < 0)
			return 0;
		if(pixval >= 1)
			return 255;
		return (u8_t)(pixval * 255.0 + 0.5);
	}
}

void array_to_scanlines(u8_t *pixels,u32_t channel)
{
	u32_t count, line_num;
	u32_t stride = SCAN_WIDTH * 4;
	u8_t *ptr, *endptr;
	vec_t *src;
	if(compensate)
		calc_scale_and_offset();
	else
	{
		pix_scale = 1;
		pix_offset = 0;
	}
	for(line_num=0;line_num<image_height;line_num++)
	{
		ptr = pixels + stride * line_num;
		endptr = ptr + stride;
		src = gray_pixels[line_num];
		for(count=0;ptr < endptr;ptr += 4,count++)
			ptr[channel] = clamp(pix_scale * src[count] + pix_offset);
	}
}

void zero_freq(vec_t *afreq,s32_t lo_pass,s32_t hi_pass)
{
	u32_t hlen;
	/*
	 * Discard frequencies below 'lo_pass'
	 */
	if(lo_pass)
		memZero(afreq,lo_pass * 2 * sizeof(vec_t));
	/*
	 * Discard frequencies above 'hi_pass'
	 */
	hi_pass <<= 1;
	hlen = SCAN_WIDTH - hi_pass;
	if(hlen)
		memZero(afreq + hi_pass,hlen * sizeof(vec_t));
}

void zero2d(matrix_t *m,s32_t lo_pass,s32_t hi_pass)
{
	u32_t i;
	matrix_t *m2 = matAlloc(m->rows,m->cols);
	for(i=0;i<image_height;i++)
		zero_freq(m->cell[i],lo_pass,hi_pass);
	matTranspose(m2,m);
	for(i=0;i<image_height;i++)
		zero_freq(m2->cell[i],lo_pass,hi_pass);
	matTranspose(m,m2);
	memFree(m2);
}

vhandle_t fft_context = LIBV_INVALID_HANDLE;

void band_pass(u8_p pixels,s32_t lo_pass,s32_t hi_pass)
{
	u32_t channel;
	for(channel=0;channel<3;channel++)
	{
		scanlines_to_array(pixels,channel);
		fftForward2D(fft_context,&pixmat);
		zero2d(&pixmat,lo_pass,hi_pass);
		fftInverse2D(fft_context,&pixmat);
		array_to_scanlines(pixels,channel);
	}
}

void masked_zero(matrix_t *m,matrix_t *mask)
{
	u32_t i,j;
	for(i=0;i<m->rows;i++)
		for(j=0;j<m->cols;j++)
			m->cell[i][j] *= mask->cell[i][j];
}

void mask_pass(u8_p pixels,matrix_t *mask)
{
	u32_t channel;
	for(channel=0;channel<3;channel++)
	{
		scanlines_to_array(pixels,channel);
		fftForward2D(fft_context,&pixmat);
		masked_zero(&pixmat,mask);
		fftInverse2D(fft_context,&pixmat);
		array_to_scanlines(pixels,channel);
	}
}

/* ==================================================================== */

void usage(char *name)
{
	puts("Band-pass image filtering by FDFT/IDFT. Version 0.2 by Vladimir A. Smirnov.\n");
	printf("Usage: %s -infile <infile> [-space-hsv -desaturate]"
		"(-bounds <low bound> <high bound> | -mask mask <maskfile>) "
		"[-outfile outfile] [-outmethod <outmethod>], where:\n",name);
	puts("\t'infile' - name of image file to process (BMP, TGA or JPG);");
	puts("\t'low bound', 'high bound' - bounds in wave-number domain;");
	puts("\t'outfile' - optional name of output file (default is 'out.bmp').");
	printf("Width of the input bitmap must be %u, height must not be greater than %u.\n",SCAN_WIDTH,SCAN_WIDTH);
	printf("Bounds in wave-number space must be in range [0..%u].\n",SCAN_WIDTH>>1);
	puts("'outmethod' must be either 'clamp' (clamp results of IDFT in range 0..255),");
	puts("or 'compensate' (increase contrast such that luma occupy range from 0 to 255),");
	puts("or 'quant' (same as compensate, but produce mono output); default is 'clamp'.");
	puts("If 'outmethod' is 'quant', it's possible to specify threshold (ex.: quant=77).");
	puts("KNOWN BUGS: Probably, a lot.");
}

u8_p LoadImage(char *name,u32_p w_ret,u32_p h_ret)
{
	char ext[_MAX_PATH];
	fsSplitPath(name,NULL,NULL,ext);
	if(!strcasecmp(ext,"bmp"))
		return bmpLoad(name,w_ret,h_ret);
	if(!strcasecmp(ext,"tga"))
		return tgaLoad(name,w_ret,h_ret);
	if(!strcasecmp(ext,"jpg"))
		return jpgLoad(name,w_ret,h_ret);
	conlPrintf(LOG_ERROR,"LoadImage(%s): unrecognized format.\n",name);
	return NULL;
}

#define	GET_PEL(vdest,pbuf,w,x,y) { \
	u8_t *ptr = pbuf + ((w * y + x) << 2); \
	rgbFromByte(vdest,ptr); }

#define	SET_PEL(pbuf,w,x,y,vsrc) { \
	u8_t *ptr = pbuf + ((w * y + x) << 2); \
	rgbByteFromFloat(ptr,vsrc); }

void rgb2hsv_pixbuf(u8_t *pixbuf,u32_t w,u32_t h)
{
	vec3_t rgb, hsv;
	u32_t i, j;
	for(i=0;i<h;i++)
		for(j=0;j<w;j++)
		{
			GET_PEL(rgb,pixbuf,w,j,i);
			hsvFromRgb(hsv,rgb);
			hsv[0] *= (vec_t)(1.0 / 360.0);
			SET_PEL(pixbuf,w,j,i,hsv);
		}
}

void hsv2rgb_pixbuf(u8_t *pixbuf,u32_t w,u32_t h)
{
	vec3_t rgb, hsv;
	u32_t i, j;
	vec_t desaturate_factor = (desaturate) ? 0 : 1;
	for(i=0;i<h;i++)
		for(j=0;j<w;j++)
		{
			GET_PEL(hsv,pixbuf,w,j,i);
			hsv[0] *= (vec_t)(360.0);
			hsv[1] *= desaturate_factor;
			rgbFromHsv(rgb,hsv);
			SET_PEL(pixbuf,w,j,i,rgb);
		}
}

u8_p LoadImageEx(char *name,u32_p w_ret,u32_p h_ret,vboolean convert_to_hsv)
{
	u8_p ret = LoadImage(name,w_ret,h_ret);
	V_ASSERT(w_ret && h_ret);
	if(convert_to_hsv)
		rgb2hsv_pixbuf(ret,*w_ret,*h_ret);
	return ret;
}

matrix_p LoadMask(char *name)
{
	u8_p ptr, pixbuf;
	matrix_p mask;
	u32_t w, h, i, j;
	pixbuf = LoadImage(name,&w, &h);
	if(!pixbuf)
		return NULL;
	if(w != SCAN_WIDTH || h != SCAN_WIDTH)
	{
		memFree(pixbuf);
		return NULL;
	}
	mask = matAlloc(w,h);
	ptr = pixbuf;	
	for(i=0;i<mask->rows;i++)
	{
		for(j=0;j<mask->cols;j++,ptr+=4)
			mask->cell[i][j] = (vec_t)((u32_t)ptr[0]+(u32_t)ptr[1]+(u32_t)ptr[2]) / (vec_t)(0xFF * 3);
	}
	return mask;
}

void SaveImage(char *name,u32_t w,u32_t h,u8_p pixels)
{
	char ext[_MAX_PATH];
	fsSplitPath(name,NULL,NULL,ext);
	if(!strcasecmp(ext,"bmp"))
		bmpSave(name,w,h,pixels);
	else if(!strcasecmp(ext,"tga"))
		tgaSave(name,w,h,pixels);
	else
	{
		conlPrintf(LOG_WARNING,"SaveImage(%s): unrecognized format, assuming TARGA.\n",name);
		tgaSave(name,w,h,pixels);
	}
}

void SaveImageEx(char *name,u32_t w,u32_t h,u8_p pixels,vboolean convert_from_hsv)
{
	if(convert_from_hsv)
		hsv2rgb_pixbuf(pixels,w,h);
	SaveImage(name,w,h,pixels);
}

vboolean set_clamp(char *arg)
{
	if(!strcmp(arg,"compensate"))
		compensate = true;
	else if(!strcmp(arg,"clamp"))
		compensate = false;
	else
	{
		char buf[64];
		memCopy(buf,arg,5);
		buf[5]=0;
		if(!strcmp(buf,"quant"))
		{
			compensate = true;
			quantize = true;
			if(arg[5] == '=')
			{
				strToInt(arg+6,&quant);
				qfloat = (vec_t)(1.0/255.0) * quant;
			}
		}
		else
			return false;
	}
	return true;
}

vboolean
init(int argc,char *argv[],
	s32_t *l_ret,s32_t *h_ret,char **inname_ret,char **outname_ret)
{
	int i;
	s32_t fmax = SCAN_WIDTH >> 1;
	init_pixmat();
	*outname_ret = NULL;
	*inname_ret = NULL;
	for(i=1;i<argc;)
	{
		char *token = argv[i++];
		if(!strcasecmp(token,"-infile"))
		{
			if(i >= argc)
				return false;
			*inname_ret = argv[i++];
			continue;
		} else if(!strcasecmp(token,"-outfile"))
		{
			if(i >= argc)
				return false;
			*outname_ret = argv[i++];
			continue;
		} else if(!strcasecmp(token,"-bounds"))
		{
			if((i+1) >= argc)
				return false;
    		if(!strToInt(argv[i++],l_ret) || !strToInt(argv[i++],h_ret))
	    	{
    			conlPuts(LOG_ERROR,
    				"init(): expected two integers after -bounds.\n");
    			return false;
    		}
    		if(*l_ret < 0 || *l_ret > fmax || *h_ret < 0 || *h_ret > fmax)
	    	{
    			conlPrintf(LOG_ERROR,
    				"init(): wrong frequency bounds '%d' and '%d'.\n"
    				"For help, run %s without any arguments.\n",
    				*l_ret,*h_ret,argv[0]);
    			return false;
    		}
			continue;
		} else if(!strcasecmp(token,"-mask"))
		{
			if(i >= argc)
				return false;
	    	mask = LoadMask(argv[i++]);
    		if(mask == NULL)
    			return false;
		} else if(!strcasecmp(token,"-outmethod"))
		{
			if(i >= argc)
				return false;
			if(!set_clamp(argv[i++]))
			{
				conlPuts(LOG_ERROR,"init(): wrong outmethod.\n");
				return false;
			}
			continue;
		} else if(!strcasecmp(token,"-space-hsv"))
		{
			space_hsv = true;
			continue;
		} else if(!strcasecmp(token,"-desaturate"))
		{
			desaturate = true;
			continue;
		}
	}
	fft_context = fftCreateContext(SCAN_POWER,true);
	return true;
}

int main(int argc,char *argv[])
{
	s32_t lo_pass,hi_pass;
	u8_t *pixels;
	char *inname, *outname;
	logStartConsoleLog();
	logUpperLevel(LOG_TRACE);
	if(!init(argc,argv,&lo_pass,&hi_pass,&inname,&outname) ||
		!inname)
	{
		usage(argv[0]);
		return -1;
	}
	if(!outname)
		outname = "out.bmp";
	pixels = LoadImageEx(inname,&image_width,&image_height,space_hsv);
	if(!pixels)
	{
		if(mask)
			memFree(mask);
		return -1;
	}
	if(image_width != SCAN_WIDTH)
	{
		conlPrintf(LOG_ERROR,"Bitmap width must be equal to %u, got %u\n",SCAN_WIDTH,image_width);
		memFree(pixels);
		if(mask)
			memFree(mask);
		return -1;
	}
	if(image_height > SCAN_WIDTH)
	{
		conlPrintf(LOG_ERROR,"Bitmap height must be less than %u, got %u\n",SCAN_WIDTH,image_height);
		memFree(pixels);
		if(mask)
			memFree(mask);
		return -1;
	}
	if(mask)
		mask_pass(pixels,mask);
	else
		band_pass(pixels,lo_pass,hi_pass);
	SaveImageEx(outname,image_width,image_height,pixels,space_hsv);
	memFree(pixels);
	if(mask)
		memFree(mask);
	if(fft_context != LIBV_INVALID_HANDLE)
		fftDestroyContext(fft_context);
	return 0;
}
