#include "dxstdafx.h"
#include "ToneMap.h"

ToneMap::ToneMap(void)
{
}

ToneMap::~ToneMap(void)
{
	Gain = 1.0f;
	Passes = 1.0f;
}

void ToneMap::SetGain(float gain) {
	Gain = gain;
}

void ToneMap::SetPasses(float pass) {
	Passes = pass;
}

void ToneMap::SetSource(LPDIRECT3DTEXTURE9 _sourceTexture) {
	this->sourceTexture = _sourceTexture;
}

HRESULT ToneMap::Init( IDirect3DDevice9* pd3dDevice, LPDIRECT3DTEXTURE9 sourceTexture, LPDIRECT3DSURFACE9 targetTextureSurface) {
	this->sourceTexture = sourceTexture;
	this->myRenderTextureSurface = targetTextureSurface;
	this->pd3dDevice = pd3dDevice;

	D3DSURFACE_DESC desc;

	sourceTexture->GetLevelDesc(0, &desc);

	int sWidth = desc.Width;
	int sHeight = desc.Height;

	// Create Luminance Texture
	HRESULT h = pd3dDevice->CreateTexture(sWidth, sHeight ,0,D3DUSAGE_RENDERTARGET,D3DFMT_G16R16F,D3DPOOL_DEFAULT,&myLuminanceTexture,NULL);
	myLuminanceTexture->GetSurfaceLevel(0,&myLuminanceTextureSurface);	

	// Create Average Luminance Texture pair
	pd3dDevice->CreateTexture(sWidth ,sHeight ,0,D3DUSAGE_RENDERTARGET,D3DFMT_G16R16F,D3DPOOL_DEFAULT,&myAverageLuminanceTexture[0],NULL);
	myAverageLuminanceTexture[0]->GetSurfaceLevel(0,&myAverageLuminanceTextureSurface[0]);	

	pd3dDevice->CreateTexture(sWidth ,sHeight ,0,D3DUSAGE_RENDERTARGET,D3DFMT_G16R16F,D3DPOOL_DEFAULT,&myAverageLuminanceTexture[1],NULL);
	myAverageLuminanceTexture[1]->GetSurfaceLevel(0,&myAverageLuminanceTextureSurface[1]);	

	// Error Buffer
	LPD3DXBUFFER errBuff;
	DWORD effectCompileFlag = 0;

	// Shader options for debuging
	#ifdef DEBUG_SHADER
	effectCompileFlag|=D3DXSHADER_DEBUG;
	effectCompileFlag |= D3DXSHADER_FORCE_PS_SOFTWARE_NOOPT;
	effectCompileFlag |= D3DXSHADER_FORCE_VS_SOFTWARE_NOOPT;
	#endif

	// Create the tone mapping effect from file
	HRESULT hr = D3DXCreateEffectFromFile(pd3dDevice,L"ToneMap.fx",NULL,NULL,effectCompileFlag,NULL,&myEffect,&errBuff);
	if(hr!=S_OK){
		int BufSize = errBuff->GetBufferSize();

		// displaying error message of arbitrary length
		wchar_t* wbuf = new wchar_t[BufSize];
		mbstowcs( wbuf, (const char*)errBuff->GetBufferPointer(), BufSize );
		MessageBox(NULL, wbuf, L".fx Compilation Error", MB_ICONERROR);		// show error message

		delete wbuf;
		return hr;
	}
	SAFE_RELEASE(errBuff);
	
	//targetTextureSurface->GetDesc(&desc);
	int tWidth = desc.Width;
	int tHeight = desc.Height;
	V(pd3dDevice->CreateDepthStencilSurface(tWidth,tHeight,D3DFMT_D16,D3DMULTISAMPLE_NONE ,0,true,&rttDepthSurface,NULL));
	V(pd3dDevice->CreateDepthStencilSurface(tWidth/4,tHeight/4,D3DFMT_D16,D3DMULTISAMPLE_NONE ,0,true,&rttDepthSurface_small,NULL));
	
	// Quad rendering declarations
/*	rttVertices[0] = D3DVERTEX_1(-1,-1,0,1, 0.0f, 2.0f);
	rttVertices[1] = D3DVERTEX_1(-1,3,0,1, 0.0f, 0.0f);
	rttVertices[2] = D3DVERTEX_1(3,-1,0,1, 2.0f, 2.0f);
*/
	rttVertices[0] = D3DVERTEX_1(-3,1,0,1, -1.0f, 0.0f);
	rttVertices[1] = D3DVERTEX_1(1,1,0,1, 1.0f, 0.0f);
	rttVertices[2] = D3DVERTEX_1(1, -3,0,1, 1.0f, 2.0f);
	//Vertex declaration
	D3DVERTEXELEMENT9 _texturedecl[]=
	{
		{0,0,D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION,0 },
		{0,16,D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,0 },
		D3DDECL_END()
	};
	V(pd3dDevice->CreateVertexDeclaration(_texturedecl,&textureDecl));
	V(pd3dDevice->CreateVertexBuffer(sizeof(D3DVERTEX_1)*3,D3DUSAGE_WRITEONLY,D3DFVF_CUSTOMVERTEX,D3DPOOL_MANAGED,&pQuadVB,NULL));
	void *pData;

	pQuadVB->Lock(0,sizeof(pData),(void**)&pData,0);
	memcpy(pData,rttVertices,sizeof(D3DVERTEX_1)*3);
	pQuadVB->Unlock();

    return S_OK;
}

void ToneMap::Destroy() {

	SAFE_RELEASE(myEffect);
	SAFE_RELEASE(myLuminanceTexture);
	SAFE_RELEASE(myLuminanceTextureSurface);
	SAFE_RELEASE(myAverageLuminanceTexture[0] );
	SAFE_RELEASE(myAverageLuminanceTextureSurface[0] );
	SAFE_RELEASE(myAverageLuminanceTexture[1] );
	SAFE_RELEASE(myAverageLuminanceTextureSurface[1] );
	SAFE_RELEASE(rttDepthSurface);
	SAFE_RELEASE(rttDepthSurface_small);
	SAFE_RELEASE(textureDecl);

	SAFE_RELEASE(myRenderTextureSurface);

	SAFE_RELEASE(pQuadVB);
}

void ToneMap::Map() {
	HRESULT hr;
	
	D3DXMATRIX modelViewProjMatrix;
	D3DXMatrixIdentity(&modelViewProjMatrix);

	V( pd3dDevice->GetRenderTarget(0,&backBuffer) );
	V( pd3dDevice->GetDepthStencilSurface(&backBufferDepthSurface) );

	// Step 1: Calculate the luminance
	D3DSURFACE_DESC desc;
	sourceTexture->GetLevelDesc(0,&desc);

	V( pd3dDevice->SetRenderTarget(0,myLuminanceTextureSurface));
	V( pd3dDevice->SetDepthStencilSurface(rttDepthSurface_small) ); 

	V( pd3dDevice->SetVertexDeclaration(textureDecl) );
	V( pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_ARGB(255, 0, 0, 0), 1.0f, 0) );
	
    // Render the scene with the Luminance effect
    if( SUCCEEDED( pd3dDevice->BeginScene() ) )
    {
		UINT numberofPasses=0;
		D3DXHANDLE techniqueHandle = myEffect->GetTechniqueByName("Luminance");
		V(myEffect->SetTechnique(techniqueHandle));
		V(myEffect->SetTexture("SourceTexture",sourceTexture));
		V(myEffect->SetMatrix("modelViewProjection",&modelViewProjMatrix));
		V(myEffect->SetFloat("dsWidth", desc.Width));
		V(myEffect->SetFloat("dsHeight", desc.Height));
		myEffect->CommitChanges();
		V(myEffect->Begin(&numberofPasses,0));
		for(UINT i = 0; i<numberofPasses; i++ )
		{
			V(myEffect->BeginPass(i));
			
			//Render triangle.
			V(pd3dDevice->SetStreamSource(0,pQuadVB,0,sizeof(D3DVERTEX_1)));
			V(pd3dDevice->DrawPrimitive(D3DPT_TRIANGLELIST,0,1));
			
			V(myEffect->EndPass());
		}
		V(myEffect->End());
    }

	// Step 2: Downsampling with Gauss filter
	//		   Initial pass

	myAverageLuminanceTextureSurface[1]->GetDesc(&desc);

	V( pd3dDevice->SetRenderTarget(0,myAverageLuminanceTextureSurface[1]));
	V( pd3dDevice->SetDepthStencilSurface(rttDepthSurface) ); 
	V( pd3dDevice->SetVertexDeclaration(textureDecl) );
	V( pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_ARGB(255, 0, 0, 0), 1.0f, 0) );
	
    // Render the scene with DownSample effect
    {
		UINT numberofPasses=0;
		D3DXHANDLE techniqueHandle = myEffect->GetTechniqueByName("Blur_v");
		V(myEffect->SetTechnique(techniqueHandle));
		V(myEffect->SetTexture("LuminanceTexture",myLuminanceTexture));
		V(myEffect->SetMatrix("modelViewProjection",&modelViewProjMatrix));
		V(myEffect->SetFloat("dsWidth", desc.Width));
		V(myEffect->SetFloat("dsHeight", desc.Height));
		myEffect->CommitChanges();
		V(myEffect->Begin(&numberofPasses,0));
		for(UINT i = 0; i<numberofPasses; i++ )
		{
			V(myEffect->BeginPass(i));
			
			//Render triangle.
			V(pd3dDevice->SetStreamSource(0,pQuadVB,0,sizeof(D3DVERTEX_1)));
			V(pd3dDevice->DrawPrimitive(D3DPT_TRIANGLELIST,0,1));

			V(myEffect->EndPass());
		}
		V(myEffect->End());
    }

	// Repetitive downsampling pass
	// Ping-pong between the two buffer
	for (int ds_pass=0; ds_pass<Passes; ds_pass++) {
	// buffer1 -> buffer0
	V( pd3dDevice->SetRenderTarget(0,myAverageLuminanceTextureSurface[0]));
	V( pd3dDevice->SetDepthStencilSurface(rttDepthSurface) ); 
	V( pd3dDevice->SetVertexDeclaration(textureDecl) );
	V( pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_ARGB(255, 0, 0, 0), 1.0f, 0) );
	
    {
		UINT numberofPasses=0;
		D3DXHANDLE techniqueHandle = myEffect->GetTechniqueByName("Blur_h");
		V(myEffect->SetTechnique(techniqueHandle));
		V(myEffect->SetTexture("LuminanceTexture",myAverageLuminanceTexture[1]));
		V(myEffect->SetMatrix("modelViewProjection",&modelViewProjMatrix));
		V(myEffect->SetFloat("dsWidth", desc.Width));
		V(myEffect->SetFloat("dsHeight", desc.Height));
		myEffect->CommitChanges();
		V(myEffect->Begin(&numberofPasses,0));
		for(UINT i = 0; i<numberofPasses; i++ )
		{
			V(myEffect->BeginPass(i));

			V(pd3dDevice->SetStreamSource(0,pQuadVB,0,sizeof(D3DVERTEX_1)));
			V(pd3dDevice->DrawPrimitive(D3DPT_TRIANGLELIST,0,1));

			V(myEffect->EndPass());
		}
		V(myEffect->End());
    }

	// buffer0->buffer1
	V( pd3dDevice->SetRenderTarget(0,myAverageLuminanceTextureSurface[1]));
	V( pd3dDevice->SetDepthStencilSurface(rttDepthSurface) ); 
	V( pd3dDevice->SetVertexDeclaration(textureDecl) );
	V( pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_ARGB(255, 0, 0, 0), 1.0f, 0) );
	
    {
		UINT numberofPasses=0;
		D3DXHANDLE techniqueHandle = myEffect->GetTechniqueByName("Blur_v");
		V(myEffect->SetTechnique(techniqueHandle));
		V(myEffect->SetTexture("LuminanceTexture",myAverageLuminanceTexture[0]));
		V(myEffect->SetMatrix("modelViewProjection",&modelViewProjMatrix));
		V(myEffect->SetFloat("dsWidth", desc.Width));
		V(myEffect->SetFloat("dsHeight", desc.Height));
		myEffect->CommitChanges();
		V(myEffect->Begin(&numberofPasses,0));
		for(UINT i = 0; i<numberofPasses; i++ )
		{
			V(myEffect->BeginPass(i));
			
			V(pd3dDevice->SetStreamSource(0,pQuadVB,0,sizeof(D3DVERTEX_1)));
			V(pd3dDevice->DrawPrimitive(D3DPT_TRIANGLELIST,0,1));

			V(myEffect->EndPass());
		}
		V(myEffect->End());
    }
	}

	// Step 3: Final composition
	V( pd3dDevice->SetRenderTarget(0,backBuffer));
	V( pd3dDevice->SetDepthStencilSurface(backBufferDepthSurface) );
	V( pd3dDevice->SetVertexDeclaration(textureDecl) );
	V( pd3dDevice->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_ARGB(255, 0, 0, 0), 1.0f, 0) );

    {
		UINT numberofPasses=0;
		D3DXHANDLE techniqueHandle = myEffect->GetTechniqueByName("FinalTechnique");
		V(myEffect->SetTechnique(techniqueHandle));
		V(myEffect->SetTexture("SourceTexture",sourceTexture));
		V(myEffect->SetTexture("LuminanceTexture",myLuminanceTexture));
		V(myEffect->SetTexture("AverageLuminanceTexture", myAverageLuminanceTexture[1]));
		V(myEffect->SetMatrix("modelViewProjection",&modelViewProjMatrix));
		V(myEffect->SetFloat("Gain", Gain));
		V(myEffect->SetFloat("dsWidth", desc.Width));
		V(myEffect->SetFloat("dsHeight", desc.Height));
		myEffect->CommitChanges();
		V(myEffect->Begin(&numberofPasses,0));
		for(UINT i = 0; i<numberofPasses; i++ )
		{
			V(myEffect->BeginPass(i));

			hr=pd3dDevice->SetRenderState(D3DRS_CULLMODE,D3DCULL_NONE);
			V(pd3dDevice->SetStreamSource(0,pQuadVB,0,sizeof(D3DVERTEX_1)));
			V(pd3dDevice->DrawPrimitive(D3DPT_TRIANGLELIST,0,1));

			V(myEffect->EndPass());
		}
		V(myEffect->End());
		
        V( pd3dDevice->EndScene() );
    }

	SAFE_RELEASE(backBuffer);
	SAFE_RELEASE(backBufferDepthSurface);
}