#include <v_lib.h>

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

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

u32_t	image_width;	/* currently may be equal to SCAN_WIDTH */
u32_t	image_height;	/* currently may be less than SCAN_WIDTH */
vec_t	gray_pixels[SCAN_WIDTH][SCAN_WIDTH];
vec_t	zero_data[SCAN_WIDTH];
vec_t	freq_real[SCAN_WIDTH][SCAN_WIDTH];
vec_t	freq_imag[SCAN_WIDTH][SCAN_WIDTH];
s32_t	quant = 128;

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

void calc_scale_and_offset(void)
{
	vec_t min = MAXFLOAT, max = -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 = 128;
		pix_scale = 0;
		return;
	}
	pix_scale = (vec_t)256.0 / (max - min);
	pix_offset =  - min * pix_scale;
}

#if 0

void freq_scale_and_offset(vec2_t scale_ret,vec2_t offset_ret)
{
	vec_t min = MAXFLOAT, max = -MAXFLOAT;
	u32_t count, line_num;
	for(line_num=0;line_num<image_height;line_num++)
	{
		vec_t *src = freq_real[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))
	{
		scale_ret[0] = 0;
		offset_ret[0] =  128;
	}
	else
	{
		scale_ret[0] = (vec_t)256.0 / (max - min);
		offset_ret[0] =  - min * scale_ret[0];
	}
	min = MAXFLOAT;
	max = -MAXFLOAT;
	for(line_num=0;line_num<image_height;line_num++)
	{
		vec_t *src = freq_imag[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))
	{
		scale_ret[1] = 0;
		offset_ret[1] =  128;
	}
	else
	{
		scale_ret[1] = (vec_t)256.0 / (max - min);
		offset_ret[1] =  - min * scale_ret[1];
	}
}

#endif

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;
	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] = ptr[channel];
	}
}

u8_t clamp(vec_t pixval)
{
	if(quantize)
	{
		if(pixval < quant)
			return 0;
		else
			return 255;
	}
	else
	{
		if(pixval < 0)
			return 0;
		if(pixval > 255)
			return 255;
		return (u8_t)(pixval + (vec_t)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);
	}
}

#if 0

void freq_to_scanlines(u8_t *pixels_real,u8_t *pixels_imag,u32_t channel,FILE *scalefile)
{
	u32_t count, line_num;
	u32_t stride = SCAN_WIDTH * 4;
	u8_t *ptr,*iptr,*endptr;
	vec_t *src,*isrc;
	vec2_t scale, offset;
	freq_scale_and_offset(scale,offset);
	fprintf(scalefile,"%g\t%g\t%g\t%g\n",
		scale[0],scale[1],offset[0],offset[1]);
	/*
	real_scales[channel] = 1 / scale[0];
	imag_scales[channel] = 1 / scale[1];
	real_offsets[channel] = offset[0];
	imag_offsets[channel] = offset[1];
	*/
	for(line_num=0;line_num<image_height;line_num++)
	{
		ptr = pixels_real + stride * line_num;
		iptr = pixels_imag + stride * line_num;
		endptr = ptr + stride;
		src = freq_real[line_num];
		isrc = freq_imag[line_num];
		for(count=0;ptr < endptr;ptr += 4,count++)
		{
			ptr[channel] = clamp(src[count] * scale[0] + offset[0]);
			iptr[channel] = clamp(isrc[count] * scale[1] + offset[1]);
		}
	}
}

#endif

void fft_forward(u32_t line_num)
{
	memZero(zero_data,SCAN_WIDTH * sizeof(vec_t));
	fftCalcNip(freq_real[line_num],freq_imag[line_num],
		gray_pixels[line_num],zero_data,SCAN_POWER,1);
}

void fft_inverse(u32_t line_num)
{
	fftCalcNip(gray_pixels[line_num],zero_data,
		freq_real[line_num],freq_imag[line_num],SCAN_POWER,-1);
}

void zero_freq(u32_t line_num,s32_t lo_pass,s32_t hi_pass)
{
	s32_t i, halfscan = SCAN_WIDTH >> 1;
	vec_t *rptr = freq_real[line_num];
	vec_t *iptr = freq_real[line_num];
	/*
	 * Discard frequencies below 'lo_pass'
	 */
	for(i=0;i<lo_pass;i++)
	{
		rptr[i] = 0;
		rptr[SCAN_WIDTH-i-1] = 0;
		iptr[i] = 0;
		iptr[SCAN_WIDTH-i-1] = 0;
	}
	/*
	 * Discard frequencies above 'hi_pass'
	 */
	for(i=hi_pass;i<halfscan;i++)
	{
		rptr[i] = 0;
		rptr[SCAN_WIDTH-i-1] = 0;
		iptr[i] = 0;
		iptr[SCAN_WIDTH-i-1] = 0;
	}
}

#if 0
u8_t *fpix_real, *fpix_imag;
void decode_partial();
#endif

void band_pass(u8_p pixels,s32_t lo_pass,s32_t hi_pass)
{
	u32_t i, channel;
#if 0
	FILE *scalefile;
	if(fpix_real && fpix_imag)
		scalefile=fopen("dft-scales.txt","wt");
	else
		scalefile=NULL;
#endif
	for(channel=0;channel<3;channel++)
	{
		scanlines_to_array(pixels,channel);
		for(i=0;i<image_height;i++)
		{
			fft_forward(i);
			zero_freq(i,lo_pass,hi_pass);
			fft_inverse(i);
		}
		array_to_scanlines(pixels,channel);
#if 0
		if(scalefile)
			freq_to_scanlines(fpix_real,fpix_imag,channel,scalefile);
#endif
	}
#if 0
	if(scalefile)
		fclose(scalefile);
#endif
}

#if 0
vec3_t real_scales, imag_scales;
vec3_t real_offsets, imag_offsets;

vboolean read_scale_file(void)
{
	u32_t channel;
	vhandle_t in = tokParseFile("dft-scales.txt","");
	vboolean ok = true;
	for(channel=0;channel<3;channel++)
	{
		ok = ok &&
			strToFloat(tokNext(in),&real_scales[channel]) &&
			strToFloat(tokNext(in),&imag_scales[channel]) &&
			strToFloat(tokNext(in),&real_offsets[channel]) &&
			strToFloat(tokNext(in),&imag_offsets[channel]);
	}
	tokDestroy(in);
	if(ok)
	{
		for(channel=0;channel<3;channel++)
		{
			real_scales[channel] = (vec_t)1.0 / real_scales[channel];
			imag_scales[channel] = (vec_t)1.0 / imag_scales[channel];
		}
	}
	for(channel=0;channel<3;channel++)
	{
		printf("channel %u:scales %g,%g; offsets %g,%g\n",
			channel,real_scales[channel],imag_scales[channel],
			real_offsets[channel],imag_offsets[channel]);
	}
	return ok;
}
	
void scanlines_to_freq(u8_t *pixels_real,u8_t *pixels_imag,u32_t channel)
{
	u32_t count, line_num;
	u32_t stride = SCAN_WIDTH * 4;
	u8_t *ptr,*iptr,*endptr;
	vec_t *dst,*idst;
	for(line_num=0;line_num<image_height;line_num++)
	{
		ptr = pixels_real + stride * line_num;
		iptr = pixels_imag + stride * line_num;
		endptr = ptr + stride;
		dst = freq_real[line_num];
		idst = freq_imag[line_num];
		for(count=0;ptr < endptr;ptr += 4,count++)
		{
			dst[count] = ((vec_t)(ptr[channel]) - real_offsets[channel]) * 300;
			idst[count] = ((vec_t)(iptr[channel]) - 128) * imag_scales[channel];
		}
	}
}

void decode(void)
{
	u8_p pixels;
	u32_t image_width;
	u32_t i, channel;
	if(!read_scale_file())
		return;
	fpix_real = bmpLoad("dft_real.bmp",&image_width,&image_height);
	if(!fpix_real)
		return;
	if(image_width != SCAN_WIDTH || image_height > SCAN_WIDTH)
	{
		memFree(fpix_real);
		return;
	}
	fpix_imag = bmpLoad("dft_imag.bmp",&image_width,&image_height);
	if(!fpix_imag)
	{
		memFree(fpix_real);
		return;
	}
	if(image_width != SCAN_WIDTH || image_height > SCAN_WIDTH)
	{
		memFree(fpix_real);
		memFree(fpix_imag);
		return;
	}
	pixels = (u8_p)memAlloc(image_width * image_height * 4);
	for(channel=0;channel<3;channel++)
	{
		scanlines_to_freq(fpix_real,fpix_imag,channel);
		for(i=0;i<image_height;i++)
			fft_inverse(i);
		compensate = true;
		quantize = false;
		array_to_scanlines(pixels,channel);
	}
	memFree(fpix_real);
	memFree(fpix_imag);
	bmpSave("decoded.bmp",image_width,image_height,pixels);
	memFree(pixels);
}

void decode_partial(void)
{
	u32_t channel;
	u8_p pixels = (u8_p)memAlloc(image_width * image_height * 4);
	for(channel=0;channel<3;channel++)
	{
		array_to_scanlines(pixels,channel);
	}
	bmpSave("decoded__.bmp",image_width,image_height,pixels);
	memFree(pixels);
}

#endif

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

void usage(char *name)
{
	conPrintf(
		"Band-pass image filtering by FDFT/IDFT. Version 0.002 by Vladimir A. Smirnov.\n\n"
		"Usage: %s <input file name> <low freq bound> <high freq bound> [output file name] [outmethod], where\n"
		"\t<input file name> - name of image file to process (with BMP, TGA or JPG extension);\n"
		"\t<low freq bound>, <high freq bound> - bounds in wave-number domain;\n"
		"\t[output file name] - optional name of output file (default is 'out.bmp').\n"
		"Width of the input bitmap must be %d, height must not be greater than %d.\n"
		"Bounds in wave-number space must be in range [0..%d].\n"
		"Last parameter must be either 'clamp' (save result of IDFT without modifications),\n"
		"or 'compensate' (increase contrast such that luma occupy all range from 0 to 255),\n"
		"or 'quant' (same as compensate, than produce mono output); default is 'clamp'.\n"
		"If 'outmethod' is 'quant', it's also possible to specify threshold (ex.: quant=77).\n"
		"KNOWN BUGS:\n"
		"1. Should perform 2D FDFT/IDFT instead of 1D.\n"
		"2. Currently general-case complex DFT is used for real-only data.\n"
		"3. Probably, other.\n",name,SCAN_WIDTH,SCAN_WIDTH,SCAN_WIDTH>>1);
}

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;
}

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);
	}
}

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);
		}
		else
			return false;
	}
	return true;
}

vboolean init(int argc,char *argv[],s32_t *l_ret,s32_t *h_ret,char **outname_ret)
{
	s32_t fmax = SCAN_WIDTH >> 1;
	if(argc < 4 || argc > 6)
	{
		usage(argv[0]);
		return false;
	}
	if(argc == 4) {
		*outname_ret = NULL;
		compensate = false;
	} else if(argc == 5) {
		if(!set_clamp(argv[4]))
			*outname_ret = argv[4];
		else
			*outname_ret = NULL;
	}
	else /* argc == 6 */
	{
		u32_t no = 4;
		if(set_clamp(argv[4]))
			no = 5;
		else
		{
			if(!set_clamp(argv[5]))
			{
				conlPrintf(LOG_ERROR,
					"init(): either 'compensate' or 'clamp' got '%s'.\n"
					"For help, run %s without any arguments.\n",
					argv[5],argv[0]);
				return false;
			}
		}
		*outname_ret = argv[no];
	}
	if(!strToInt(argv[2],l_ret) || !strToInt(argv[3],h_ret))
	{
		conlPrintf(LOG_ERROR,
			"init(): expected two integers, got '%s' and '%s'.\n"
			"For help, run %s without any arguments.\n",
			argv[2],argv[3],argv[0]);
		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;
	}
	return true;
}

int main(int argc,char *argv[])
{
	s32_t lo_pass,hi_pass;
	u8_t *pixels;
	char *outname;
	logStartConsoleLog();
	logUpperLevel(LOG_TRACE);
#if 0
	if(argc == 2 && !strcmp(argv[1],"decode"))
	{
		decode();
		return 0;
	}
#endif
	if(!init(argc,argv,&lo_pass,&hi_pass,&outname))
		return -1;
	if(!outname)
		outname = "out.bmp";
	pixels = LoadImage(argv[1],&image_width,&image_height);
	if(!pixels)
		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);
		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);
		return -1;
	}
#if 0
	fpix_real = (u8_t*)memAlloc(image_width * image_height * 4);
	fpix_imag = (u8_t*)memAlloc(image_width * image_height * 4);
#endif
	band_pass(pixels,lo_pass,hi_pass);
	SaveImage(outname,image_width,image_height,pixels);
#if 0
	SaveImage("dft_real.bmp",image_width,image_height,fpix_real);
	SaveImage("dft_imag.bmp",image_width,image_height,fpix_imag);
	memFree(fpix_real);
	memFree(fpix_imag);
#endif
	memFree(pixels);
	return 0;
}
