/**
* @file			white_balance.c
* @brief		white balance algorithm
* @author		Patrick Roth - roth@stettbacher.ch
* @copyright	Stettbacher Signal Processing AG
* 
* @remarks
*
* <PRE>
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with this library; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
* </PRE>
*
*/

#include <stdio.h>
#include <math.h>
#include <stdlib.h>

#include "color_pipe_private.h"
#include "color.h"

/**
 * Minimum color gain
 */
#define GAIN_MIN		0.1

/**
 * Maximum color gain
 */
#define GAIN_MAX		10.0



/**
 * Calculate the average value of U and V component on given YUV image.
 * 
 * @param img_yuv YUV image buffer
 * @param height image height in number of pixels
 * @param width image width in number of pixels
 * @param threshold gray value threshold
 * @param u_avg On return: average value from all gray pixels (U-component)
 * @param v_avg On return: average value from all gray pixels (V-component)
 * @return -1 if no gray pixel is available otherwise 0
 */
static int get_u_v_avg(const int16_t *img_yuv, const int height, const int width, const float threshold, float *u_avg, float *v_avg) {
	int index, num_gray_pixel, y, x;
	int sum_u, sum_v;
	int pixel_y, pixel_u, pixel_v, color_val;
	const int shift_fact = 10;
	const int thresh = threshold * ((float)(1 << shift_fact));
	
	/*
	 * Accumulate all gray pixels below given threshold
	 */
	index = 0;
	num_gray_pixel = 0;
	sum_u = 0;
	sum_v = 0;
	for(y = 0; y < height; y++) {
		for(x = 0; x < width; x++) {
			
			pixel_y = img_yuv[index];
			if(pixel_y == 0) {
				index += 3;
				continue;
			}
			
			pixel_u = img_yuv[index+1];
			pixel_v = img_yuv[index+2];
			
			color_val = ((abs(pixel_u) + abs(pixel_v)) << shift_fact)/pixel_y;
			
			if(color_val < thresh) {
				// it's a gray pixel
				num_gray_pixel++;
				sum_u += pixel_u;
				sum_v += pixel_v;
			}
			index += 3;
		}
	}
	
	// check whether if there's any gray pixel
	if(num_gray_pixel == 0) {
		// no gray pixel --> nothing to do
		return -1;
	}
	
	// calculate avarage of U and V component over all gray pixels
	*u_avg = ((float)sum_u)/num_gray_pixel;
	*v_avg = ((float)sum_v)/num_gray_pixel;
	
// 	printf("XXX num_gray_pixel = %d, u_avg = %f, v_avg = %f\n", num_gray_pixel, *u_avg, *v_avg);
	
	return 0;
}


/**
 * Apply red and blue pixel gain on input image.
 * 
 * @param img_out On return: Gain is applied on this image
 * @param img_in input RGB image
 * @param height image height in number of pixels
 * @param width image width in number of pixels
 * @param bit_channel bits per color channel
 * @param gain_blue blue gain to apply
 * @param gain_red red gain to apply
 */
static void apply_pixel_gain_rgb(void *img_out, const void *img_in, const int height, const int width,
								 const int bit_channel, const float gain_blue, const float gain_red) {
	
	int index, x, y, pixel_red, pixel_blue;
	const int shift_fact = 8;
	const int blue = gain_blue * ((float)(1 << shift_fact));
	const int red = gain_red * ((float)(1 << shift_fact));
	const uint8_t *img8_in = img_in;
	const uint16_t *img16_in = img_in;
	uint8_t *img8_out = img_out;
	uint16_t *img16_out = img_out;
	const int max_pix = (1 << bit_channel) - 1;
	
	
	index = 0;
	pixel_red = 0;
	pixel_blue = 0;
	for(y = 0; y < height; y++) {
		for(x = 0; x < width; x++) {
			
			// apply gain on red pixel
			if(bit_channel <= 8) {
				pixel_red = (img8_in[index]*red) >> shift_fact;
			}
			else if(bit_channel <= 16) {
				pixel_red = (img16_in[index]*red) >> shift_fact;
			}
			if(pixel_red > max_pix) {
				pixel_red = max_pix;
			}
			else if (pixel_red < 0) {
				pixel_red = 0;
			}
			
			if(bit_channel <= 8) {
				img8_out[index] = pixel_red;
			}
			else if(bit_channel <= 16) {
				img16_out[index] = pixel_red;
			}
			index++;
			
			// no gain is applied on green pixel
			if(bit_channel <= 8) {
				img8_out[index] = img8_in[index];
			}
			else if(bit_channel <= 16) {
				img16_out[index] = img16_in[index];
			}
			index++;
			
			// apply gain on blue pixel
			if(bit_channel <= 8) {
				pixel_blue = (img8_in[index]*blue) >> shift_fact;
			}
			else if(bit_channel <= 16) {
				pixel_blue = (img16_in[index]*blue) >> shift_fact;
			}
			if(pixel_blue > max_pix) {
				pixel_blue = max_pix;
			}
			else if (pixel_blue < 0) {
				pixel_blue = 0;
			}
			
			if(bit_channel <= 8) {
				img8_out[index] = pixel_blue;
			}
			else if(bit_channel <= 16) {
				img16_out[index] = pixel_blue;
			}
			
			index++;
		}
	}
}


/**
 * Local white balance algorithm.
 * 
 * Reference:
 * J.-y Huo et al, "Robust Automatic White Balance Algorithm using for Gray Color Points in Images", IEEE Xplore, 2006
 * 
 * @param awb_data white balancing algorithm data
 * @return 0 on success otherwise -1
 */
int white_balance(struct awb_data_t *awb_data) {
	
	float u_avg, v_avg;
	float ctrl_k, gray_threshold;
	int bit_channel, height, width;
	void *img_balanced, *img_in;
	int16_t *img_yuv;
	int uv_div;
	
	
	// following parameter may be scaled
	ctrl_k = awb_data->ctrl_k;
	gray_threshold = awb_data->gray_threshold;
	bit_channel = awb_data->bit_channel;
	img_balanced = awb_data->img_rgb_balanced;
	img_in = awb_data->img_in;
	img_yuv = awb_data->img_yuv;
	height = awb_data->height;
	width = awb_data->width;
	
	
	// Apply red and blue gain on input image.
	apply_pixel_gain_rgb(img_balanced, img_in, height, width, bit_channel, awb_data->gain_blue, awb_data->gain_red);
	
	
	// RGB to YUV color space conversion
	if(color_rgb_to_yuv(img_yuv, img_balanced, height, width, bit_channel)) {
		// conversion failed --> do abort
		return -1;
	}
	
	if(get_u_v_avg(img_yuv, height, width, gray_threshold, &u_avg, &v_avg)) {
		// no gray pixel --> nothing to do
		return 0;
	}
	
	
	/*
	 * Norm U/V channel average value to 8 bit per color channel image.
	 * 
	 */
	uv_div = 1 << (bit_channel - 8);
	u_avg /= uv_div;
	v_avg /= uv_div;
	
	// adjust blue gain if image is bluish
	if(fabs(u_avg) > 0.4) {
		awb_data->gain_blue -= ctrl_k*u_avg;
	}
	
	// adjust red gain if image is reddish
	if(fabs(v_avg) > 0.4) {
		awb_data->gain_red -= ctrl_k*v_avg;
	}
	
	
	// range check of blue red
	if(awb_data->gain_blue < GAIN_MIN) {
		awb_data->gain_blue = GAIN_MIN;
	}
	else if(awb_data->gain_blue > GAIN_MAX) {
		awb_data->gain_blue = 3.0;
	}
	
	// range check of red gain
	if(awb_data->gain_red < GAIN_MIN) {
		awb_data->gain_red = GAIN_MIN;
	}
	else if(awb_data->gain_red > GAIN_MAX) {
		awb_data->gain_red = 3.0;
	}
	
// 	printf("XXX u_avg = %.4f, v_avg = %.4f, red = %.4f, blue = %.4f, bit_channel = %d\n", u_avg, v_avg, awb_data->gain_red, awb_data->gain_blue, bit_channel);
	return 0;
}