// Copyright (c) Valve Corporation, All rights reserved. ======================================================================================================
//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// These functions are an attempt to standardize the lighting models across HLVR scene shaders, with the goal of accelerating shader authoring and
// homogenizing lighting calculations and making code extension easier.
//
// PS_InitFinalCombiner or PS_CommonProcessing should be called at the top of all pixel shaders and sets up a common structure that is used to communicate to our standardized
// final combiner shader. Custom work can be done to modify attributes in the FinalCombinerInput_t after this is called.
// PS_InitFinalCombiner assumes that you want to control the "optional" parameters yourself. These *need* to be set up by you if you want correct lighting.
//
// PS_FinalCombiner should be called at the end and works on the FinalCombinerInput_t data only. 
// This does lighting, tonemapping, etc. in a standardized way.
//
//
// NOTE: These functions do the bare minimum of vertex and pixel setup.
//		 For standard world shaders that automatically support features like texture mapping, alpha transparency and normal mapping
//       there is a better starting point in the vr_shared_standard_PS_code.fxc headers
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------

#ifndef VR_COMMON_PS_CODE_FXC
#define VR_COMMON_PS_CODE_FXC

#include "math_general.fxc"
#include "encoded_normals.fxc"
#include "vr_gradient_fog.fxc"
#include "vr_cubemap_fog.fxc"
#include "vr_lighting.fxc"
#include "volumetric_fog.fxc"

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// Common Params
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------

#if ( S_UNLIT )
	BoolAttribute( unlit, true );
#else
	BoolAttribute( NeedsLightProbe, true );
#endif

bool g_bFogEnabled < Default( 1.0f ); UiType( CheckBox ); UiGroup( "Fog,300" ); >;

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// Common Render States
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------

//
// Skip pixels that have stencil set to 0x1 (to support vr checkerboard rendering)
// Set up enough state that if we set the PassOp to Replace, we'll write to stencil bit 2 (used to mark dynamic objects)
//
#if !defined( STENCIL_ALREADY_SET )
	RenderState( StencilEnable, true );
	RenderState( StencilPassOp, KEEP );
	RenderState( StencilFunc, NOT_EQUAL );
	RenderState( StencilRef, 3 );
	RenderState( StencilReadMask, 1 );
	RenderState( StencilWriteMask, 254 );

	#if ( S_RENDER_BACKFACES )
		RenderState( BackStencilPassOp, KEEP );
		RenderState( BackStencilFunc, NOT_EQUAL );
	#endif
#endif

#if !defined( DEPTH_STATE_ALREADY_SET )
	#define DEPTH_STATE_ALREADY_SET
	RenderState( DepthEnable, true );
	RenderState( DepthWriteEnable, true );
	RenderState( DepthFunc, GREATER_EQUAL );
#endif

//
// Disable backface culling
//
#if !defined( CULL_MODE_ALREADY_SET )
	#if ( S_RENDER_BACKFACES )
		RenderState( CullMode, NONE );
	#endif
#endif

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// Common Pixel Processing Code
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------

FinalCombinerInput_t PS_InitFinalCombiner()
{
	FinalCombinerInput_t o;

	//
	// Setup default values
	//
	o.vDiffuseColor					= float3( 1.0, 1.0, 1.0 );
	o.vSpecularColor				= float3( 0.0, 0.0, 0.0 );
	o.flOpacity						= 1.0;
	o.vNormalTs						= float3( 0.0, 0.0, 1.0 );
	o.vRoughness					= float2( 0.0, 0.0 );
	o.vEmissive						= float3( 0.0, 0.0, 0.0 );
	o.vTangentUWs					= float3( 1.0, 0.0, 0.0 );
	o.vTangentVWs					= float3( 0.0, 1.0, 0.0 );
	o.flAmbientOcclusion			= 1.0;
	o.flFresnelExponent				= 5.0;
	o.flRetroReflectivity			= 0.0;

	// S_AMBIENT_OCCLUSION_COLOR_BLEED
	o.vAmbientOcclusionColorBleed = float3( 1.0, 1.0, 1.0 );
	o.flAmbientOcclusionDirectDiffuse = 0.0;
	o.flAmbientOcclusionDirectSpecular = 1.0;

	// S_TWO_LOBE_SPECULAR
	o.flTwoLobeSpecularRatio		= 0.0;

	// Subsurface Scattering
	o.vSSSNormalWs					= float3( 0.0, 0.0, 0.0 );
	o.flSSSCurvature				= 0.0;
	o.vSSSWrapColor					= float3( 1.0, 1.0, 1.0 );
	o.vSSSWrapParameters			= SUBSURFACE_SCATTERING_DEFAULT_WRAP_PARAMETERS;
	o.flSSSMask						= 1.0;

	// S_TRANSMISSIVE
	o.vTransmissiveMask				= float3( 0.0, 0.0, 0.0 );
	o.vTransmissiveFalloff			= float3( 0.0, 0.0, 0.0 );

	o.vPositionSs					= float4( 0.0, 0.0, 0.0, 0.0 );
	o.vPositionWs					= float3( 0.0, 0.0, 0.0 );
	o.vPositionWithOffsetWs			= float3( 0.0, 0.0, 0.0 );
	o.bUseBentNormalsForIndirect	= false;
	o.vBentNormalWs					= float3( 1.0, 1.0, 1.0 );
	o.vBentDiffuseNormalWs			= float3( 1.0, 1.0, 1.0 );
	o.vAnisotropy					= float2( 0.0, 0.0 );
	o.vPerPixelTangentUWs			= float3( 0.0, 0.0, 0.0 );
	o.vPerPixelTangentVWs			= float3( 0.0, 0.0, 0.0 );
	o.flWrinkle						= 0.0;

	o.vLightmapUV = float2( 0.0f, 0.0f );

	// Tools Vis
	{
		o.vTextureCoords.xy = float2( 0, 0 );
	}

	return o;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
#ifdef COMMON_PS_INPUT_DEFINED

FinalCombinerInput_t PS_CommonProcessing( const PS_INPUT i )
{
	FinalCombinerInput_t o = PS_InitFinalCombiner();

	o.vDiffuseColor = i.vVertexColor.rgb;
	o.flOpacity = i.vVertexColor.a;

	float3 vNormalWs = i.vNormalWs.xyz;

	// Negate the world normal if we are rendering the back face
	#if ( S_RENDER_BACKFACES )
	{
		vNormalWs *= ( ( i.face ? 1.0 : -1.0 ) );
	}
	#endif

	o.vNormalWs = normalize( vNormalWs );

	#if ( S_USE_PER_VERTEX_CURVATURE )
	{
		o.flSSSCurvature = i.flSSSCurvature;
	}
	#else
	{
		o.flSSSCurvature = length( fwidth( i.vNormalWs.xyz ) ) / length( fwidth( i.vPositionWithOffsetWs.xyz ) );
	}
	#endif

	o.vPositionWithOffsetWs = i.vPositionWithOffsetWs;
	o.vPositionWs = i.vPositionWithOffsetWs + g_vHighPrecisionLightingOffsetWs.xyz;
	o.vPositionSs = i.vPositionSs;


	o.vLightmapUV = i.vLightmapUV.xy;

	#if ( S_WRINKLE )
	{
		o.flWrinkle = i.vNormalWs.w;
		o.flWrinkle = clamp( o.flWrinkle, -1.0, 1.0 );
	}
	#endif

	// Tools Vis
	{
		#ifdef g_tColor
			o.vTextureCoords.xy = i.vTextureCoords.xy;
		#else
			o.vTextureCoords.xy = float2( 0, 0 );
		#endif	
	}

	return o;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
#if ( PS_INPUT_HAS_TANGENT_BASIS )

void PS_CommonTransformNormal( const PS_INPUT i, inout FinalCombinerInput_t finalCombinerInput, float3 vNormalTs )
{
	float3 vNormalWs = finalCombinerInput.vNormalWs.xyz;

	float3 vTangentUWs = normalize( i.vTangentUWs.xyz );
	float3 vTangentVWs = normalize( i.vTangentVWs.xyz );

	// HACK: Tools still generate tangent space the inverted Source1 way where positive y is down. Flipping the normal here to compensate.
	vNormalTs.y = -vNormalTs.y;

	//
	// Transform from tangent space into world space
	//
	finalCombinerInput.vTangentUWs = vTangentUWs;
	finalCombinerInput.vTangentVWs = vTangentVWs;
	finalCombinerInput.vNormalTs = vNormalTs;
	finalCombinerInput.vNormalWs = Vec3TsToWsNormalized( vNormalTs.xyz, vNormalWs.xyz, vTangentUWs.xyz, vTangentVWs.xyz ); // Need to normalize because tangent space basis may not be orthogonal due to interpolation
}

#endif

#endif // COMMON_PS_INPUT_DEFINED

//-------------------------------------------------------------------------------------------------------------------------------------------------------------

struct PS_OUTPUT
{
	float4 vColor : SV_Target0;
};

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
float3 CalculateDiffuseAmbientOcclusion( FinalCombinerInput_t finalCombinerInput, LightingTerms_t lightingTerms )
{
	float flAmbientOcclusion = min( finalCombinerInput.flAmbientOcclusion, sqrt( lightingTerms.flBakedAmbientOcclusion ) );
	flAmbientOcclusion = min( lightingTerms.flDynamicAmbientOcclusion, flAmbientOcclusion );

	// Jimenez et al. 2016, "Practical Realtime Strategies for Accurate Indirect Occlusion"
	float3 albedo = finalCombinerInput.vDiffuseColor.rgb;
    float3 a =  2.0404 * albedo - 0.3324;
    float3 b = -4.7951 * albedo + 0.6417;
    float3 c =  2.7552 * albedo + 0.6903;

    return max( float3( flAmbientOcclusion.xxx ), ( ( flAmbientOcclusion * a + b ) * flAmbientOcclusion + c ) * flAmbientOcclusion );
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
float3 CalculateSpecularAmbientOcclusion( FinalCombinerInput_t finalCombinerInput, LightingTerms_t lightingTerms )
{
	float flAmbientOcclusion = min( finalCombinerInput.flAmbientOcclusion, lightingTerms.flBakedAmbientOcclusion );
	flAmbientOcclusion = min( lightingTerms.flDynamicAmbientOcclusion, flAmbientOcclusion );

	if ( finalCombinerInput.bUseBentNormalsForIndirect )
	{
		return PowerSpecularOcclusion( finalCombinerInput.vPositionWs.xyz, finalCombinerInput.vBentNormalWs.xyz, flAmbientOcclusion ).xxx;
	}
	else
	{
		return PowerSpecularOcclusion( finalCombinerInput.vPositionWs.xyz, finalCombinerInput.vNormalWs.xyz, flAmbientOcclusion ).xxx;
	}
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
FinalCombinerInput_t CalculateDiffuseAndSpecularFromAlbedoAndMetalness( FinalCombinerInput_t f, float3 vAlbedo, float flMetalness )
{
	f.vAlbedo = vAlbedo.rgb;
	f.vDiffuseColor = vAlbedo.rgb * ( 1.0f - flMetalness );
	f.vSpecularColor = lerp( float3( 0.04f, 0.04f, 0.04f ), vAlbedo.rgb, flMetalness );
	return f;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// Final Combiner. 
//
// This function should be run by all pixel shaders prior to framebuffer submission, will perform lighting, tonemapping, fogging etc.
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------
PS_OUTPUT PS_FinalCombinerDoLighting( FinalCombinerInput_t f, inout LightingTerms_t lightingTerms )
{
	PS_OUTPUT o;

	#if ( !S_UNLIT )
	{
		#if ( S_ANISOTROPIC_GLOSS )
		{
			// Calculate anisotropy
			float2 vAnisotropicRatio = f.vRoughness.xy / f.vRoughness.yx;
			f.vAnisotropy.x = step( vAnisotropicRatio.y, vAnisotropicRatio.x ); // 1 if anisotropy is greatest along tangentU, 0 otherwise
			f.vAnisotropy.y = 1.0 - min( vAnisotropicRatio.x, vAnisotropicRatio.y );

			// Calculate per-pixel tangent space
			f.vPerPixelTangentUWs = normalize( cross( f.vTangentVWs.xyz, f.vNormalWs.xyz ) );
			f.vPerPixelTangentVWs = normalize( cross( f.vNormalWs.xyz, f.vTangentUWs.xyz ) );
		}
		#endif

		// Roughness adjusted to fix geometric specular aliasing
		f.vRoughness.xy = AdjustRoughnessByGeometricNormal( f.vRoughness.xy, f.vNormalWs.xyz );

		// Compute lighting
		ComputeDirectLighting( lightingTerms, f );
		CalculateIndirectLighting( lightingTerms, f );

		//
		// Diffuse AO
		//
		float3 vDiffuseAO = CalculateDiffuseAmbientOcclusion( f, lightingTerms );
		lightingTerms.vIndirectDiffuse.rgb *= vDiffuseAO.rgb;
		lightingTerms.vDiffuse.rgb *= lerp( float3( 1.0, 1.0, 1.0 ), vDiffuseAO.rgb, f.flAmbientOcclusionDirectDiffuse );

		//
		// Specular AO
		//
		float3 vSpecularAO = CalculateSpecularAmbientOcclusion( f, lightingTerms );
		lightingTerms.vIndirectSpecular.rgb *= vSpecularAO.rgb;
		lightingTerms.vSpecular.rgb *= lerp( float3( 1.0, 1.0, 1.0 ), vSpecularAO.rgb, f.flAmbientOcclusionDirectSpecular );

		// Zero-out specular if no specular lighting
		#if defined( S_SPECULAR ) && ( S_SPECULAR == 0 ) 
		{
			lightingTerms.vSpecular = float3( 0.0, 0.0, 0.0 );
			lightingTerms.vIndirectSpecular = float3( 0.0, 0.0, 0.0 );
		}
		#endif // S_SPECULAR
	}
	#endif // S_UNLIT

	// Store translucency value in alpha
	o.vColor.a = f.flOpacity;

	//
	// Add Diffuse Lighting Contribution
	//
	o.vColor.rgb = ( lightingTerms.vDiffuse.rgb + lightingTerms.vIndirectDiffuse.rgb ) * f.vDiffuseColor.rgb;

	//
	// Emissive
	//
	o.vColor.rgb += f.vEmissive.rgb;

	//
	// Add Specular Contribution
	//
	o.vColor.rgb += lightingTerms.vSpecular.rgb;
	o.vColor.rgb += lightingTerms.vIndirectSpecular.rgb; 

	//
	// Add transmissive
	//
	o.vColor.rgb += lightingTerms.vTransmissive.rgb * f.vTransmissiveMask.rgb;

	return o;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
PS_OUTPUT PS_FinalCombinerDoPostProcessing( FinalCombinerInput_t f, LightingTerms_t lightingTerms, PS_OUTPUT o )
{
	float3 vPositionWs = f.vPositionWs;
	float4 vPositionSs = f.vPositionSs;

	if ( g_bFogEnabled )
	{
		// S_ADDITIVE_BLEND This is only used in one mode of static_overlay!!! Can ffuck off soon
		#if ( S_ADDITIVE_BLEND )
		{
			//
			// When blending additively, dest pixel contains the fog term.
			// Therefore, we want to scale alpha with fog amount in front of us,
			// rather than double-adding the fog
			//
			const float3 vPositionToCameraWs = vPositionWs.xyz - g_vCameraPositionWs;

			if ( g_bGradientFogEnabled )
			{
				o.vColor.a *= 1.0 - CalculateGradientFog( vPositionWs, vPositionToCameraWs ).a;
			}

			if ( g_bCubemapFogEnabled )
			{
				o.vColor.a *= 1.0 - CalculateCubemapFog( vPositionWs, vPositionToCameraWs ).a;
			}

			if ( g_bVolumetricFogEnabled )
			{
				o.vColor.a *= CalculateVolumetricFog( vPositionWs.xyz, vPositionSs.xy ).a;
			}
		}
		#else
		{
			o.vColor.rgb = Fog::Apply( vPositionWs.xyz, vPositionSs.xy, o.vColor.rgb );
		}
		#endif
	}

	if( DepthNormals::WantsDepthNormals() )
	{
		o.vColor = DepthNormals::Output( f.vNormalWs.xyz, f.vRoughness.x, f.flOpacity );
		return o;
	}

	[branch]
	if ( ToolsVis::WantsToolsVis() )
	{
		//
		// Calculate albedo
		//
		float3 vAlbedo = f.vAlbedo.rgb;

        ToolsVis toolVis = ToolsVis::Init(o.vColor.rgba, lightingTerms.vDiffuse.rgb, lightingTerms.vSpecular.rgb, lightingTerms.vIndirectDiffuse.rgb, lightingTerms.vIndirectSpecular.rgb, lightingTerms.vTransmissive.rgb );

		toolVis.HandleFlatOverlayColor( vAlbedo.rgb, o.vColor.rgba );
		toolVis.HandleFullbright( o.vColor.rgba, vAlbedo.rgb, f.vPositionWs.xyz, f.vNormalWs.xyz );
		toolVis.HandleDiffuseLighting( o.vColor.rgba );
		toolVis.HandleSpecularLighting( o.vColor.rgba );
		toolVis.HandleTransmissiveLighting( o.vColor.rgba );
		toolVis.HandleLightingComplexity( o.vColor.rgba, f.vPositionSs.xy, f.vPositionWs.xyz, f.vNormalWs.xyz );

		toolVis.HandleAlbedo( o.vColor.rgba, vAlbedo.rgb );
		toolVis.HandleReflectivity( o.vColor.rgba, vAlbedo.rgb );
		toolVis.HandleRoughness( o.vColor.rgba, f.vRoughness.xy );
		toolVis.HandleDiffuseAmbientOcclusion( o.vColor.rgba, CalculateDiffuseAmbientOcclusion( f, lightingTerms ) );
		toolVis.HandleSpecularAmbientOcclusion( o.vColor.rgba, CalculateSpecularAmbientOcclusion( f, lightingTerms ) );
		toolVis.HandleShaderIDColor( o.vColor.rgba );
		toolVis.HandleCubemapReflections( o.vColor.rgba, f.vPositionWs.xyz, f.vNormalWs.xyz, f.vPositionSs.xy );

		toolVis.HandleNormalTs( o.vColor.rgba, f.vNormalTs.xyz );
		toolVis.HandleNormalWs( o.vColor.rgba, f.vNormalWs.xyz );
		toolVis.HandleTangentUWs( o.vColor.rgba, f.vTangentUWs.xyz );
		toolVis.HandleTangentVWs( o.vColor.rgba, f.vTangentVWs.xyz );

		if ( f.bUseBentNormalsForIndirect )
		{
			toolVis.HandleBentNormalWs( o.vColor.rgba, f.vBentNormalWs.xyz );
		}

		toolVis.HandleGeometricRoughness( o.vColor.rgba, f.vNormalWs.xyz );
		toolVis.HandleCurvature( o.vColor.rgba, f.flSSSCurvature );
		#if defined( COLOR_SAMPLER_DEFINED )
			toolVis.ShowUVs( o.vColor.rgba, vAlbedo.rgb, g_tColor, f.vTextureCoords.xy );
			toolVis.ShowMipUtilization( o.vColor.rgba, vAlbedo.rgb, g_tColor, f.vTextureCoords.xy );
		#endif

		toolVis.HandleWrinkle( o.vColor, f.flWrinkle );

		// Facepunch Pipeline
		toolVis.HandleTiledRenderingColors( o.vColor, vAlbedo.rgb, f.vPositionSs.xy );

		o.vColor.rgb = saturate( o.vColor.rgb ); // turns off bloom in tools vis mode
	}
	
	// Merge up into toolsvis
	if ( g_bWireframeMode )
	{
		o.vColor = g_vWireframeColor;
	}
	return o;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------
//
// Including this version as typically we dont want to declare or access the ligthing ourput
//
//-------------------------------------------------------------------------------------------------------------------------------------------------------------
PS_OUTPUT PS_FinalCombiner( FinalCombinerInput_t finalCombinerInput )
{
	LightingTerms_t lightingTerms = InitLightingTerms();

	PS_OUTPUT o;

	o = PS_FinalCombinerDoLighting( finalCombinerInput, lightingTerms );
	o = PS_FinalCombinerDoPostProcessing( finalCombinerInput, lightingTerms, o );

	return o;
}

//-------------------------------------------------------------------------------------------------------------------------------------------------------------

#endif
