//-----------------------------------------------------------------------------
// File:		LPVCommon.usf
//
// Summary:		Common functionality for LPV shaders
//
// Created:		2013-03-01
//
// Author:		mailto:benwood@microsoft.com
//
//				Copyright (C) Microsoft. All rights reserved.
//-----------------------------------------------------------------------------

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

#define REFINE_OVER_TIME	1

#if REFINE_OVER_TIME
  #define PREV_FRAME_MULTIPLIER			0.9f

  // With REFINE_OVER_TIME enabled and a previous frame multiplier of 0.9, lighting 
  // is effectively multiplied by 10. 
  // Final lighting therefore needs to be divided by 10 to account for this
  #define REFINE_LIGHTING_MULTIPLIER	0.1f
#else
	#define REFINE_LIGHTING_MULTIPLIER	1.0f
#endif
 
#define TILED_ADDRESSING		0

// HACK!
#ifdef LPV_AMBIENT_CUBE
	#undef	LPV_AMBIENT_CUBE 
#endif
 
// TODO: pass these in from a common compile enviroment setup
#define LPV_AMBIENT_CUBE		0				// This has to match the value in LightPropagationVolume.cpp
#define LPV_AMBIENT_CUBE_WITH_SH_BUFFER_HACK 0	// This has to match the value in LightPropagationVolume.cpp
#define LPV_VOLUME_TEXTURE		1				// This has to match the value in LightPropagationVolume.h

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

#define AMBIENT_CUBE_NEGATIVE_X 0
#define AMBIENT_CUBE_POSITIVE_X 1
#define AMBIENT_CUBE_NEGATIVE_Y 2
#define AMBIENT_CUBE_POSITIVE_Y 3
#define AMBIENT_CUBE_NEGATIVE_Z 4
#define AMBIENT_CUBE_POSITIVE_Z 5

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

struct LPVCellPackedAmbientCube
{
#if LPV_AMBIENT_CUBE_WITH_SH_BUFFER_HACK
	uint packedData[14];
#else
	uint packedData[6];
#endif
};

struct LPVCellPackedSH
{
	uint packedCoeffs[14];
};

#if LPV_VOLUME_TEXTURE || LPV_GV_VOLUME_TEXTURE
SamplerState								gLpv3DTextureSampler;
#endif

#if LPV_VOLUME_TEXTURE
Texture3D<float4>							gLpv3DTexture0;
Texture3D<float4>							gLpv3DTexture1;
Texture3D<float4>							gLpv3DTexture2;
Texture3D<float4>							gLpv3DTexture3;
Texture3D<float4>							gLpv3DTexture4;
Texture3D<float4>							gLpv3DTexture5;
Texture3D<float4>							gLpv3DTexture6;
#elif LPV_AMBIENT_CUBE
StructuredBuffer<LPVCellPackedAmbientCube>	gLpvBuffer;
#else
StructuredBuffer<LPVCellPackedSH>			gLpvBuffer;
#endif

// Write shaders only
#if LPV_WRITE_SHADER

#if LPV_VOLUME_TEXTURE
RWTexture3D<float4>							gLpv3DTextureRW0;
RWTexture3D<float4>							gLpv3DTextureRW1;
RWTexture3D<float4>							gLpv3DTextureRW2;
RWTexture3D<float4>							gLpv3DTextureRW3;
RWTexture3D<float4>							gLpv3DTextureRW4;
RWTexture3D<float4>							gLpv3DTextureRW5;
RWTexture3D<float4>							gLpv3DTextureRW6;
#elif LPV_AMBIENT_CUBE
RWStructuredBuffer<LPVCellPackedAmbientCube> gLpvBufferRW;
#else
RWStructuredBuffer<LPVCellPackedSH>			gLpvBufferRW;
#endif

#endif // LPV_WRITE_SHADER

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

struct LPVCell
{
#if LPV_AMBIENT_CUBE
	float3 ambientCube[6];
#else
	float3 coeffs[9];
#endif
};

//-------------------------------------------------------------------------------------------------
// Functions
//-------------------------------------------------------------------------------------------------

void ClearCell( in out LPVCell cell )
{ 
#if LPV_AMBIENT_CUBE
	[unroll]
	for ( int i=0; i<6; i++ ) 
		cell.ambientCube[i] = float3( 0.0f, 0.0f, 0.0f );
#else
	[unroll]
	for ( int i=0; i<9; i++ ) 
		cell.coeffs[i] = float3( 0.0f, 0.0f, 0.0f );
#endif
}

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

void MultiplyCell( in out LPVCell cell, float multiplier )
{
#if LPV_AMBIENT_CUBE
	[unroll]
	for ( int i=0; i<6; i++ ) 
		cell.ambientCube[i] *= multiplier;
#else
	[unroll]
	for ( int i=0; i<9; i++ ) 
		cell.coeffs[i] *= multiplier;
#endif
}

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

void MultiplyCellV( in out LPVCell cellA, in out LPVCell cellB )
{
#if LPV_AMBIENT_CUBE
	[unroll]
	for ( int i=0; i<6; i++ ) 
		cellA.ambientCube[i] *= cellB.ambientCube[i];
#else
	[unroll]
	for ( int i=0; i<9; i++ ) 
		cellA.coeffs[i] *= cellB.coeffs[i];
#endif
}

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

void AddCell( in out LPVCell cellA, in LPVCell cellB )
{
#if LPV_AMBIENT_CUBE
	[unroll]
	for ( int i=0; i<6; i++ ) 
		cellA.ambientCube[i] += cellB.ambientCube[i];
#else
	[unroll]
	for ( int i=0; i<9; i++ ) 
		cellA.coeffs[i] += cellB.coeffs[i];
#endif
}

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

int3 IndexToGridPos( uint cellIndex )
{
	int3 texIndex;
	texIndex.x = ( cellIndex % 32 );
	texIndex.y = ( cellIndex / 32 ) %  32;
	texIndex.z = cellIndex / ( 32 * 32 );
	return texIndex;
}

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

uint GetGridAddress( uint3 gridPos )
{
#if TILED_ADDRESSING
	uint3 tileAddr = gridPos / 4;
	uint3 rem = gridPos % 4;
	return ( tileAddr.z*64 + tileAddr.y*8 + tileAddr.x ) * 64 +
		   rem.z*16 + rem.y * 4 + rem.x;
#else
	return gridPos.z*32*32+gridPos.y*32+gridPos.x;
#endif
}

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

void PackSH3( float3 coeffs[9], out LPVCellPackedSH cell )
{
	uint3 compressedCoeffs[9]; 
	[unroll]
	for ( int i=0; i<9; i++ )
	{
		compressedCoeffs[i] = f32tof16( coeffs[i] );
	}
	
	cell.packedCoeffs[0] = compressedCoeffs[0].x | ( compressedCoeffs[0].y << 16 );
	cell.packedCoeffs[1] = compressedCoeffs[0].z | ( compressedCoeffs[1].x << 16 );
	cell.packedCoeffs[2] = compressedCoeffs[1].y | ( compressedCoeffs[1].z << 16 );

	cell.packedCoeffs[3] = compressedCoeffs[2].x | ( compressedCoeffs[2].y << 16 );
	cell.packedCoeffs[4] = compressedCoeffs[2].z | ( compressedCoeffs[3].x << 16 );
	cell.packedCoeffs[5] = compressedCoeffs[3].y | ( compressedCoeffs[3].z << 16 );

	cell.packedCoeffs[6] = compressedCoeffs[4].x | ( compressedCoeffs[4].y << 16 );
	cell.packedCoeffs[7] = compressedCoeffs[4].z | ( compressedCoeffs[5].x << 16 );
	cell.packedCoeffs[8] = compressedCoeffs[5].y | ( compressedCoeffs[5].z << 16 );

	cell.packedCoeffs[9] = compressedCoeffs[6].x | ( compressedCoeffs[6].y << 16 );
	cell.packedCoeffs[10] = compressedCoeffs[6].z | ( compressedCoeffs[7].x << 16 );
	cell.packedCoeffs[11] = compressedCoeffs[7].y | ( compressedCoeffs[7].z << 16 );

	cell.packedCoeffs[12] = compressedCoeffs[8].x | ( compressedCoeffs[8].y << 16 );
	cell.packedCoeffs[13] = compressedCoeffs[8].z;
}

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

void UnpackSH3( LPVCellPackedSH cell, out float3 coeffs[9] )
{
	uint3 compressedCoeffs;
	compressedCoeffs.x = cell.packedCoeffs[0];
	compressedCoeffs.y = cell.packedCoeffs[0] >> 16;
	compressedCoeffs.z = cell.packedCoeffs[1];
	coeffs[0] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[1] >> 16;
	compressedCoeffs.y = cell.packedCoeffs[2];
	compressedCoeffs.z = cell.packedCoeffs[2] >> 16;
	coeffs[1] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[3];
	compressedCoeffs.y = cell.packedCoeffs[3] >> 16;
	compressedCoeffs.z = cell.packedCoeffs[4];
	coeffs[2] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[4] >> 16;
	compressedCoeffs.y = cell.packedCoeffs[5];
	compressedCoeffs.z = cell.packedCoeffs[5] >> 16;
	coeffs[3] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[6];
	compressedCoeffs.y = cell.packedCoeffs[6] >> 16;
	compressedCoeffs.z = cell.packedCoeffs[7];
	coeffs[4] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[7] >> 16;
	compressedCoeffs.y = cell.packedCoeffs[8];
	compressedCoeffs.z = cell.packedCoeffs[8] >> 16;
	coeffs[5] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[9];
	compressedCoeffs.y = cell.packedCoeffs[9] >> 16;
	compressedCoeffs.z = cell.packedCoeffs[10];
	coeffs[6] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[10] >> 16;
	compressedCoeffs.y = cell.packedCoeffs[11];
	compressedCoeffs.z = cell.packedCoeffs[11] >> 16;
	coeffs[7] = f16tof32( compressedCoeffs );

	compressedCoeffs.x = cell.packedCoeffs[12];
	compressedCoeffs.y = cell.packedCoeffs[12] >> 16;
	compressedCoeffs.z = cell.packedCoeffs[13];
	coeffs[8] = f16tof32( compressedCoeffs );
}

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

// Unpacks an HDR colour from RGY32 
float3 UnpackRGY32( uint packedColour )
{
	const float ONE_OVER_255 = 1.0f/255.0f;
	float3 rgb;
	rgb.r = ( packedColour & 0xff000000 ) >> 24;
	rgb.g = ( packedColour & 0x00ff0000 ) >> 16;
	rgb.b = 255.0f-(rgb.r+rgb.g);
	float luminance = f16tof32( packedColour );
	return rgb * luminance * ONE_OVER_255;
}

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

// Packs an HDR colour to RGY32 
uint PackRGY32( float3 colour )
{
	// @TODO: could normalise for more precision, but greater ALU in the decode
	float luminance = (colour.r+colour.g+colour.b);
	colour.rg /= luminance;

	uint packedValue =	uint( colour.r * 255.0f + 0.5f ) << 24 |
						uint( colour.g * 255.0f + 0.5f ) << 16; 

	packedValue |= f32tof16( luminance );
	return packedValue;
}
//-------------------------------------------------------------------------------------------------

void PackAmbientCube( in float3 unpackedCube[6], out LPVCellPackedAmbientCube packedCube )
{
#if LPV_AMBIENT_CUBE_WITH_SH_BUFFER_HACK
	[unroll]
	for ( int i=0; i<14; i++ ) packedCube.packedData[i] = 0;
#endif

	[unroll]
	for ( int i=0; i<6; i++ )
	{
		packedCube.packedData[i] = PackRGY32( unpackedCube[i] );
	} 

}

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

void UnpackAmbientCube( LPVCellPackedAmbientCube packedCube, out float3 unpackedCube[6] )
{
	[unroll]
	for ( int i=0; i<6; i++ )
	{
		unpackedCube[i] = UnpackRGY32( packedCube.packedData[i] );
	}
}

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

float3 AmbientCubeLookup( float3 ambientCube[6], float3 n ) 
{
	float weight[6];
	weight[0] = saturate(-n.x );
	weight[1] = saturate( n.x );
	weight[2] = saturate(-n.y );
	weight[3] = saturate( n.y );
	weight[4] = saturate(-n.z );
	weight[5] = saturate( n.z );

	float3 colour = 0;
	[unroll]
	for ( int i=0; i<6; i++ )
{
		colour += ambientCube[i] * weight[i];
	}
	return colour;
}

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

float3 SHLookup( float3 coeffs[9], float3 n ) 
{
	const float1 c1 = 0.429043 ;
	const float1 c2 = 0.511664 ;
	const float1 c3 = 0.743125 ;
	const float1 c4 = 0.886227 ;
	const float1 c5 = 0.247708 ;

	float3 n2 = n*n;
	float xy = n.x*n.y ; 
	float yz = n.y*n.z ; 
	float xz = n.x*n.z ;

	float3 colour = c1*coeffs[8]*(n2.x-n2.y) + c3*coeffs[6]*n2.z + c4*coeffs[0] - c5*coeffs[6] 
            + 2.0f*c1*(coeffs[4]*xy + coeffs[7]*xz + coeffs[5]*yz) 
            + 2.0f*c2*(coeffs[3]*n.x+coeffs[1]*n.y+coeffs[2]*n.z) ;

	//return colour;
	return max( colour, float3(0,0,0) );

}

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

float3 LPVCellLookup( in LPVCell cell, in float3 normal ) 
{
#if LPV_AMBIENT_CUBE
	return AmbientCubeLookup( cell.ambientCube, normal );
#else
	return SHLookup( cell.coeffs, normal );
#endif
}

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

float3 SHGetDominantDirectionApproxScalar( in float coeffs[9] )
{
	float3 dir = float3( -coeffs[3], -coeffs[1], -coeffs[2] );
	float len = length( dir );
	float den = ( len > 0.01 ) ? len : 1.0f;
	return dir/den;
}

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

float3 WorldToGrid( float3 worldPos )
{
#if LPV_WRITE_SHADER
	return worldPos * LpvWrite.OneOverLpvScale + LpvWrite.mLpvGridOffset;
#else
	return worldPos * LpvRead.OneOverLpvScale + LpvRead.mLpvGridOffset;
#endif
}

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

float3 GridToWorld( float3 gridPos )
{
#if LPV_WRITE_SHADER
	return ( gridPos - LpvWrite.mLpvGridOffset ) * LpvWrite.LpvScale;
#else
	return ( gridPos - LpvRead.mLpvGridOffset ) * LpvRead.LpvScale;
#endif
}

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

float MaxGridExtent( float3 gridPos )
{
	float3 av = abs(gridPos-15.5)+0.5;
	return max(av.x,max(av.y,av.z));
}

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

int GetGridIndex( float3 worldPos )
{
	float3 p = WorldToGrid( worldPos );

	float maxExtent = MaxGridExtent( p );
	if ( maxExtent > 15.5 )
	{
		return -1;
	}
	else
	{
		int3 ui = int3( p );
		return GetGridAddress( ui );
	}
}

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

LPVCell ReadLpvCell( in uint cellIndex, const bool bUseReadWriteBuffer = false )
{
	LPVCell cell;
#if LPV_VOLUME_TEXTURE
	int4 texIndex = int4( IndexToGridPos( cellIndex ), 0 );
	float4 tex0 = gLpv3DTexture0.Load( texIndex );
	float4 tex1 = gLpv3DTexture1.Load( texIndex );
	float4 tex2 = gLpv3DTexture2.Load( texIndex );
	float4 tex3 = gLpv3DTexture3.Load( texIndex );
	float4 tex4 = gLpv3DTexture4.Load( texIndex );
	float4 tex5 = gLpv3DTexture5.Load( texIndex );
	float4 tex6 = gLpv3DTexture6.Load( texIndex );

	cell.coeffs[0] = tex0.rgb; // tex0.a is reserved for secondary occlusion
	cell.coeffs[1] = tex1.rgb; 
	cell.coeffs[2] = float3( tex1.a, tex2.rg );
	cell.coeffs[3] = float3( tex2.ba, tex3.r );
	cell.coeffs[4] = tex3.gba;
	cell.coeffs[5] = tex4.rgb;
	cell.coeffs[6] = float3( tex4.a, tex5.rg );
	cell.coeffs[7] = float3( tex5.ba, tex6.r );
	cell.coeffs[8] = tex6.gba;
#elif LPV_AMBIENT_CUBE
	LPVCellPackedAmbientCube PackedCube;
  #if LPV_WRITE_SHADER
	if ( bUseReadWriteBuffer )
	{
		PackedCube = gLpvBufferRW[cellIndex];
	}
	else
  #endif
	{
		PackedCube = gLpvBuffer[cellIndex];
	}
	UnpackAmbientCube( PackedCube, cell.ambientCube );
#else
	LPVCellPackedSH packedCell;
  #if LPV_WRITE_SHADER
	if ( bUseReadWriteBuffer )
	{
		packedCell = gLpvBufferRW[cellIndex];
	}
	else
  #endif
	{
		packedCell = gLpvBuffer[cellIndex];
	}
	UnpackSH3( packedCell, cell.coeffs );
#endif
	return cell;
}

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

#if LPV_VOLUME_TEXTURE
LPVCell ReadLpvCellVolumeTextureFiltered( float3 gridPos )
{
	float3 texPos = gridPos / 32.0f; // Half pixel offset?
	float4 tex0 = gLpv3DTexture0.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex1 = gLpv3DTexture1.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex2 = gLpv3DTexture2.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex3 = gLpv3DTexture3.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex4 = gLpv3DTexture4.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex5 = gLpv3DTexture5.SampleLevel( gLpv3DTextureSampler, texPos, 0 );
	float4 tex6 = gLpv3DTexture6.SampleLevel( gLpv3DTextureSampler, texPos, 0 );

	LPVCell cell; 
	cell.coeffs[0] = tex0.rgb; // tex0.a is reserved for secondary occlusion
	cell.coeffs[1] = tex1.rgb; 
	cell.coeffs[2] = float3( tex1.a, tex2.rg );
	cell.coeffs[3] = float3( tex2.ba, tex3.r );
	cell.coeffs[4] = tex3.gba;
	cell.coeffs[5] = tex4.rgb;
	cell.coeffs[6] = float3( tex4.a, tex5.rg );
	cell.coeffs[7] = float3( tex5.ba, tex6.r );
	cell.coeffs[8] = tex6.gba;

	return cell;
}
#endif

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