#include "Macros.fxh"

DECLARE_TEXTURE(AlbedoTexture, 0);
DECLARE_TEXTURE(NormalTexture, 1);
DECLARE_TEXTURE(MetallicRoughnessTexture, 2);

float4 UVOffsetAndDimensions;

BEGIN_CONSTANTS

    float3 AlbedoValue                         _ps(c0)    _cb(c0);
    float MetallicValue                        _ps(c1)    _cb(c1);
    float RoughnessValue                       _ps(c2)    _cb(c2);

    int NumTextureRows              _vs(c9)               _cb(c3);
    int NumTextureColumns           _vs(c10)              _cb(c4);

MATRIX_CONSTANTS

    float4x4 World                  _vs(c0)               _cb(c7);
    float4x4 ViewProjection         _vs(c4)               _cb(c11);

END_CONSTANTS

struct VertexInput
{
    float4 Position : POSITION;
    float3 Normal   : NORMAL;
    float2 TexCoord : TEXCOORD;
};

struct PixelInput
{
    float4 Position      : SV_POSITION;
    float3 PositionWorld : TEXCOORD0;
    float3 NormalWorld   : TEXCOORD1;
    float2 TexCoord      : TEXCOORD2;
};

struct PixelOutput
{
    float4 gPosition          : COLOR0;
    float4 gNormal            : COLOR1;
    float4 gAlbedo            : COLOR2;
    float4 gMetallicRoughness : COLOR3;
};

// Vertex Shader

PixelInput main_vs(VertexInput input)
{
    PixelInput output;

    output.PositionWorld = mul(input.Position, World).xyz;
    output.NormalWorld = normalize(mul(input.Normal, World));

    float2 texCoord;
    texCoord.x = (input.TexCoord.x * UVOffsetAndDimensions.z) + UVOffsetAndDimensions.x;
    texCoord.y = (input.TexCoord.y * UVOffsetAndDimensions.w) + UVOffsetAndDimensions.y;
    output.TexCoord = texCoord;

    float4x4 worldViewProjection = mul(World, ViewProjection);
    output.Position = mul(input.Position, worldViewProjection);

    return output;
}

PixelInput instanced_vs(
    VertexInput input,
    float3 Translation : TEXCOORD2,
    float2 UVOffset    : TEXCOORD5
) {
    PixelInput output;

    float4x4 world = float4x4(
        float4(1, 0, 0, 0),
        float4(0, 1, 0, 0),
        float4(0, 0, 1, 0),
        float4(Translation.x, Translation.y, Translation.z, 1)
    );

    float4x4 worldViewProjection = mul(world, ViewProjection);

    output.PositionWorld = mul(input.Position, world);
    output.NormalWorld = mul(input.Normal, world);

    float2 texCoord;
    texCoord.x = (input.TexCoord.x / NumTextureColumns + UVOffset.x);
    texCoord.y = (input.TexCoord.y / NumTextureRows + UVOffset.y);
    output.TexCoord = texCoord;

    output.Position = mul(input.Position, worldViewProjection);

    return output;
}

// Pixel Shaders

// Easy trick to get tangent-normals to world-space to keep PBR code simplified.
float3 GetNormalFromMap(float3 worldPos, float2 texCoords, float3 normal)
{
    float3 tangentNormal = SAMPLE_TEXTURE(NormalTexture, texCoords).xyz * 2.0 - 1.0;

    float3 Q1  = ddx(worldPos);
    float3 Q2  = ddy(worldPos);
    float2 st1 = ddx(texCoords);
    float2 st2 = ddy(texCoords);

    float3 N   = normalize(normal);
    float3 T  = normalize(Q1*st2.y - Q2*st1.y);
    float3 B  = -normalize(cross(N, T));
    float3x3 TBN = float3x3(T, B, N);

    return normalize(mul(tangentNormal, TBN));
}

PixelOutput NonePS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(normalize(input.NormalWorld), 1.0);
    output.gAlbedo = float4(AlbedoValue, 1.0);
    output.gMetallicRoughness = float4(MetallicValue, RoughnessValue, 0.0, 1.0);

    return output;
}

PixelOutput AlbedoPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(normalize(input.NormalWorld), 1.0);
    output.gAlbedo = SAMPLE_TEXTURE(AlbedoTexture, input.TexCoord);
    output.gMetallicRoughness = float4(MetallicValue, RoughnessValue, 0.0, 1.0);

    if (output.gAlbedo.a == 0.0) { discard; }

    return output;
}

PixelOutput MetallicRoughnessPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(normalize(input.NormalWorld), 1.0);
    output.gAlbedo = float4(AlbedoValue, 1.0);
    output.gMetallicRoughness = SAMPLE_TEXTURE(MetallicRoughnessTexture, input.TexCoord);

    return output;
}

PixelOutput NormalPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(GetNormalFromMap(input.PositionWorld, input.TexCoord, input.NormalWorld), 1.0);
    output.gAlbedo = float4(AlbedoValue, 1.0);
    output.gMetallicRoughness = float4(MetallicValue, RoughnessValue, 0.0, 1.0);

    return output;
}

PixelOutput AlbedoMetallicRoughnessPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(normalize(input.NormalWorld), 1.0);
    output.gAlbedo = SAMPLE_TEXTURE(AlbedoTexture, input.TexCoord);
    output.gMetallicRoughness = SAMPLE_TEXTURE(MetallicRoughnessTexture, input.TexCoord);

    if (output.gAlbedo.a == 0.0) { discard; }

    return output;
}

PixelOutput AlbedoNormalPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(GetNormalFromMap(input.PositionWorld, input.TexCoord, input.NormalWorld), 1.0);
    output.gAlbedo = SAMPLE_TEXTURE(AlbedoTexture, input.TexCoord);
    output.gMetallicRoughness = float4(MetallicValue, RoughnessValue, 0.0, 1.0);

    if (output.gAlbedo.a == 0.0) { discard; }

    return output;
}

PixelOutput MetallicRoughnessNormalPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(GetNormalFromMap(input.PositionWorld, input.TexCoord, input.NormalWorld), 1.0);
    output.gAlbedo = float4(AlbedoValue, 1.0);
    output.gMetallicRoughness = SAMPLE_TEXTURE(MetallicRoughnessTexture, input.TexCoord);

    return output;
}

PixelOutput AlbedoMetallicRoughnessNormalMapPS(PixelInput input)
{
    PixelOutput output;

    output.gPosition = float4(input.PositionWorld, 1.0);
    output.gNormal = float4(GetNormalFromMap(input.PositionWorld, input.TexCoord, input.NormalWorld), 1.0);
    output.gAlbedo = SAMPLE_TEXTURE(AlbedoTexture, input.TexCoord);
    output.gMetallicRoughness = SAMPLE_TEXTURE(MetallicRoughnessTexture, input.TexCoord);

    if (output.gAlbedo.a == 0.0) { discard; }

    return output;
}

VertexShader VSArray[2] =
{
    compile vs_3_0 main_vs(),
    compile vs_3_0 instanced_vs()
};

int VSIndices[2] =
{
    0, 1
};

int VertexShaderIndex = 0;

PixelShader PSArray[8] =
{
    compile ps_3_0 NonePS(),

    compile ps_3_0 AlbedoPS(),
    compile ps_3_0 MetallicRoughnessPS(),
    compile ps_3_0 NormalPS(),

    compile ps_3_0 AlbedoMetallicRoughnessPS(),
    compile ps_3_0 AlbedoNormalPS(),
    compile ps_3_0 MetallicRoughnessNormalPS(),

    compile ps_3_0 AlbedoMetallicRoughnessNormalMapPS()
};

int PSIndices[8] =
{
    0, 1, 2, 3, 4, 5, 6, 7
};

int PixelShaderIndex = 0;

Technique GBuffer
{
    Pass
    {
        VertexShader = (VSArray[VSIndices[VertexShaderIndex]]);
        PixelShader  = (PSArray[PSIndices[PixelShaderIndex]]);
    }
}