#include "dxstdafx.h"
#include "PostProcManager.h"
#include "TargetableTexture.h"
#include "PPEffectsCommon.h"
#include "PostProcFXFilesManager.h"
#include "PostProcFXChain.h"
#include "ScaledPPEffect.h"
#include "EffectParam.h"

#pragma once


// ----------------------------- GEOMETRY FOR POST PROCESS ---------------------------
struct PPVert
{
	float x, y, z, rhw;	// Already transformed vertices list to feed the driver
	float tu0, tv0;		// Texcoord for post-process source
	float tu1, tv1;		// Texcoord for the original scene

	const static D3DVERTEXELEMENT9 Decl[4];

	void setVertex(float fX, float fY, float fZ, float fRHW, float fTU0, float fTV0, float fTU1, float fTV1)
	{
		x=fX; y=fY; z=fZ; rhw=fRHW;
		tu0=fTU0; tv0=fTV0;
		tu1=fTU1; tv1=fTV1;
	}	
};


// Vertex declaration for post-processing
const D3DVERTEXELEMENT9 PPVert::Decl[4] =
{
	{ 0, 0,  D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITIONT,	0 },
	{ 0, 16, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,	0 },
	{ 0, 24, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,	1 },
	D3DDECL_END()
};


//------------------------------------------------------------------------------------------
// Constructor
// Check the max simultaneous RTs and create an Array of NULL pointers to surfaces
//------------------------------------------------------------------------------------------
PostProcManager::PostProcManager()
{
	// Init everything to NULL...

	// The quad
	m_pQuad = NULL;
	m_pQuadVB = NULL;
	m_pVertDeclPP = NULL;
	
	// The RTs
	m_nNbAvailableRTs = 0;
	m_pTextureTargetsSwapSet[0] = m_pTextureTargetsSwapSet[1] = m_pInitialTextureTargets = NULL;
	m_nNbSimultRTs = 0;

	// The helpers
	m_pSrcTextures = NULL;
	m_pOldRT = NULL;
	m_pCurrentEffect = NULL;
	m_nCurrentSrcSet = 0;

	// Retrieve the D3D device and quit if no good
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return;

	// Retrieve the device caps about RTs
	D3DCAPS9 deviceCaps;
	pd3dDevice->GetDeviceCaps(&deviceCaps);
	m_nNbSimultRTs = deviceCaps.NumSimultaneousRTs;

	// If we can use render targets
	if(m_nNbSimultRTs != 0)
	{
		// Create the pre-post-process render targets
		m_pInitialTextureTargets = new TargetableTexture[m_nNbSimultRTs];

		// Init the sets of targetable textures
		m_pTextureTargetsSwapSet[0] = new TargetableTexture[m_nNbSimultRTs];
		m_pTextureTargetsSwapSet[1] = new TargetableTexture[m_nNbSimultRTs];
	}

	// Can we use 64 bits floating point render targets?
	LPDIRECT3D9 pD3D = DXUTGetD3DObject();
	if(!pD3D)
		m_TargetFormat = D3DFMT_UNKNOWN;
	else
	{
		D3DDISPLAYMODE DisplayMode;
		pd3dDevice->GetDisplayMode(0, &DisplayMode);
		if(FAILED(pD3D->CheckDeviceFormat(deviceCaps.AdapterOrdinal, deviceCaps.DeviceType, DisplayMode.Format, D3DUSAGE_RENDERTARGET, D3DRTYPE_TEXTURE, D3DFMT_A16B16G16R16F)))
			m_TargetFormat = D3DFMT_A8R8G8B8;
		else
			m_TargetFormat = D3DFMT_A16B16G16R16F;
	}
}


//--------------------------------------------------------------------------------------
// Destructor
//--------------------------------------------------------------------------------------
PostProcManager::~PostProcManager()
{
	// Deleting the RTs
	EmptyRTsArrays();		
	delete [] m_pInitialTextureTargets;
	delete [] m_pTextureTargetsSwapSet[0];
	delete [] m_pTextureTargetsSwapSet[1];
	
	// Deleting the quad
	ReleaseQuad();
}


//--------------------------------------------------------------------------------------
// Creates a new RT with given parameters and store it in the targets sets nIndex slot
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::CreateRTTextureSets(const UINT Width, const UINT Height, const int nIndex)
{
	HRESULT hr;

	// Checking the nIndex bounds
	if(nIndex<0 || nIndex>(m_nNbSimultRTs-1)) return D3DERR_INVALIDCALL;

	//Init the original set with the maximum precision available (manually capped to 64bits for memory usage matters)
	V_RETURN(m_pInitialTextureTargets[nIndex].Init(Width, Height, m_TargetFormat));

	// The post-process working sets have the same precision has the backbuffer
	V_RETURN( m_pTextureTargetsSwapSet[0][nIndex].Init(Width, Height, D3DFMT_A8R8G8B8) );
	return m_pTextureTargetsSwapSet[1][nIndex].Init(Width, Height, D3DFMT_A8R8G8B8);
}


//--------------------------------------------------------------------------------------
// Creates nCount identical RTs ( 0<nCount<=m_nNbSimultRTs )
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::Init(const UINT Width, const UINT Height, const int nCount)
{
	// Checking the nIndex bounds
	if(nCount<1 || nCount>m_nNbSimultRTs) return D3DERR_INVALIDCALL;

	// Cleaning the previous inits
	EmptyRTsArrays();

	// Creating the RT textures sets
	HRESULT hr;
	for(int i=0; i<nCount; ++i)
		V_RETURN(CreateRTTextureSets(Width, Height, i));

	// Init the helpers for the post processing stages
	m_pSrcTextures = new LPDIRECT3DTEXTURE9[nCount];

	// Saving the number of created RT per set
	m_nNbAvailableRTs = nCount;

	return S_OK;
}


//--------------------------------------------------------------------------------------
// Swap m_nCurrentSrcSet and change the render targets
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::SwapRTSets()
{
	HRESULT hr;

	// Device cheking
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return D3DERR_INVALIDDEVICE;

	// Swap m_nCurrentSet
	m_nCurrentSrcSet=!m_nCurrentSrcSet;

	// And affect the new RTs (only the one we'll writes to!
	// It'll avoid the GPU to clean unused targets (one of the most ressource hungry call)
	if(m_pTextureTargetsSwapSet[m_nCurrentSrcSet])
	{
		// Retrieve the list of the RT we'll write to
		const bool * pRTWriteList = m_pCurrentEffect->getRTWriteList();

		// Init the used RT and set the unused to NULL
		for(int i=0; i<m_nNbAvailableRTs; ++i)
			V_RETURN(pd3dDevice->SetRenderTarget(i, (pRTWriteList[i])?m_pTextureTargetsSwapSet[m_nCurrentSrcSet][i].GetSurface():NULL));
	}

	return S_OK;
}


//--------------------------------------------------------------------------------------
// Empty the RT array
//--------------------------------------------------------------------------------------
void PostProcManager::EmptyRTsArrays()
{
	// Release the initial render targets
	for(int i=0; i<m_nNbAvailableRTs; ++i)
		m_pInitialTextureTargets[i].ReleaseAll();

	// Release the render targets sets
	for(int i=0; i<m_nNbAvailableRTs; ++i)	
		m_pTextureTargetsSwapSet[0][i].ReleaseAll();
	for(int i=0; i<m_nNbAvailableRTs; ++i)
		m_pTextureTargetsSwapSet[1][i].ReleaseAll();

	// And the helpers
	delete [] m_pSrcTextures;
	m_pSrcTextures = NULL;

	m_nNbAvailableRTs = 0;
}


//--------------------------------------------------------------------------------------
// Backup RT0 in m_pOldRT
// Return invalidcall if m_pOldRT already contain something
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::BackupRT0()
{
	// Is there already an unfreed buffer backup?
	if(m_pOldRT) return D3DERR_INVALIDCALL;

	// Device checking
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return D3DERR_INVALIDDEVICE;

	return(pd3dDevice->GetRenderTarget(0, &m_pOldRT));
}


//--------------------------------------------------------------------------------------
// Create a quad the same resolution as the screen
//--------------------------------------------------------------------------------------
void PostProcManager::CreateQuad()
{
	// Device cheking
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return;
	
	// Creating the quad mesh containers
	m_pQuad = new PPVert[4];			// Allocating 4 vertices and the vertex buffer
	pd3dDevice->CreateVertexBuffer(sizeof(PPVert)*4, D3DUSAGE_WRITEONLY|D3DUSAGE_DYNAMIC, 0, D3DPOOL_DEFAULT, &m_pQuadVB, NULL); // The vertex buffer
	pd3dDevice->CreateVertexDeclaration(PPVert::Decl, &m_pVertDeclPP);	// Finally, create the vertex declaration to use this mesh

	// Setup a quad the same size as the backbuffer
	SetQuadDimToScreen();
}


//--------------------------------------------------------------------------------------
// Setup the quad to be the same size as the screen
//--------------------------------------------------------------------------------------
void PostProcManager::SetQuadDimToScreen()
{
	// Safe exit tests
	if(!m_pQuad || !m_pQuadVB) return;

	// Setup a quad the same size as the backbuffer
	const D3DSURFACE_DESC* pd3dsdBackBuffer = DXUTGetBackBufferSurfaceDesc();
	m_pQuad[0].setVertex(-0.5f							, -0.5f								, 1.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f);
	m_pQuad[1].setVertex(pd3dsdBackBuffer->Width - 0.5f	, -0.5f								, 1.0f, 1.0f, 1.0f, 0.0f, 1.0f, 0.0f);
	m_pQuad[2].setVertex(-0.5f							, pd3dsdBackBuffer->Height - 0.5f	, 1.0f, 1.0f, 0.0f, 1.0f, 0.0f, 1.0f);
	m_pQuad[3].setVertex(pd3dsdBackBuffer->Width - 0.5f	, pd3dsdBackBuffer->Height - 0.5f	, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f);

	// Update the vertex buffer
	UpdateQuadVB();
}


//--------------------------------------------------------------------------------------
// Modify the quad vertices list, as if sclaing the quad using fScaleX & fScaleY
//--------------------------------------------------------------------------------------
void PostProcManager::ScaleQuad(const D3DXVECTOR2 * pScaler)
{
	bool bUpdateVB = false;

	// Should we update the post-process U data?
	if( m_pQuad[1].tu0 != m_vUVExtend.x )
	{
		m_pQuad[1].tu0 = m_pQuad[3].tu0 = m_vUVExtend.x;
		bUpdateVB = true;
	}

	// Should we update the post-process V data?
	if( m_pQuad[2].tv0 != m_vUVExtend.y )
	{
		m_pQuad[2].tv0 = m_pQuad[3].tv0 = m_vUVExtend.y;
		bUpdateVB = true;
	}

	// Scaling along X
	if( pScaler->x != 1.0f )
	{
		m_pQuad[1].x = (m_pQuad[1].x + 0.5f) * pScaler->x - 0.5f;
		m_pQuad[3].x = (m_pQuad[3].x + 0.5f) * pScaler->x - 0.5f;
		bUpdateVB = true;
	}

	// Scaling along Y
	if( pScaler->y != 1.0f )
	{
		m_pQuad[2].y = (m_pQuad[2].y + 0.5f) * pScaler->y - 0.5f;
		m_pQuad[3].y = (m_pQuad[3].y + 0.5f) * pScaler->y - 0.5f;
		bUpdateVB = true;
	}

	// Update the VB ?
	if(bUpdateVB)
	{
		// We compute the new extends
		m_vUVExtend.x *= pScaler->x;
		m_vUVExtend.y *= pScaler->y;		

		// And then update the quad VB data
		UpdateQuadVB();
	}
}


//--------------------------------------------------------------------------------------
// Update the quad vertex buffer with vertices list (create the VB if not exist)
//--------------------------------------------------------------------------------------
void PostProcManager::UpdateQuadVB()
{
	// Device cheking
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return;

	if(m_pQuadVB)
	{
		// Fill in the vertex buffer
		LPVOID pVBData;
		if( SUCCEEDED(m_pQuadVB->Lock(0, 0, &pVBData, D3DLOCK_DISCARD)) )
		{
			CopyMemory(pVBData, m_pQuad, sizeof(m_pQuad));
			m_pQuadVB->Unlock();
		}
	}
}


//--------------------------------------------------------------------------------------
// Destroy the quad vertices
//--------------------------------------------------------------------------------------
void PostProcManager::ReleaseQuad()
{
	// Deleting the quad VB
	SAFE_RELEASE(m_pQuadVB);

	// Deleting the quad vertices
	if(m_pQuad)
	{
		delete m_pQuad;
		m_pQuad = NULL;
	}

	// Release the quad vertex declaration
	SAFE_RELEASE(m_pVertDeclPP);
}


//--------------------------------------------------------------------------------------
// Prepare the render targets (and other things?) before rendering
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::SetupInitialRenderTargets(const bool bClearRTs)
{
	HRESULT hr;

	// Retrieve the current device
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return D3DERR_INVALIDDEVICE;

	// Backup the RT0 if asked
	V_RETURN(BackupRT0());

	// Setup the render targets
	for(int i=0; i<m_nNbAvailableRTs; ++i)
	{
		// Retrieving surfaces from the textures and allocating them as render target
		pd3dDevice->SetRenderTarget(i, m_pInitialTextureTargets[i].GetSurface());
	}

	// Clear the RTs
	if(bClearRTs)
		V_RETURN(pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET, D3DCOLOR_ARGB(0, 0, 0, 0), 1.0f, 0));

	return S_OK;
}


//--------------------------------------------------------------------------------------
// Perform all post process operations in the list
// At the end, restore RT0 with m_pOldRT if not NULL
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::PerformPostProcess(PostProcFXChain * pFXChain, float fTime)
{
	// If we aren't allowed to do post processing do nothing
	if(!pFXChain) return D3DERR_INVALIDCALL;
	if(pFXChain->GetSize()==0) return S_OK;

	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return D3DERR_INVALIDDEVICE;

	// --------------------- TEST ONLY -----------------------//
	HRESULT hr;
	int i;
	int j;

	// We first re-init the quad values and its extend, and flag the post-process stage
	m_vUVExtend.x = m_vUVExtend.y = 1.0f;
	SetQuadDimToScreen();

	// From now on we'll use the post-process quad vertex declaration
	pd3dDevice->SetVertexDeclaration(m_pVertDeclPP);

	// We'll use a special array to avoid loosing a source texture during the process
	for(i=0; i<m_nNbAvailableRTs; ++i)
		m_pSrcTextures[i] = m_pInitialTextureTargets[i].GetTexture();	// We'll use the initial texture set for the 1st process

	// We retrieve an access pointer on the Chainlist
	const list<ScaledPPEffect*> * pFXList = pFXChain->GetList();

	// And now we parse the post-processes list
	const int nLastEffectIndex = (int)(pFXList->size()-1);
	list<ScaledPPEffect*>::const_iterator EffectIterator;
	for(i=0, EffectIterator=pFXList->begin(); EffectIterator!=pFXList->end(); ++EffectIterator, ++i)
	{
		// Retrieve the currently processed ScaledPPEffect
		ScaledPPEffect * pCurrentScaledPPEffect = (*EffectIterator);

		// Retrieve the current effect
		m_pCurrentEffect = pCurrentScaledPPEffect->GetEffect();
		if(!m_pCurrentEffect) return D3DERR_NOTAVAILABLE;
		
		// Scaling the PP-quad if needed
		D3DXVECTOR2 vPPScaler = pCurrentScaledPPEffect->GetScaler();
		ScaleQuad(&vPPScaler);

		// We always bind the initial renders to every effect
		V_RETURN(m_pCurrentEffect->BindEffectSceneTextures(m_pInitialTextureTargets, m_nNbAvailableRTs));

		// And we bind the new sources textures
		V_RETURN(m_pCurrentEffect->BindEffectRTTextures(m_pSrcTextures, m_nNbAvailableRTs));

		// We bind the time value
		V_RETURN(m_pCurrentEffect->BindTime(fTime));

		//if we're not doing the last post-process effect...
		if(i<nLastEffectIndex)
		{
			// Swap the render target sets
			SwapRTSets();

			// After an effect is done we update the sources for the next effect
			const bool * bRTWrites = m_pCurrentEffect->getRTWriteList();
			for(j=0; j<m_nNbAvailableRTs; ++j)
			{
				// If the effect wrote on RT[i]
				if(bRTWrites[j])
					m_pSrcTextures[j] = m_pTextureTargetsSwapSet[m_nCurrentSrcSet][j].GetTexture(); // The next source is the result stored on the current RT swap set				
			}
		}
		else
		{
			// Setting up the render targets >0 to NULL
			for(j=1; j<m_nNbAvailableRTs; ++j)
				pd3dDevice->SetRenderTarget(j, NULL);

			// We'll finally ouput to the original backbuffer
			if(m_pOldRT)
			{
				V_RETURN(pd3dDevice->SetRenderTarget(0, m_pOldRT));
				SAFE_RELEASE(m_pOldRT);
			}
		}

		// Perform a single post process using this effect
		V_RETURN(PerformSinglePostProcess(pCurrentScaledPPEffect));
	}

	// Alright baby! No more effect to use, we just finished PP!
	m_pCurrentEffect = NULL;
	return S_OK;
}


//--------------------------------------------------------------------------------------
// Perform a single post process operation
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::PerformSinglePostProcess(ScaledPPEffect * pScaledPPEffect) const
{
	HRESULT hr;

	// Checking parameter
	if(!pScaledPPEffect) return D3DERR_INVALIDCALL;

	// Device cheking
	LPDIRECT3DDEVICE9 pd3dDevice = DXUTGetD3DDevice();
	if(!pd3dDevice) return D3DERR_INVALIDDEVICE;

	// Setup the effect technique
	V_RETURN(m_pCurrentEffect->SetEffectTechnique());

	// We'll need to access the effect
	LPD3DXEFFECT pD3DXEffect = m_pCurrentEffect->getEffect();
	if(!pD3DXEffect) return D3DERR_NOTAVAILABLE;

	// Setup the parameters if some are found
	for(int i=0; i<m_pCurrentEffect->GetNbAvailableParams(); ++i)
	{
		// We should have some available parameters but we actually don't have any (!?!)
		EffectParam * pParam = m_pCurrentEffect->getParameter(i);		
		if(!pParam) return D3DERR_MOREDATA;

		// Setup the effect parameters using the values stored in the scaled effect
		pD3DXEffect->SetFloatArray(pParam->GetParamHandle(), pScaledPPEffect->GetParamValues(i), pParam->GetParamSize());
	}

	// Then we can render the quad
    if( SUCCEEDED( pd3dDevice->BeginScene() ) )
    {
		// Clearing the RTs
		V_RETURN(pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET, 0x00000000, 1.0f, 0));

		// Draw the quad
        UINT cPasses, p;
		pD3DXEffect->Begin(&cPasses, 0);
        for(p=0; p<cPasses; ++p)
		{
			pD3DXEffect->BeginPass(p);
			
			// Draw the quad (under a thousand of vertices DrawPrimitiveUp is as good as DrawPrimitive)		
			V_RETURN(pd3dDevice->DrawPrimitiveUP(D3DPT_TRIANGLESTRIP, 2, m_pQuad, sizeof(PPVert)));
			
			// End of this pass
			pD3DXEffect->EndPass();
        }
		pD3DXEffect->End();

		pd3dDevice->EndScene();
	}

	// Single post process operation successfully done
	return S_OK;
}


//--------------------------------------------------------------------------------------
// OnResetDevice
//--------------------------------------------------------------------------------------
HRESULT PostProcManager::OnResetDevice(const int nWidth, const int nHeight)
{	
	// Setup the post-processing quad
	CreateQuad();

	return S_OK;
}


//--------------------------------------------------------------------------------------
// OnLostDevice
//--------------------------------------------------------------------------------------
void PostProcManager::OnLostDevice()
{	
	// Release the quad and the RTs array
	ReleaseQuad();
	EmptyRTsArrays();
}