#include "common.h"
#include "night_vision.h"

float4 blur_params;

float4 main(p_screen I) : SV_Target
{
	// Defines
	float3 image;
	const float GAUSS_SIGMA = 2.5;
	const int support = int(GAUSS_SIGMA * 1.5);
	float lua_param_nvg_generation = floor(shader_param_8.x);             // 1, 2, or 3
	float lua_param_nvg_num_tubes = frac(shader_param_8.x) * 10.0f;  
	float lua_param_nvg_washout_thresh = frac(shader_param_8.y);
	float lua_param_vignette_current = floor(shader_param_8.z) / 100.0f;

	if (lua_param_nvg_generation > 0.0f && compute_lens_mask(aspect_ratio_correction(I.tc0), lua_param_nvg_num_tubes) == 1.0f) // if NVGs are enabled ... 
	{
	
		// HALF RES BLUR - PASS 1 - R = LUMA, G = LIGHTMAP, B = NORMALIZED DEPTH OUT TO FARTHEST_DEPTH (defined in night_vision.h)
		if (blur_params.x == 1 && blur_params.z == (screen_res.x / 2.0f))
		{
			image = float3(s_image.Sample(smp_rtlinear, I.tc0).rg,saturate(s_position.Load( int3( (I.tc0)*screen_res.xy, 0 ), 0 ).z / farthest_depth));
			if (image.b == 0.0f) { image.b = 1.0f; } float scale = 0.05f; scale *= (IGN_calc(I.tc0/2.0*screen_res.xy)); float3 total = 0.0f; float divisor = 0.0f;
			float total_depth = image.b; float divisor_depth = 1.0f; float r = 1.0f; float2x2 G = rot(2.399996); float2 offset = float2(scale,scale);
				for (int i = 0; i < 16; ++i) {
					r += 1.0f / r; offset = mul(offset, G);	float3 color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (offset * (r - 1.0f) * 1.0f/screen_res.xy),0).rgb;
					color.b =  saturate(s_position.Load( int3( (I.tc0*screen_res.xy) + (offset * (r - 1.0f)), 0 ), 0 ).z / (farthest_depth));
					if (color.b == 0.0f) { color.b = 1.0f; } float weight = 1.0f; total.rg += color.rg * weight; divisor += weight; if (color.b <= image.b - 0.04) 
					{ total_depth += color.b; divisor_depth += weight; }					
				}
			image = float3(max((total.rg / divisor),0),max((total_depth / divisor_depth),0)); 	
			float vignette = calc_vignette(lua_param_nvg_num_tubes, I.tc0, lua_param_vignette_current); image = clamp(image,0.0,1.0); image *= vignette;
			return float4(image, 1.0f);
		}
		
		// HALF RES BLUR - PASS 2 - R = LUMA, G = LIGHTMAP, B = BLURRED DEPTH
		else if (blur_params.x == 0 && blur_params.z == (screen_res.x / 2.0f)) // half-res fixed DOF blur pass 2
		{
			image = s_image.Sample(smp_rtlinear, I.tc0).rgb; float scale = 0.3f; float center_depth = image.b; scale *= (IGN_calc(I.tc0/2.0*screen_res.xy));
			float3 total = 0.0f; float divisor = 0.0f; float r = 1.0f; float2x2 G = rot(2.399996); float2 offset = float2(scale,scale);
			for (int i = 0; i < 24; ++i) { r += 1.0f / r; offset = mul(offset, G); 
				float falloff = saturate( ( 1.2f * center_depth ) - log(center_depth * 6.4f) - ( 0.4f * pow(center_depth, 2.3f) ) );
				float3 color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (offset * (falloff) * (r - 1.02f) * 1.0f/screen_res.xy),0).rgb;
				float weight = 1.0f; total += color * weight; divisor += weight;
			}
			image = max((total.rgb / divisor),0);
			float vignette = calc_vignette(lua_param_nvg_num_tubes, I.tc0, lua_param_vignette_current); image = clamp(image,0.0,1.0); image *= vignette;
			return float4(image, 1.0f);
		}
		
		// QUARTER RES BLUR - PASS 1 - R = LUMA, G = LIGHTMAP, B = ALBEDO
		else if (blur_params.x == 1 && blur_params.z == (screen_res.x / 4))
		{
			float general_threshold = 0.05f; float light_thresh = lua_param_nvg_washout_thresh; float Pi = 6.28318530718; float blur_directions = 12.0; float blur_quality = 6.0; 
			float blur_radius = 1.5; float3 color = 0.0f; float3 color_average = s_image.Sample(smp_rtlinear, I.tc0).rgb; float weight = 0.0f;
			float light_avgs = 0.0f; 
			for(float i=1.0; i<= blur_quality; i++) { for( float d=0.0; d<Pi; d+=Pi/blur_directions) {
				color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (float2(cos(d),sin(d))*blur_radius*i/screen_res.xy),0).rgb; color.rb *= color.rb;
				color.rb = saturate((color.rb - general_threshold) / (1.0f - general_threshold)); color.g = saturate((color.g - light_thresh) / (1.0f - light_thresh));
				if (color.g > 0.0f) { color.g *= 8.0f; light_avgs += 1.0f;} else {light_avgs += 0.1 * i;} 
				color.rb = sqrt(color.rb); color_average += color;  } }
			image.rb = color_average.rb / (blur_directions*blur_quality); image.g = color_average.g / light_avgs;
			return float4(image, 1.0f);
		}
		
		// QUARTER RES BLUR - PASS 2 - R = LUMA, G = LIGHTMAP, B = ALBEDO
		else if (blur_params.x == 0 && blur_params.z == (screen_res.x / 4)) // quarter and eighth res circular bokeh pass 2
		{
			float Pi = 6.28318530718; float blur_directions = 16.0; float blur_quality = 4.0; float blur_radius = 4; float3 color = 0.0f;
			float3 color_average = s_image.Sample(smp_rtlinear, I.tc0).rgb;	float light_avgs = 0.0f; float weight = 0.0f; float total_weight = 0.0f;
			for(float i=1.0; i<=	blur_quality; i++) { for( float d=0.0; d<Pi; d+=Pi/blur_directions) {
				color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (float2(cos(d),sin(d))*blur_radius*i/screen_res.xy),0).rgb;
				if (color.g > 0.0f) { color.g *= 4.0f; light_avgs += 1.0f;}	else {light_avgs += 0.1 * i;} color_average += color; }	}
			image.rb = color_average.rb / (blur_directions*blur_quality); image.g = color_average.g / light_avgs;
			return float4(image, 1.0f);
		}
		
		// EIGHTH RES BLUR - PASS 1 - R = LUMA, G = LIGHTMAP, B = ALBEDO
		else if (blur_params.x == 1 && blur_params.z == (screen_res.x / 8)) // quarter and eighth res circular bokeh pass 1
		{
		float general_threshold = 0.05f; float light_thresh = lua_param_nvg_washout_thresh;	float Pi = 6.28318530718; float blur_directions = 12.0; float blur_quality = 4.0;
		float blur_radius = 1; float3 color = 0.0f; float3 color_average = s_image.Sample(smp_rtlinear, I.tc0).rgb; float weight = 0.0f; float total_weight = 0.0f;
		float light_avgs = 0.0f;
		for(float i=1.0; i<= blur_quality; i++) { for( float d=0.0; d<Pi; d+=Pi/blur_directions) { 
			color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (float2(cos(d),sin(d))*blur_radius*i/screen_res.xy),0).rgb; color.rb *= color.rb;
			color.rb = saturate((color.rb - general_threshold) / (1.0f - general_threshold)); color.g = saturate((color.g - light_thresh) / (1.0f - light_thresh));
			if (color.g > 0.0f) { color.g *= 8.0f; light_avgs += 1.0f;} else {light_avgs += 0.1 * i;} color.rb = sqrt(color.rb); color_average += color; } }
		image.rb = color_average.rb / (blur_directions*blur_quality); image.g = color_average.g / light_avgs;
		return float4(image, 1.0f);
		}
		
		// EIGHTH RES BLUR - PASS 2 - R = LUMA, G = LIGHTMAP, B = ALBEDO
		else if (blur_params.x == 0 && blur_params.z == (screen_res.x / 8)) // quarter and eighth res circular bokeh pass 2
		{
		float Pi = 6.28318530718; float blur_directions = 16.0;	float blur_quality = 4.0; float blur_radius = 4; float3 color = 0.0f;
		float3 color_average = s_image.Sample(smp_rtlinear, I.tc0).rgb;	float light_avgs = 0.0f; float weight = 0.0f; float total_weight = 0.0f;
		for(float i=1.0; i<=	blur_quality; i++) { for( float d=0.0; d<Pi; d+=Pi/blur_directions) {
			color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (float2(cos(d),sin(d))*blur_radius*i/screen_res.xy),0).rgb;
			if (color.g > 0.0f) { color.g *= 4.0f; light_avgs += 1.0f;} else {light_avgs += 0.1 * i;} color_average += color; } }
		image.rb = color_average.rb / (blur_directions*blur_quality); image.g = color_average.g / light_avgs;
		return float4(image, 1.0f);
		}
	}
	
	// OTHERWISE, DO NORMAL BLUR
	float scale = 0.3f;	scale *= (IGN_calc(I.tc0/2.0*screen_res.xy)); float3 total =  s_image.Sample(smp_rtlinear, I.tc0).rgb; float divisor = 1.0f;
	float r = 0.9f; float2x2 G = rot(2.399996);	float2 offset = float2(scale,scale);
	for (int i = 0; i < 16; ++i) {r += 1.0f / r; offset = mul(offset, G); 
		float3 color = s_image.SampleLevel(smp_rtlinear, I.tc0 + (offset * (r - 1.0f) * 1.0f/screen_res.xy),0).rgb;	float weight = 1.0f; 
		total += color * weight; divisor += weight; } image = max((total / divisor),0);
	return float4(image, 1.0f);
}