#include "Main.h"

ShadowMap::ShadowMap()
{
    _DepthStencil = NULL;
}

ShadowMap::~ShadowMap()
{
    FreeMemory();
}

void ShadowMap::FreeMemory()
{
    for(UINT Index = 0; Index < ShadowMapCount; Index++)
    {
        _DepthMapTextures[Index].FreeMemory();
    }
    _DepthMapBlurTexture.FreeMemory();
    if(_DepthStencil != NULL)
    {
        _DepthStencil->Release();
        _DepthStencil = NULL;
    }
    
    _VShaderCreate.FreeMemory();
    _PShaderCreate.FreeMemory();

    _VShaderRender.FreeMemory();
    _PShaderRender.FreeMemory();

    _VShaderGaussian.FreeMemory();
    _PShaderGaussianX.FreeMemory();
    _PShaderGaussianY.FreeMemory();
}

void ShadowMap::SetAsRenderTarget(UINT ShadowMapIndex, IDirect3DDevice9 *Device)
{
    _DepthMapTextures[ShadowMapIndex].SetAsRenderTarget(0);
    Device->SetDepthStencilSurface(_DepthStencil);

    Device->Clear( 0, NULL, D3DCLEAR_ZBUFFER | D3DCLEAR_TARGET, D3DCOLOR_ARGB(0, 255, 0, 0), 1.0f, 0 );

    _VShaderCreate.Set();
    _PShaderCreate.Set();
}

void ShadowMap::SetAsTextures(IDirect3DDevice9 *Device)
{
    for(UINT ShadowMapIndex = 0; ShadowMapIndex < ShadowMapCount; ShadowMapIndex++)
    {
        _DepthMapTextures[ShadowMapIndex].Set(ShadowMapIndex + 1);
    }

    _VShaderRender.Set();
    _PShaderRender.Set();
}

void ShadowMap::SetCameraMatrices(const MatrixController &MC, ShadowMapShaderType Shader)
{
    if(Shader == ShadowMapShaderCreate)
    {
        Matrix4 WorldView = MC.World * MC.View;
        Matrix4 WorldViewProj = WorldView * MC.Perspective;
        _VShaderCreate.SetMatrix("WorldView", WorldView);
        _VShaderCreate.SetMatrix("WorldViewProj", WorldViewProj);
    }
    else if(Shader == ShadowMapShaderRender)
    {
        _VShaderRender.SetMatrix("World", MC.World);
        _VShaderRender.SetMatrix("WorldViewProj", MC.TotalMatrix());
    }
}

void ShadowMap::SetLightMatrices(UINT LightIndex, const Camera &LightCamera, const Matrix4 &Perspective)
{
    Matrix4 View = LightCamera.Matrix();
    Matrix4 ViewProj = View * Perspective;

    _StringLightView[_StringLightView.Length() - 1] = '0' + LightIndex;
    _StringLightViewProj[_StringLightViewProj.Length() - 1] = '0' + LightIndex;
    _StringLightDir[_StringLightDir.Length() - 1] = '0' + LightIndex;

    _VShaderRender.SetMatrix(_StringLightView.CString(), View);
    _VShaderRender.SetMatrix(_StringLightViewProj.CString(), ViewProj);

    _PShaderRender.SetVec3(_StringLightDir.CString(), LightCamera.VecLookDir());

    if(LightIndex == 0)
    {
        _PShaderRender.SetVec3("LightPos", LightCamera.VecEye());
    }
}

void ShadowMap::Init(D3D9GraphicsDevice &GD, UINT Resolution)
{
    FreeMemory();

    _StringLightDir = "LightDir0";
    _StringLightView = "LightView0";
    _StringLightViewProj = "LightViewProj0";
    
    for(UINT Index = 0; Index < ShadowMapCount; Index++)
    {
        _DepthMapTextures[Index].Init(GD);
        _DepthMapTextures[Index].Allocate(D3DFMT_G32R32F, Resolution, Resolution, true, 1);
    }
    
    _DepthMapBlurTexture.Init(GD);
    _DepthMapBlurTexture.Allocate(D3DFMT_G32R32F, Resolution, Resolution, true, 1);
    
    D3DAlwaysValidate( GD.GetDevice()->CreateDepthStencilSurface(Resolution, Resolution, D3DFMT_D24S8, D3DMULTISAMPLE_NONE, 0, TRUE, &_DepthStencil, NULL), "CreateDepthStencilSurface" );
    
    _VShaderCreate.Init(GD, "Assets\\ShadowMapCreate.vs");
    _PShaderCreate.Init(GD, "Assets\\ShadowMapCreate.ps");

    _VShaderRender.Init(GD, "Assets\\ShadowMapRender.vs");
    _PShaderRender.Init(GD, "Assets\\ShadowMapRender.ps");

    _VShaderGaussian.Init(GD, "Assets\\GaussianBlur.vs");
    _PShaderGaussianX.Init(GD, "Assets\\GaussianBlurX.ps");
    _PShaderGaussianY.Init(GD, "Assets\\GaussianBlurY.ps");

    _SimpleQuadMesh.SetGD(GD);
    _SimpleQuadMesh.CreateCanonicalRenderPlane();
}

void LoadChannel(const Grid<Vec2f> &Input, Grid<float> &Output)
{
    Output.Allocate(Input.Rows(), Input.Cols());
    for(UINT Row = 0; Row < Input.Rows(); Row++)
    {
        for(UINT Col = 0; Col < Input.Cols(); Col++)
        {
            Output(Row, Col) = Input(Row, Col).x;
        }
    }
}

void ShadowMap::SaveTextureToFile(D3D9Texture &T, const String &Filename)
{
    Grid<Vec2f> Data;
    Grid<float> SingleChannel;
    Bitmap Bmp;

    T.ReadData(Data);
    LoadChannel(Data, SingleChannel);
    Bmp.LoadGrid(SingleChannel);
    Bmp.SavePNG(Filename);
}

void ShadowMap::BlurAllShadowMaps(IDirect3DDevice9 *Device)
{
    const bool DumpShadowMap = false;
    const UINT DumpShadowMapIndex = 1;
    
    D3D9ProtectRenderTarget Protector(Device, true, true);

    
    if(DumpShadowMap)
    {
        SaveTextureToFile(_DepthMapTextures[DumpShadowMapIndex], "Input.png");
    }

    Device->SetRenderState(D3DRS_ZENABLE, D3DZB_FALSE);
    Device->SetSamplerState(0, D3DSAMP_MIPFILTER, D3DTEXF_POINT);
	Device->SetSamplerState(0, D3DSAMP_MINFILTER, D3DTEXF_POINT);
	Device->SetSamplerState(0, D3DSAMP_MAGFILTER, D3DTEXF_POINT);
    Device->SetSamplerState(0, D3DSAMP_ADDRESSU, D3DTADDRESS_CLAMP);
    Device->SetSamplerState(0, D3DSAMP_ADDRESSV, D3DTADDRESS_CLAMP);

    for(UINT Index = 0; Index < ShadowMapCount; Index++)
    {
        _DepthMapTextures[Index].Set(0);
        _DepthMapBlurTexture.SetAsRenderTarget(0);
        Device->SetDepthStencilSurface(_DepthStencil);

        _VShaderGaussian.Set();
        _PShaderGaussianX.Set();

        _SimpleQuadMesh.Render();

        if(DumpShadowMap && Index == DumpShadowMapIndex)
        {
            SaveTextureToFile(_DepthMapBlurTexture, "Intermediate.png");
        }

        _DepthMapTextures[Index].SetAsRenderTarget(0);
        _DepthMapBlurTexture.Set(0);
    
        _PShaderGaussianY.Set();

        _SimpleQuadMesh.Render();

        if(DumpShadowMap && Index == DumpShadowMapIndex)
        {
            SaveTextureToFile(_DepthMapTextures[DumpShadowMapIndex], "Output.png");
        }

        Device->SetRenderState(D3DRS_ZENABLE, D3DZB_TRUE);
    }

    /*Console::File() << "Start\n";
    for(UINT Row = 0; Row < Data0.Rows(); Row++)
    {
        for(UINT Col = 0; Col < Data0.Cols(); Col++)
        {
            if(Data0.GetElement(Row, Col).x - Data1.GetElement(Row, Col).x != 0.0f)
            {
                Console::File() << Data0.GetElement(Row, Col).x << '\t' << Data1.GetElement(Row, Col).x << endl;
            }
        }
    }
    Console::File() << "End\n";*/
}

void ShadowMap::SaveDebugBitmap(const String &Filename)
{
    Grid<float> Data;
    Bitmap Bmp;
    _DepthMapTextures[0].ReadData(Data);
    Bmp.LoadGrid(Data);
    Bmp.SavePNG(Filename);
}