#include <v_lib.h>

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

#define	SCAN_POWER	8
#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	zero_data[SCAN_WIDTH];
vec_t	freq_real[SCAN_WIDTH];
vec_t	freq_imag[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;
}

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

void fft_forward(u32_t line_num)
{
	memZero(zero_data,SCAN_WIDTH * sizeof(vec_t));
	fftCalcNip(freq_real,freq_imag,
		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,freq_imag,SCAN_POWER,-1);
}

void zero_freq(s32_t lo_pass,s32_t hi_pass)
{
	s32_t i, halfscan = SCAN_WIDTH >> 1;
	vec_t *rptr = freq_real;
	vec_t *iptr = freq_imag;
	/*
	 * 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;
	}
}

void band_pass(u8_p pixels,s32_t lo_pass,s32_t hi_pass)
{
	u32_t i, channel;
	for(channel=0;channel<3;channel++)
	{
		scanlines_to_array(pixels,channel);
		for(i=0;i<image_height;i++)
		{
			fft_forward(i);
			zero_freq(lo_pass,hi_pass);
			fft_inverse(i);
		}
		array_to_scanlines(pixels,channel);
	}
}

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

void usage(char *name)
{
	puts("Band-pass image filtering by FDFT/IDFT. Version 0.002 by Vladimir A. Smirnov.\n");
	printf("Usage: %s <infile> <low bound> <high bound> [outfile] [outmethod], where:\n",name);
	puts("\t<infile> - name of image file to process (BMP, TGA or JPG);");
	puts("\t<low freq bound>, <high freq 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("Last parameter 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, than produce mono output); default is 'clamp'.");
	puts("If 'outmethod' is 'quant', it's possible to specify threshold (ex.: quant=77).");
	puts("KNOWN BUGS:");
	puts("1. Should perform 2D FDFT/IDFT instead of 1D.");
	puts("2. Currently general-case complex DFT is used for real-only data.");
	puts("3. Probably, other.");
}

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(!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;
	}
	band_pass(pixels,lo_pass,hi_pass);
	SaveImage(outname,image_width,image_height,pixels);
	memFree(pixels);
	return 0;
}
