/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "DXHelper.h"
#include "TemporalBlendPass.h"
#include "TemporalAA.h"

// Constant buffer
struct ConstantBuf
{
    float param[4];
};

/**
*/
TemporalBlendPass::TemporalBlendPass() : m_pQuadLayout( NULL ), m_pVB( NULL ),
                                         m_pQuadVS( NULL ), m_pQuadPS( NULL ), m_pSampleState( NULL ), m_pConstantBuf( NULL ),
                                         m_pSceneTexturePrevSRV( NULL ), m_pSceneTextureCurSRV( NULL )
{
}

/**
*/
TemporalBlendPass::~TemporalBlendPass()
{
}

/**
    @param pd3dDevice    [in]    D3D Device
*/
HRESULT TemporalBlendPass::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    ID3DBlob* pBlob = NULL;

    // vertex shader
    V_RETURN( CompileShaderFromFile( L"TemporalBlendPass.hlsl", "QuadVS", "vs_4_0", &pBlob ) );
    V_RETURN( pd3dDevice->CreateVertexShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pQuadVS ) );
    
    // input layout
    const D3D11_INPUT_ELEMENT_DESC quadlayout[] =
    {
        { "POSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
        { "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, 0, 16, D3D11_INPUT_PER_VERTEX_DATA, 0 },
    };
    V_RETURN( pd3dDevice->CreateInputLayout( quadlayout, 2, pBlob->GetBufferPointer(), pBlob->GetBufferSize(), &m_pQuadLayout ) );
    SAFE_RELEASE( pBlob );    

    // pixel shader
    V_RETURN( CompileShaderFromFile( L"TemporalBlendPass.hlsl", "QuadPS", "ps_4_0", &pBlob ) );
    V_RETURN( pd3dDevice->CreatePixelShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pQuadPS ) );
    SAFE_RELEASE( pBlob );

    // point sampler state
    D3D11_SAMPLER_DESC SamplerDesc;
    ZeroMemory( &SamplerDesc, sizeof(SamplerDesc) );
    SamplerDesc.AddressU = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.AddressV = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_POINT;
    V_RETURN( pd3dDevice->CreateSamplerState( &SamplerDesc, &m_pSampleState ) );

    // full screen quad
    SCREEN_VERTEX svQuad[4];
    svQuad[0].pos = D3DXVECTOR4( -1.0f, 1.0f, 0.5f, 1.0f );
    svQuad[0].tex = D3DXVECTOR2( 0.0f, 0.0f );
    svQuad[1].pos = D3DXVECTOR4( 1.0f, 1.0f, 0.5f, 1.0f );
    svQuad[1].tex = D3DXVECTOR2( 1.0f, 0.0f );
    svQuad[2].pos = D3DXVECTOR4( -1.0f, -1.0f, 0.5f, 1.0f );
    svQuad[2].tex = D3DXVECTOR2( 0.0f, 1.0f );
    svQuad[3].pos = D3DXVECTOR4( 1.0f, -1.0f, 0.5f, 1.0f );
    svQuad[3].tex = D3DXVECTOR2( 1.0f, 1.0f );

    D3D11_BUFFER_DESC vbdesc =
    {
        4 * sizeof( SCREEN_VERTEX ),
        D3D11_USAGE_DEFAULT,
        D3D11_BIND_VERTEX_BUFFER,
        0,
        0
    };
    D3D11_SUBRESOURCE_DATA InitData;
    InitData.pSysMem = svQuad;
    InitData.SysMemPitch = 0;
    InitData.SysMemSlicePitch = 0;
    V_RETURN( pd3dDevice->CreateBuffer( &vbdesc, &InitData, &m_pVB ) );

    // Setup constant buffers
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;    
    Desc.ByteWidth = sizeof( ConstantBuf );
    V_RETURN( pd3dDevice->CreateBuffer( &Desc, NULL, &m_pConstantBuf ) );

    return hr;
}

/**
*/
void TemporalBlendPass::OnDestroyDevice()
{
    SAFE_RELEASE( m_pQuadPS );
    SAFE_RELEASE( m_pSampleState );
    SAFE_RELEASE( m_pVB );
    SAFE_RELEASE( m_pQuadVS );
    SAFE_RELEASE( m_pQuadLayout );
    SAFE_RELEASE( m_pConstantBuf );
}

/**
    @param pd3dDevice                 [in]    D3D Device
    @param pd3dImmediateContext       [in]    D3D Device Immediate Context
    @param pBackBufferDesc            [in]    Description of backbuffer
*/
HRESULT TemporalBlendPass::OnFrameRender( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext, const DXGI_SURFACE_DESC* pBackBufferDesc, 
                                      int nBufWidth, int nBufHeight )
{
    HRESULT hr = S_OK;

    // set constant buffer
    D3D11_MAPPED_SUBRESOURCE MappedResource;            
    V( pd3dImmediateContext->Map( m_pConstantBuf, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
    float* p = reinterpret_cast< float* >( MappedResource.pData );
    p[0] = static_cast< float >( nBufWidth );
    p[1] = static_cast< float >( nBufHeight );
    p[2] = static_cast< float >( g_temporalAA.GetMode() );
    pd3dImmediateContext->Unmap( m_pConstantBuf, 0 );
    ID3D11Buffer* ppCB[1] = { m_pConstantBuf };
    pd3dImmediateContext->PSSetConstantBuffers( 0, 1, ppCB );

    DrawFullScreenQuad( pd3dImmediateContext, m_pQuadPS, pBackBufferDesc->Width, pBackBufferDesc->Height );

    return hr;
}

/**
    @param pd3dImmediateContext       [in]    D3D Device Immediate Context
    @param pPS                        [in]    Pixel Shader
    @param nWidth                     [in]    Backbuffer width
    @param nHeight                    [in]    Backbuffer height
*/
void TemporalBlendPass::DrawFullScreenQuad( ID3D11DeviceContext* pd3dImmediateContext, ID3D11PixelShader* pPS,
                                           UINT nWidth, UINT nHeight )
{
    //D3DPERF_BeginEvent( D3DCOLOR_RGBA( 1, 0, 0, 1 ), L"FullScreenQuad::DrawFullScreenQuad" );

    // Save the old viewport
    D3D11_VIEWPORT vpOld[D3D11_VIEWPORT_AND_SCISSORRECT_MAX_INDEX];
    UINT nViewPorts = 1;
    pd3dImmediateContext->RSGetViewports( &nViewPorts, vpOld );

    // Setup the viewport to match the backbuffer
    D3D11_VIEWPORT vp;
    vp.Width = static_cast< float >( nWidth );
    vp.Height = static_cast< float >( nHeight );
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    vp.TopLeftX = 0;
    vp.TopLeftY = 0;
    pd3dImmediateContext->RSSetViewports( 1, &vp );

    UINT strides = sizeof( SCREEN_VERTEX );
    UINT offsets = 0;
    ID3D11Buffer* pBuffers[1] = { m_pVB };

    pd3dImmediateContext->IASetInputLayout( m_pQuadLayout );
    pd3dImmediateContext->IASetVertexBuffers( 0, 1, pBuffers, &strides, &offsets );
    pd3dImmediateContext->IASetPrimitiveTopology( D3D11_PRIMITIVE_TOPOLOGY_TRIANGLESTRIP );

    ID3D11SamplerState* aSamplers[] = { m_pSampleState };
    pd3dImmediateContext->PSSetSamplers( 0, 1, aSamplers );

    ID3D11ShaderResourceView* aRViews[ 2 ] = { m_pSceneTexturePrevSRV, m_pSceneTextureCurSRV };
    pd3dImmediateContext->PSSetShaderResources( 0, 2, aRViews );

    pd3dImmediateContext->VSSetShader( m_pQuadVS, NULL, 0 );
    pd3dImmediateContext->PSSetShader( pPS, NULL, 0 );
    pd3dImmediateContext->Draw( 4, 0 );

    // Restore the Old viewport
    pd3dImmediateContext->RSSetViewports( nViewPorts, vpOld );

    //
    ID3D11ShaderResourceView* aRViewsNull[ 2 ] = { NULL, NULL };
    pd3dImmediateContext->PSSetShaderResources( 0, 2, aRViewsNull );

    //D3DPERF_EndEvent();
}