Deferred Shading

#pragma once

#include "Shader.h"

class DeferredRenderer
{
private:
  struct QuadVertex
  {
    float x, y, z;
    float tx, ty;

    QuadVertex( float ax, float ay, float az, float atx, float aty )
    {
      x = ax;
      y = ay;
      z = az;

      tx = atx;
      ty = aty;
    }
  };

  IDirect3DTexture9 * positionMap;
  IDirect3DTexture9 * normalMap;
  IDirect3DTexture9 * diffuseMap;

  IDirect3DSurface9 * positionSurface;
  IDirect3DSurface9 * normalSurface;
  IDirect3DSurface9 * diffuseSurface;
  IDirect3DSurface9 * backSurface;

  VertexShader * vertexShaderPassOne;
  PixelShader * pixelShaderPassOne;

  VertexShader * vertexShaderPassTwo;
  PixelShader * pixelShaderPassTwo;

  IDirect3DVertexBuffer9 * screenQuad;
  IDirect3DVertexDeclaration9 * vertexDeclaration;

  D3DXHANDLE v1World;
  D3DXHANDLE v1Proj;
  D3DXHANDLE v1View;

  D3DXHANDLE v2World;
  D3DXHANDLE v2Proj;
  D3DXHANDLE v2View;

  D3DXHANDLE p2LightPos;
  D3DXHANDLE p2LightRange;
  D3DXHANDLE p2CameraPos;
public:

  void CreateScreenQuad()
  {
    g_device->CreateVertexBuffer( 6 * sizeof( QuadVertex ), D3DUSAGE_WRITEONLY, D3DFVF_XYZ | D3DFVF_TEX1, D3DPOOL_DEFAULT, &screenQuad, 0 );

    QuadVertex vertices[ ] = { QuadVertex( -0.5, -0.5, 0, 0, 0 ), QuadVertex( g_width - 0.5, 0 - 0.5, 0, 1, 0 ), QuadVertex ( 0 - 0.5, g_height - 0.5, 0, 0, 1 ),
      QuadVertex( g_width - 0.5, 0 - 0.5, 0, 1, 0 ), QuadVertex( g_width - 0.5, g_height - 0.5, 0, 1, 1 ), QuadVertex ( 0 - 0.5, g_height - 0.5, 0, 0, 1 ) };

    void * data = 0;

    screenQuad->Lock( 0, 0, &data, 0 );
    memcpy( data, vertices, sizeof( QuadVertex ) * 6 );
    screenQuad->Unlock( );

    D3DVERTEXELEMENT9 quadVertexDeclation[ ] =
    {
      { 0,  0, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0 },
      { 0, 12, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0 },
      D3DDECL_END()
    };

    g_device->CreateVertexDeclaration( quadVertexDeclation, &vertexDeclaration );
  }

  void CreateRenderTargets()
  {
    D3DCAPS9 caps;  g_device->GetDeviceCaps( &caps );

    if( caps.NumSimultaneousRTs < 4 )
      MessageBoxA( 0, "Your video card does not support MRT.", 0, MB_OK | MB_ICONERROR );

    D3DFORMAT format = D3DFMT_A16B16G16R16F;

    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, format, D3DPOOL_DEFAULT, &positionMap )))
      MessageBoxA( 0, "Failed to create 'PositionMap' texture.", 0, MB_OK | MB_ICONERROR );
    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, format, D3DPOOL_DEFAULT, &normalMap )))
      MessageBoxA( 0, "Failed to create 'PositionMap' texture.", 0, MB_OK | MB_ICONERROR );
    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, format, D3DPOOL_DEFAULT, &diffuseMap )))
      MessageBoxA( 0, "Failed to create 'PositionMap' texture.", 0, MB_OK | MB_ICONERROR );

    positionMap->GetSurfaceLevel( 0, &positionSurface );
    normalMap->GetSurfaceLevel( 0, &normalSurface );
    diffuseMap->GetSurfaceLevel( 0, &diffuseSurface );

    g_device->GetRenderTarget( 0, &backSurface );
  }

  void InitPassOneShaders()
  {
    // PASS 1 - Filling the G-Buffer
    string vertexSourcePassOne =
      "float4x4 g_world;\n"
      "float4x4 g_view;\n"
      "float4x4 g_projection;\n"

      "struct VS_INPUT {\n"
      "  float4 position : POSITION;\n"
      "  float4 normal   : NORMAL;\n"
      "  float2 texCoord : TEXCOORD0;\n"
      "};\n"

      "struct VS_OUTPUT {\n"
      "  float4 position : POSITION;\n"
      "  float3 worldPos : TEXCOORD0;\n"
      "  float3 normal : TEXCOORD1;\n"
      "  float2 diffuseMapCoord : TEXCOORD2;\n"
      "};\n"

      "VS_OUTPUT main(VS_INPUT input) {\n"
      "  VS_OUTPUT output;\n"
      "  float4x4 worldViewProj = mul(mul(g_world, g_view), g_projection);\n"
      "  output.position = mul(input.position, worldViewProj);\n"
      "  output.normal = normalize(mul(input.normal, (float3x3)g_world));\n"
      "  output.worldPos = mul(input.position, g_world);\n"
      "  output.diffuseMapCoord = input.texCoord;\n"
      "  return output;\n"
      "};\n";

    vertexShaderPassOne = new VertexShader( vertexSourcePassOne );

    v1World = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_world" );
    v1Proj = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_projection" );
    v1View = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_view" );

    string pixelSourcePassOne =

      "texture diffuseMap;\n"

      "sampler diffuseSampler : register(s0) = sampler_state\n"
      "{\n"
      "   texture = <diffuseMap>;\n"
      "};\n"

      "struct PS_INPUT {\n"
      "  float3 worldPos : TEXCOORD0;\n"
      "  float3 normal : TEXCOORD1;\n"
      "  float2 diffuseMapCoord : TEXCOORD2;\n"
      "};\n"

      "struct MRT_OUTPUT {\n"
      "  float4 position : COLOR0;\n"
      "  float4 normal : COLOR1;\n"
      "  float4 diffuseMap : COLOR2;\n"
      "};\n"

      "MRT_OUTPUT main( PS_INPUT input )\n"
      "{\n"
      "   MRT_OUTPUT output;\n"
      "   float3 normal = normalize( input.normal );\n"
      "   output.position = float4( input.worldPos.x, input.worldPos.y, input.worldPos.z, 1.0 );\n"
      "   output.normal = float4( normal.x, normal.y, normal.z, 1.0 );\n"
      "   output.diffuseMap = tex2D( diffuseSampler, input.diffuseMapCoord );\n"
      "   return output;"
      "};\n"
      ;

    pixelShaderPassOne = new PixelShader( pixelSourcePassOne );
  }

  DeferredRenderer()
  {
    CreateScreenQuad();
    CreateRenderTargets();
    InitPassOneShaders();
    InitPassTwoShaders();
  }

  void InitPassTwoShaders()
  {
    // PASS 2 - Final render
    string vertexSourcePassTwo =
      "float4x4 g_world;\n"
      "float4x4 g_view;\n"
      "float4x4 g_projection;\n"
      "struct VS_INPUT {\n"
      "  float4 position : POSITION;\n"
      "  float3 texcoord : TEXCOORD0;\n"
      "};\n"

      "struct VS_OUTPUT {\n"
      "  float4 position : POSITION;\n"
      "  float3 texcoord : TEXCOORD0;\n"
      "};\n"

      "VS_OUTPUT main(VS_INPUT input) {\n"
      "  VS_OUTPUT output;\n"
      "  float4x4 worldViewProj = mul(mul(g_world, g_view), g_projection);\n"
      "  output.position   = mul(input.position, worldViewProj);\n"
      "  output.texcoord   = input.texcoord;\n"
      "  return output;\n"
      "};\n";

    vertexShaderPassTwo = new VertexShader( vertexSourcePassTwo );

    v2World = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_world" );
    v2Proj = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_projection" );
    v2View = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_view" );

    string pixelSourcePassTwo =

      "texture positionMap;\n"
      "texture normalMap;\n"
      "texture diffuseMap;\n"

      "sampler positionSampler : register(s0) = sampler_state\n"
      "{\n"
      "   texture = <positionMap>;\n"
      "};\n"

      "sampler normalSampler : register(s1) = sampler_state\n"
      "{\n"
      "   texture = <normalMap>;\n"
      "};\n"

      "sampler diffuseSampler : register(s2) = sampler_state\n"
      "{\n"
      "   texture = <diffuseMap>;\n"
      "};\n"

      // light props
      "float3 lightPos;\n"
      "float lightRange;\n"

      // camera props
      "float3 cameraPosition;\n"

      "float4 main( float3 texcoord : TEXCOORD0 ) : COLOR0\n"
      "{\n"
      "   float4 diffuseTexel = tex2D( diffuseSampler, texcoord.xy );\n"
      "   float3 p = tex2D( positionSampler, texcoord.xy ).xyz;\n"
      "   float3 n = tex2D( normalSampler, texcoord.xy ).xyz;\n"

      "   float3 lightDirection = lightPos - p;"
      "   float3 l = normalize( lightDirection );\n"

      // specular
      "   float3 v = normalize( cameraPosition - p );\n"
      "   float3 r = reflect( -v, n );\n"
      "   float spec = pow( saturate( dot( l, r ) ), 20.0 );\n"

      // diffuse
      "   float diff = saturate(dot( l, n ));\n"

      "   float falloff = lightRange / dot( lightDirection, lightDirection );\n"

      "   float o = falloff * ( diff + spec );\n"

      "   return float4( diffuseTexel.x * o, diffuseTexel.y * o, diffuseTexel.z * o, diffuseTexel.w );\n"
      "};\n";

    pixelShaderPassTwo = new PixelShader( pixelSourcePassTwo );

    p2LightPos = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "lightPos" );
    p2LightRange = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "lightRange" );
    p2CameraPos = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "cameraPosition" );
  }

  ~DeferredRenderer()
  {
    if( positionMap )
      positionMap->Release();
    if( normalMap )
      normalMap->Release();
    if( diffuseMap )
      diffuseMap->Release();

    if( positionSurface )
      positionSurface->Release();
    if( normalSurface )
      normalSurface->Release();
    if( diffuseSurface )
      diffuseSurface->Release();
  }

  void Begin()
  {
    g_device->SetRenderTarget( 0, positionSurface );
    g_device->SetRenderTarget( 1, normalSurface );
    g_device->SetRenderTarget( 2, diffuseSurface );

    g_device->Clear( 0, 0, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER , D3DCOLOR_XRGB( 0, 0, 0 ), 1.0, 0 );
    g_device->BeginScene( );

    vertexShaderPassOne->Bind();
    pixelShaderPassOne->Bind();

    g_device->SetRenderState(D3DRS_ALPHABLENDENABLE, FALSE );
  }

  void RenderMesh( Mesh * mesh )
  {
    g_device->SetStreamSource( 0, mesh->vb, 0, sizeof( Vertex ));
    g_device->SetIndices( mesh->ib );

    D3DXMATRIX view; g_device->GetTransform( D3DTS_VIEW, &view );
    D3DXMATRIX proj; g_device->GetTransform( D3DTS_PROJECTION, &proj );
    D3DXMATRIX world; GetD3DMatrixFromBulletTransform( mesh->parent->globalTransform, world );

    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1View, &view );
    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1Proj, &proj );
    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1World, &world );

    g_device->DrawIndexedPrimitive( D3DPT_TRIANGLELIST, 0, 0, mesh->vertexCount, 0, mesh->faceCount );
  }

  void End()
  {
    g_device->EndScene();

    g_device->SetRenderTarget( 0, backSurface );
    g_device->SetRenderTarget( 1, 0 );
    g_device->SetRenderTarget( 2, 0 );

    g_device->Clear( 0, 0, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER , D3DCOLOR_XRGB( 0, 0, 0 ), 1.0, 0 );
    g_device->BeginScene( );

    g_device->SetVertexDeclaration( vertexDeclaration );

    vertexShaderPassTwo->Bind();
    pixelShaderPassTwo->Bind();

    g_device->SetRenderState( D3DRS_LIGHTING, FALSE );
    g_device->SetRenderState( D3DRS_CULLMODE, D3DCULL_NONE );

    D3DXMATRIX ortho; D3DXMatrixOrthoOffCenterLH ( &ortho, 0, g_width, g_height, 0, 0, 1024 );
    D3DXMATRIX identity; D3DXMatrixIdentity( &identity );

    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2World, &identity );
    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2View, &identity );
    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2Proj, &ortho );

    g_device->SetTexture( 0, positionMap );
    g_device->SetTexture( 1, normalMap );
    g_device->SetTexture( 2, diffuseMap );

    g_device->SetStreamSource( 0, screenQuad, 0, sizeof( QuadVertex ));

    g_device->SetRenderState(D3DRS_ALPHABLENDENABLE, TRUE);
    g_device->SetRenderState(D3DRS_SRCBLEND, D3DBLEND_ONE );
    g_device->SetRenderState(D3DRS_DESTBLEND, D3DBLEND_ONE );

    float temp[ 3 ];

    btVector3 lightPos = g_camera->base->globalTransform.getOrigin();

    temp[ 0 ] = lightPos.x();
    temp[ 1 ] = lightPos.y();
    temp[ 2 ] = lightPos.z();

    pixelShaderPassTwo->GetConstantTable()->SetFloatArray( g_device, p2CameraPos, temp, 3 );

    for( int i = 0; i < g_lights.size(); i++ )
    {
      Light * light = g_lights.at( i );

      btVector3 lightPos = light->base->globalTransform.getOrigin();

      temp[ 0 ] = lightPos.x();
      temp[ 1 ] = lightPos.y();
      temp[ 2 ] = lightPos.z();

      pixelShaderPassTwo->GetConstantTable()->SetFloatArray( g_device, p2LightPos, temp, 3 );
      pixelShaderPassTwo->GetConstantTable()->SetFloat( g_device, p2LightRange, light->radius );


      g_device->DrawPrimitive( D3DPT_TRIANGLELIST, 0, 2 );
    }

    g_device->SetTexture( 0, 0 );
    g_device->SetTexture( 1, 0 );
    g_device->SetTexture( 2, 0 );
  }
};