in vec2 fragmentUv01;
in vec3 fragmentNormal;
in vec4 fragmentPosition;

#include "drive/water/water.shared.glsl"


void main()
{
	// 0 = default
	// 1 = height
	// 2 = normal
	// 3 = baselayer
	// 4 = reflection 
	// 5 = bottomlayer
	// 6 = caustics

	uint stage = uint(u_stage);
	uint substage = uint(u_substage);
    float time = u_time;

	vec3 sunDir = -normalize(vec3(1.0, 1.0, 1.0));

	float iblStrength = u_iblStrength;
	float sunStrength = u_sunStrength;

	out_color0 = vec4(0.0, 0.0, 0.0, 1.0);

	float waterHeight = sampleWaterHeight(fragmentPosition.xz, time);
	vec3 waterNormal = sampleWaterNormal(fragmentPosition.xz, time);
	vec3 waterViewNormal = (u_view * vec4(waterNormal, 0.0)).xyz;

	vec2 gridPosition = fragmentPosition.xz * 10.0;
	vec2 gridIndex = floor(gridPosition);
	vec2 gridFraction = gridPosition - gridIndex;
	float gridWidth = 0.025;
	float gridDistA = step(gridFraction.x, 1.0 - gridWidth);
	float gridDistB = step(gridFraction.y, 1.0 - gridWidth);
	float gridDistC = step(gridWidth, gridFraction.x);
	float gridDistD = step(gridWidth, gridFraction.y);
	float gridDist = 1.0 - (gridDistA * gridDistB * gridDistC * gridDistD);

	// HEIGHT
	if(stage == uint(1))
	{
		out_color0.rgb = vec3(waterHeight);
		out_color0.rgb *= iblStrength;
		return;
	}


	// NORMAL
	if(stage == uint(2))
	{
		out_color0.rgb = waterNormal * vec3(1.0, 1.0,1.0);
		return;
	}

	vec2 screenUv = getScreenUv(gl_FragCoord.xy);
	vec3 frontPosition_OG = texture(u_frontPosition, screenUv).xyz;
	vec3 frontColor_OG =  texture(u_frontColor, screenUv).xyz;
	float depth_OG = distance(frontPosition_OG, fragmentPosition.xyz);

	float depthScale_OG = smoothstep(0.0, 1.0,  depth_OG);
	vec2 pixelScale = vec2(1.0) / u_resolution;
	vec2 off_OG = waterViewNormal.xz * depthScale_OG * pixelScale * u_refractionStrength;

	vec3 frontColor = texture(u_frontColor, screenUv + off_OG).rgb;
	vec3 frontPosition = texture(u_frontPosition, screenUv + off_OG).xyz;

	float causticsCoeff = sampleWaterHeight(frontPosition.xz, time);
	causticsCoeff = clamp( smoothstep(0.0, 0.5, causticsCoeff), 0.0, 1.0);

	vec3 lightToSurface = sunDir;
	vec3 ndotl = vec3(clamp(dot(waterNormal, lightToSurface), 0.0, 1.0));
	
	vec3 viewToSurface = normalize(fragmentPosition.xyz - (inverse(u_view) * vec4(0.0, 0.0, 0.0, 1.0)).xyz);
	viewToSurface.y *= 0.25;
	viewToSurface = normalize(viewToSurface);
	float vdotn = clamp(dot(viewToSurface, waterNormal), 0.0, 1.0);

	vec3 lightWater = vec3(28.0, 200.0, 255.0) / 255.0;

	vec3 darkWater = vec3(0.0, 60.0, 107.0) / 255.0;
	vec3 shoreDark  = pow(darkWater, vec3(2.4));
	lightWater = darkWater*10.0;
	
	vec3 r = reflect(viewToSurface, waterNormal);
	vec3 ry = r;
	ry.y *= -1.0;

	vec3 hdr = textureLod(u_reflectionHdr, r, 0.0).rgb;

	float baseHeightCoeff = waterHeight *0.5;
	vec3 baseSunCoeff = (ndotl * 0.5 + 0.5)*0.5;
	vec3 baseColor = mix(shoreDark, lightWater,  baseHeightCoeff * baseSunCoeff);
	
	// BASELAYER
	if(stage == uint(3))
	{
		if(substage == uint(0))
		{
			out_color0.rgb = shoreDark * iblStrength;
		}
		else if(substage == uint(1))
		{
			out_color0.rgb = mix(shoreDark, lightWater, baseHeightCoeff) * iblStrength;
		}
		else if(substage == uint(2))
		{
			out_color0.rgb = baseSunCoeff* iblStrength;
		}
		else if(substage >= uint(3))
		{
			out_color0.rgb = baseColor* iblStrength;
		}
		
		return;
	}

	float viewDistanceToDepth = distance(frontPosition, fragmentPosition.xyz);
	float deepCoeff = 1.0 - smoothstep(0.0, 0.95, viewDistanceToDepth);
	float deepCoeff2 = 1.0 - smoothstep(0.0, 0.05, viewDistanceToDepth);
	vec3 frontColorCaustics = mix(frontColor, frontColor + (frontColor*sunStrength), causticsCoeff *(1.0 -deepCoeff2));

	// apply deep color 
	float shallowCoeff = 1.0 - smoothstep(0.0, 0.35, viewDistanceToDepth);
	vec3 shallowColor = mix(baseColor, frontColorCaustics, 0.15);
	vec3 deepColor = mix(frontColorCaustics*baseColor, shallowColor, shallowCoeff);

	// basecolor
	out_color0.rgb = mix(baseColor, deepColor,  0.1 + deepCoeff*0.9) * iblStrength;

	float viewDepthCoeff2 = 1.0 - smoothstep(0.0, 0.2, viewDistanceToDepth);
	out_color0.rgb = mix(out_color0.rgb, frontColor, viewDepthCoeff2);
    
	float reflCoeff = smoothstep(0.1, 0.7, waterHeight);

	// fresnel reflection
	float reflectionStrength = 10.0;
	out_color0.rgb += hdr * vdotn * reflectionStrength * reflCoeff * sunStrength;
	out_color0.a = 1.0;

	// REFLECTION
	if(stage == uint(4))
	{
		if(substage == uint(0))
		{
			out_color0.rgb = hdr * sunStrength;
		}
		else if(substage == uint(1))
		{
			out_color0.rgb = vec3(vdotn * sunStrength);
		}
		else if(substage == uint(2))
		{
			out_color0.rgb = vec3(vdotn * sunStrength* reflCoeff);
		}
		else if(substage >= uint(3))
		{
			out_color0.rgb = baseColor* iblStrength;
			out_color0.rgb += hdr * vdotn * reflectionStrength * sunStrength* reflCoeff;
		}
		return;
	}

	// BOTTOMLAYER
	if(stage == uint(5))
	{
		vec3 baseColorDBG = baseColor;
		float viewDistanceToDepthDBG = distance(frontPosition_OG, fragmentPosition.xyz);
		float deepCoeffDBG = 1.0 - smoothstep(0.0, 1.0, viewDistanceToDepthDBG);
		float shallowCoeffDBG = 1.0 - smoothstep(0.0, 0.35, viewDistanceToDepthDBG);
		vec3 shallowColor2 = mix(baseColorDBG, frontColor_OG, 0.15);
		vec3 deepColor2 = mix(frontColor_OG*baseColorDBG, shallowColor2, shallowCoeffDBG);
		vec3 deepColor3 =  mix(baseColorDBG, deepColor2,  0.1 + deepCoeffDBG*0.9);

		if(substage == uint(0))
		{
			out_color0.rgb = shallowColor2* iblStrength;
		}
		else if(substage == uint(1))
		{
			out_color0.rgb = deepColor2* iblStrength;
		}
		else if(substage == uint(2))
		{
			out_color0.rgb = deepColor3* iblStrength;
			
			out_color0.rgb += hdr * vdotn * reflectionStrength * sunStrength* reflCoeff;
		}
		else if (substage >= uint(3))
		{
			float viewDepthCoeff2DBG = 1.0 - smoothstep(0.0, 0.2, viewDistanceToDepthDBG);

			out_color0.rgb = deepColor3* iblStrength;
			out_color0.rgb = mix(out_color0.rgb, frontColor_OG, viewDepthCoeff2DBG);
			out_color0.rgb += hdr * vdotn * reflectionStrength * sunStrength* reflCoeff;

			
		}

		return;
	}

	// CAUSTICS
	if(stage ==uint(6))
	{

		float causticsCoeffDBG = sampleWaterHeight(frontPosition_OG.xz, time);
		causticsCoeffDBG = smoothstep(0.0, 0.5, causticsCoeffDBG);

		vec3 baseColorDBG = baseColor;
		float viewDistanceToDepthDBG = distance(frontPosition_OG, fragmentPosition.xyz);
		float deepCoeffDBG = 1.0 - smoothstep(0.0, 1.0, viewDistanceToDepthDBG);
		float shallowCoeffDBG = 1.0 - smoothstep(0.0, 0.35, viewDistanceToDepthDBG);

		float deepCoeff2DBG = 1.0 - smoothstep(0.0, 0.2, viewDistanceToDepthDBG);
		vec3 frontColorCausticsDBG = mix(frontColor_OG, frontColor_OG + (frontColor_OG*sunStrength), causticsCoeffDBG *(1.0 -deepCoeff2DBG));

		vec3 shallowColor2 = mix(baseColorDBG, frontColorCausticsDBG, 0.15);
		vec3 deepColor2 = mix(frontColorCausticsDBG*baseColorDBG, shallowColor2, shallowCoeffDBG);
		vec3 deepColor3 =  mix(baseColorDBG, deepColor2,  0.1 + deepCoeffDBG*0.9);
		float viewDepthCoeff2DBG = 1.0 - smoothstep(0.0, 0.026, viewDistanceToDepthDBG);

		if(substage == uint(0))
		{
			out_color0.rgb = vec3(causticsCoeffDBG *(1.0 -deepCoeff2DBG));
		}
		else if(substage == uint(1))
		{
			out_color0.rgb = frontColorCausticsDBG;
		}
		else if(substage >= uint(2))
		{
			out_color0.rgb = deepColor3* iblStrength;
			out_color0.rgb = mix(out_color0.rgb, frontColor_OG, viewDepthCoeff2DBG);
			out_color0.rgb += hdr * vdotn * reflectionStrength * sunStrength* reflCoeff;
		}

		return;
	}

	// REFRACTION
	if(stage == uint(7))
	{
		if(substage == uint(0))
		{
			vec2 refracUv = (waterViewNormal.xz) * 0.5 + 0.5;
			out_color0.rgb = vec3(refracUv.x, refracUv.y, 0.0);

		}
		else if(substage == uint(1))
		{
			out_color0.rgb = frontColor.rgb;
		}

		return;
	}
}