float4x4 g_WorldMatrix;
float4x4 g_ViewMatrix;
float4x4 g_ProjectionMatrix;
float4x4 g_NormalMatrix;

Texture2D g_MRT0;


float  g_fNormalScale;
float  g_fAcceptAngle;
float2 g_f2RTSize;                  
float  g_fHDAORejectRadius;       
float  g_fHDAOIntensity;    
float  g_fHDAOAcceptRadius;


SamplerState samPoint
{
    Filter = MIN_MAG_MIP_POINT;
    
    AddressU = Wrap;
    AddressV = Wrap;
};

SamplerState samLinear
{
    Filter = MIN_MAG_MIP_LINEAR;
    
    AddressU = Wrap;
    AddressV = Wrap;
};

SamplerState samAniso
{
    Filter = ANISOTROPIC;
    MaxAnisotropy = 16;
    
    AddressU = Wrap;
    AddressV = Wrap;
};

struct VS_INPUT
{
    float4 Pos  : POSITION;
    float3 Norm : NORMAL;
    float3 Tang : TANGENT; 
    float2 Tex  : TEXCOORD;
};

struct PS_INPUT
{
    float4 Pos  : SV_POSITION;
    float2 Tex  : TEXCOORD0;
};

BlendState AlphaBlending
{
    BlendEnable[0] = TRUE;
    SrcBlend = SRC_ALPHA;
    DestBlend = INV_SRC_ALPHA;
    BlendOp = ADD;
    RenderTargetWriteMask[0] = 0x0F;
};

BlendState DisableBlend 
{ 
	BlendEnable[0] = FALSE; 
};

DepthStencilState DepthEnabling 
{ 
	DepthEnable = TRUE; 
};

DepthStencilState DepthDisabling
{
	DepthEnable = FALSE;
};




//--------------------------------------------------------------------------------------
// Gather pattern
//--------------------------------------------------------------------------------------

// Gather defines
#define RING_1    (1)
#define RING_2    (2)
#define RING_3    (3)
#define RING_4    (4)
#define NUM_RING_1_GATHERS    (2)
#define NUM_RING_2_GATHERS    (6)
#define NUM_RING_3_GATHERS    (12)
#define NUM_RING_4_GATHERS    (20)

// Ring sample pattern
static const float2 g_f2HDAORingPattern[NUM_RING_4_GATHERS] = 
{
    // Ring 1
    { 1, -1 },
    { 0, 1 },
    
    // Ring 2
    { 0, 3 },
    { 2, 1 },
    { 3, -1 },
    { 1, -3 },
        
    // Ring 3
    { 1, -5 },
    { 3, -3 },
    { 5, -1 },
    { 4, 1 },
    { 2, 3 },
    { 0, 5 },
    
    // Ring 4
    { 0, 7 },
    { 2, 5 },
    { 4, 3 },
    { 6, 1 },
    { 7, -1 },
    { 5, -3 },
    { 3, -5 },
    { 1, -7 }
};

// Ring weights
static const float4 g_f4HDAORingWeight[NUM_RING_4_GATHERS] = 
{
    // Ring 1 (Sum = 5.30864)
    { 1.00000, 0.50000, 0.44721, 0.70711 },
    { 0.50000, 0.44721, 0.70711, 1.00000 },
    
    // Ring 2 (Sum = 6.08746)
    { 0.30000, 0.29104, 0.37947, 0.40000 },
    { 0.42426, 0.33282, 0.37947, 0.53666 },
    { 0.40000, 0.30000, 0.29104, 0.37947 },
    { 0.53666, 0.42426, 0.33282, 0.37947 },
    
    // Ring 3 (Sum = 6.53067)
    { 0.31530, 0.29069, 0.24140, 0.25495 },
    { 0.36056, 0.29069, 0.26000, 0.30641 },
    { 0.26000, 0.21667, 0.21372, 0.25495 },
    { 0.29069, 0.24140, 0.25495, 0.31530 },
    { 0.29069, 0.26000, 0.30641, 0.36056 },
    { 0.21667, 0.21372, 0.25495, 0.26000 },
    
    // Ring 4 (Sum = 7.00962)
    { 0.17500, 0.17365, 0.19799, 0.20000 },
    { 0.22136, 0.20870, 0.24010, 0.25997 },
    { 0.24749, 0.21864, 0.24010, 0.28000 },
    { 0.22136, 0.19230, 0.19799, 0.23016 },
    { 0.20000, 0.17500, 0.17365, 0.19799 },
    { 0.25997, 0.22136, 0.20870, 0.24010 },
    { 0.28000, 0.24749, 0.21864, 0.24010 },
    { 0.23016, 0.22136, 0.19230, 0.19799 }
};

static const float g_fRingWeightsTotal[RING_4] =
{
    5.30864,
    11.39610,
    17.92677,
    24.93639
};

#define NUM_NORMAL_LOADS (4)
static const int2 g_i2NormalLoadPattern[NUM_NORMAL_LOADS] = 
{
    { 1, 8 },
    { 8, -1 },
    { 5, 4 },
    { 4, -4 }
};

//----------------------------------------------------------------------------------------
// Helper function to Gather samples
//----------------------------------------------------------------------------------------
void GatherZNZSamples( Texture2D tex, float2 texCoord, out float4 z, out float4 nz )
{
    float4 sample0 = tex.SampleLevel( samPoint, texCoord, 0, int2( 0, 1 ) );
    float4 sample1 = tex.SampleLevel( samPoint, texCoord, 0, int2( 1, 1 ) );
    float4 sample2 = tex.SampleLevel( samPoint, texCoord, 0, int2( 1, 0 ) );
    float4 sample3 = tex.SampleLevel( samPoint, texCoord, 0, int2( 0, 0 ) );
   
    z.x = sample0.a;
    z.y = sample1.a;
    z.z = sample2.a;
    z.w = sample3.a;
    
    nz.x = sample0.b;
    nz.y = sample1.b;
    nz.z = sample2.b;
    nz.w = sample3.b;
}

//--------------------------------------------------------------------------------------
// Helper function to compute camera space XYZ from depth and texture coords
//--------------------------------------------------------------------------------------
float3 GetCameraXYZFromDepth( float depth, int2 i2ScreenCoord )
{
    return normalize(float3(2.0 * i2ScreenCoord.x / 1600.0 - 1.0, 2.0 * i2ScreenCoord.y / 900.0 - 1.0, 1.0)) * depth;
}


//--------------------------------------------------------------------------------------
// Used as an early rejection test - based on geometry
//--------------------------------------------------------------------------------------
float GeometryRejectionTest( int2 i2ScreenCoord )
{
    float3 f3N[3];
    float3 f3Pos[3];
    float3 f3Dir[2];
    float fDot;
    float fSummedDot = 0.0f;
    int2 i2MirrorPattern;
    int2 i2OffsetScreenCoord;
    int2 i2MirrorOffsetScreenCoord;
    
    float4 sample = g_MRT0.Sample(samPoint, float2(i2ScreenCoord.x / 1600.0, i2ScreenCoord.y / 900.0));
	float3 normal = sample.rgb;
	float depth = sample.a;
    
    f3Pos[0] = GetCameraXYZFromDepth( depth, i2ScreenCoord );
    f3N[0] = normal;

    f3Pos[0] -= ( f3N[0] * g_fNormalScale );

    for( int iNormal=0; iNormal<NUM_NORMAL_LOADS; iNormal++ ) {
        i2MirrorPattern = ( g_i2NormalLoadPattern[iNormal] + int2( 1, 1 ) ) * int2( -1, -1 );
        i2OffsetScreenCoord = i2ScreenCoord + g_i2NormalLoadPattern[iNormal];
        i2MirrorOffsetScreenCoord = i2ScreenCoord + i2MirrorPattern;
        
        // Clamp our test to screen coordinates
        i2OffsetScreenCoord = ( i2OffsetScreenCoord > ( g_f2RTSize - float2( 1.0f, 1.0f ) ) ) ? ( g_f2RTSize - float2( 1.0f, 1.0f ) ) : ( i2OffsetScreenCoord );
        i2MirrorOffsetScreenCoord = ( i2MirrorOffsetScreenCoord > ( g_f2RTSize - float2( 1.0f, 1.0f ) ) ) ? ( g_f2RTSize - float2( 1.0f, 1.0f ) ) : ( i2MirrorOffsetScreenCoord );
        i2OffsetScreenCoord = ( i2OffsetScreenCoord < 0 ) ? ( 0 ) : ( i2OffsetScreenCoord );
        i2MirrorOffsetScreenCoord = ( i2MirrorOffsetScreenCoord < 0 ) ? ( 0 ) : ( i2MirrorOffsetScreenCoord );
        
        sample = g_MRT0.Sample(samPoint, float2(i2OffsetScreenCoord.x / 1600.0, i2OffsetScreenCoord.y / 900.0));
        normal = sample.rgb;
        depth = sample.a;
        
        f3Pos[1] = GetCameraXYZFromDepth( depth, i2OffsetScreenCoord );
        f3N[1] = normal;
        
        sample = g_MRT0.Sample(samPoint, float2(i2MirrorOffsetScreenCoord.x / 1600.0, i2MirrorOffsetScreenCoord.y / 900.0));
        normal = sample.rgb;
        depth = sample.a;
        
        f3Pos[2] = GetCameraXYZFromDepth( depth, i2MirrorOffsetScreenCoord );
        f3N[2] = normal;

                
        f3Pos[1] -= ( f3N[1] * g_fNormalScale );
        f3Pos[2] -= ( f3N[2] * g_fNormalScale );
        
        f3Dir[0] = f3Pos[1] - f3Pos[0];
        f3Dir[1] = f3Pos[2] - f3Pos[0];
        
        f3Dir[0] = normalize( f3Dir[0] );
        f3Dir[1] = normalize( f3Dir[1] );
        
        fDot = dot( f3Dir[0], f3Dir[1] );
        
        fSummedDot += ( fDot + 2.0f );
    }
        
    return ( fSummedDot * 0.125f );
}


//--------------------------------------------------------------------------------------
// Used as an early rejection test - based on normals
//--------------------------------------------------------------------------------------
float NormalRejectionTest( int2 i2ScreenCoord )
{
    float3 f3N1;
    float3 f3N2;
    float fDot;
    float fSummedDot = 0.0f;
    int2 i2MirrorPattern;
    int2 i2OffsetScreenCoord;
    int2 i2MirrorOffsetScreenCoord;
  

    for( int iNormal=0; iNormal<NUM_NORMAL_LOADS; iNormal++ )
    {
        i2MirrorPattern = ( g_i2NormalLoadPattern[iNormal] + int2( 1, 1 ) ) * int2( -1, -1 );
        i2OffsetScreenCoord = i2ScreenCoord + g_i2NormalLoadPattern[iNormal];
        i2MirrorOffsetScreenCoord = i2ScreenCoord + i2MirrorPattern;
        
        // Clamp our test to screen coordinates
        i2OffsetScreenCoord = ( i2OffsetScreenCoord > ( g_f2RTSize - float2( 1.0f, 1.0f ) ) ) ? ( g_f2RTSize - float2( 1.0f, 1.0f ) ) : ( i2OffsetScreenCoord );
        i2MirrorOffsetScreenCoord = ( i2MirrorOffsetScreenCoord > ( g_f2RTSize - float2( 1.0f, 1.0f ) ) ) ? ( g_f2RTSize - float2( 1.0f, 1.0f ) ) : ( i2MirrorOffsetScreenCoord );
        i2OffsetScreenCoord = ( i2OffsetScreenCoord < 0 ) ? ( 0 ) : ( i2OffsetScreenCoord );
        i2MirrorOffsetScreenCoord = ( i2MirrorOffsetScreenCoord < 0 ) ? ( 0 ) : ( i2MirrorOffsetScreenCoord );
                        
        f3N1 = g_MRT0.Sample(samPoint, float2(i2OffsetScreenCoord.x / 1600.0, i2OffsetScreenCoord.y / 900.0)).rgb;
		f3N2 = g_MRT0.Sample(samPoint, float2(i2MirrorOffsetScreenCoord.x / 1600.0, i2MirrorOffsetScreenCoord.y / 900.0)).rgb;

        fDot = dot( f3N1, f3N2 );
        
        fSummedDot += ( fDot > g_fAcceptAngle ) ? ( 0.0f ) : ( 1.0f - ( abs( fDot ) * 0.25f ) );
    }
        
    return ( 0.5f + fSummedDot * 0.25f  );
}


//--------------------------------------------------------------------------------------
// Vertex Shader
//--------------------------------------------------------------------------------------
PS_INPUT SSAO_VS( VS_INPUT input )
{
    PS_INPUT output = (PS_INPUT)0;

    output.Pos = mul(input.Pos, g_WorldMatrix);
    output.Pos = mul(output.Pos, g_ViewMatrix);
    output.Pos = mul(output.Pos, g_ProjectionMatrix);
    
    output.Tex = input.Tex;

    return output;
}

float4 SSAO_PS( PS_INPUT input, uniform int iNumRingGathers, uniform int iNumRings ) : SV_Target
{
    int2 i2ScreenCoord;
    float2 f2ScreenCoord_10_1;
    float2 f2ScreenCoord;
    float2 f2MirrorScreenCoord;
    float2 f2TexCoord_10_1;
    float2 f2MirrorTexCoord_10_1;
    float2 f2TexCoord;
    float2 f2MirrorTexCoord;
    float2 f2InvRTSize;
    float fCenterZ;
    float fOffsetCenterZ;
    float fCenterNormalZ;
    float4 f4SampledZ[2];
    float4 f4OffsetSampledZ[2];
    float4 f4SampledNormalZ[2];
    float4 f4Diff;
    float4 f4Compare[2];
    float4 f4Occlusion = 0.0f;
    float fOcclusion;
    int iGather;
    float fDot = 1.0f;
    float2 f2KernelScale = float2( g_f2RTSize.x / 1600.0, g_f2RTSize.y / 900.0 );
    float3 f3CameraPos;
                            
    // Compute integer screen coord, and store off the inverse of the RT Size
    f2InvRTSize = 1.0f / g_f2RTSize;
    f2ScreenCoord = input.Tex * g_f2RTSize;
    i2ScreenCoord = int2( f2ScreenCoord );
                
    // Test the normals to see if we should apply occlusion
    fDot = NormalRejectionTest( i2ScreenCoord );    
    
    float4 sample;

    if( fDot > 0.5f )
    {                
        // Sample the center pixel for camera Z
        
        sample = g_MRT0.Sample(samPoint, input.Tex);
		fCenterNormalZ = sample.b;
		fCenterZ = sample.a;
	
        fOffsetCenterZ = fCenterZ + fCenterNormalZ * g_fNormalScale;
        
            
        // Loop through each gather location, and compare with its mirrored location
        for( iGather=0; iGather<iNumRingGathers; iGather++ )
        {
            f2MirrorScreenCoord = ( ( f2KernelScale * g_f2HDAORingPattern[iGather] ) + float2( 1.0f, 1.0f ) ) * float2( -1.0f, -1.0f );
            
            f2TexCoord = float2( ( f2ScreenCoord + ( f2KernelScale * g_f2HDAORingPattern[iGather] ) ) * f2InvRTSize );
            f2MirrorTexCoord = float2( ( f2ScreenCoord + ( f2MirrorScreenCoord ) ) * f2InvRTSize );
            
            // Sample
           
			GatherZNZSamples( g_MRT0, f2TexCoord, f4SampledZ[0], f4SampledNormalZ[0] );
            GatherZNZSamples( g_MRT0, f2MirrorTexCoord, f4SampledZ[1], f4SampledNormalZ[1] );
            
                        
            // Detect valleys
            f4Diff = fCenterZ.xxxx - f4SampledZ[0];
            f4Compare[0] = ( f4Diff < g_fHDAORejectRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            f4Compare[0] *= ( f4Diff > g_fHDAOAcceptRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            
            f4Diff = fCenterZ.xxxx - f4SampledZ[1];
            f4Compare[1] = ( f4Diff < g_fHDAORejectRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            f4Compare[1] *= ( f4Diff > g_fHDAOAcceptRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            
            f4Occlusion.xyzw += ( g_f4HDAORingWeight[iGather].xyzw * ( f4Compare[0].xyzw * f4Compare[1].zwxy ) * fDot );    

                
            // Scale normals
            f4OffsetSampledZ[0] = f4SampledZ[0] + ( f4SampledNormalZ[0] * g_fNormalScale );
            f4OffsetSampledZ[1] = f4SampledZ[1] + ( f4SampledNormalZ[1] * g_fNormalScale );
                        
            // Detect valleys
            f4Diff = fOffsetCenterZ.xxxx - f4OffsetSampledZ[0];
            f4Compare[0] = ( f4Diff < g_fHDAORejectRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            f4Compare[0] *= ( f4Diff > g_fHDAOAcceptRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            
            f4Diff = fOffsetCenterZ.xxxx - f4OffsetSampledZ[1];
            f4Compare[1] = ( f4Diff < g_fHDAORejectRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            f4Compare[1] *= ( f4Diff > g_fHDAOAcceptRadius.xxxx ) ? ( 1.0f ) : ( 0.0f );
            
            f4Occlusion.xyzw += ( g_f4HDAORingWeight[iGather].xyzw * ( f4Compare[0].xyzw * f4Compare[1].zwxy ) * fDot );   
            
            
        }
    }
                    
    // Finally calculate the HDAO occlusion value
    //if( bUseNormals )
    //{
    //    fOcclusion = ( ( f4Occlusion.x + f4Occlusion.y + f4Occlusion.z + f4Occlusion.w ) / ( 3.0f * g_fRingWeightsTotal[iNumRings - 1] ) );
    //}
    //else
    //{
        fOcclusion = ( ( f4Occlusion.x + f4Occlusion.y + f4Occlusion.z + f4Occlusion.w ) / ( 2.0f * g_fRingWeightsTotal[iNumRings - 1] ) );
    //}
    fOcclusion *= ( g_fHDAOIntensity );
    fOcclusion = 1.0f - saturate( fOcclusion );
    return float4(0.0, 0.0, 0.0, fOcclusion);
}


//--------------------------------------------------------------------------------------
technique10 Render
{
    pass P0
    {
		SetBlendState( AlphaBlending, float4( 0.0f, 0.0f, 0.0f, 0.0f ), 0xFFFFFFFF );
		SetDepthStencilState( DepthDisabling, 0 );
		
        SetVertexShader( CompileShader( vs_4_0, SSAO_VS() ) );
        SetGeometryShader( NULL );
        SetPixelShader( CompileShader( ps_4_0, SSAO_PS( NUM_RING_4_GATHERS, RING_4 ) ) );
    }
}