#include "ibl_common.hlsli"

struct Block
{
    uint OutputTextureHandle;
};

[[vk::push_constant]]
Block Pcb;

float GeometrySchlickGGX(float NdotV, float Roughness)
{
    float R = Roughness;
    // (Rough + 1)^2 / 8 for Punctual Lights
    // Rough^2 / 2 for IBL
    float K = (R * R) / 2.0;

    float Numerator = NdotV;
    float Denominator = NdotV * (1.0f - K) + K;
	
    return Numerator / Denominator;
}

float GeometrySmith(float NdotV, float NdotL, float Roughness)
{
    float GGX1 = GeometrySchlickGGX(NdotV, Roughness);
    float GGX2 = GeometrySchlickGGX(NdotL, Roughness);
	
    return GGX1 * GGX2;
}
          
float2 IntegrateBRDF(float NdotV, float Roughness)
{
    float3 ViewDir;
    ViewDir.x = sqrt(1.0f - NdotV * NdotV);
    ViewDir.y = 0.0f;
    ViewDir.z = NdotV;

    float A = 0.0f;
    float B = 0.0f;

    float3 Normal = float3(0.0f, 0.0f, 1.0f);

    const uint SAMPLE_COUNT = 1024u;
    
    for (uint i = 0u; i < SAMPLE_COUNT; ++i)
    {
        float2 Xi = Hammersley(i, SAMPLE_COUNT);
        float3 Halfway = ImportanceSampleGGX(Xi, Normal, Roughness);
        float3 LightDir = normalize(2.0f * dot(ViewDir, Halfway) * Halfway - ViewDir);

        float NdotL = max(LightDir.z, 0.0f);
        float NdotH = max(Halfway.z, 0.0f);
        float VdotH = max(dot(ViewDir, Halfway), 0.0f);

        if (NdotL > 0.0f)
        {
            float G = GeometrySmith(NdotV, NdotL, Roughness);
            float G_Vis = (G * VdotH) / max((NdotH * NdotV), 0.0001f);
            float Fc = pow(1.0f - VdotH, 5.0f);

            A += (1.0f - Fc) * G_Vis;
            B += Fc * G_Vis;
        }
    }
    A /= float(SAMPLE_COUNT);
    B /= float(SAMPLE_COUNT);
    
    return float2(A, B);
}

[numthreads(16, 16, 1)]
void main(uint3 GlobalInvocationID : SV_DispatchThreadID)
{
    float Width, Height;
    StorageTexturesRG[Pcb.OutputTextureHandle].GetDimensions(Width, Height);
    
    float2 UV = GlobalInvocationID.xy / float2(Width, Height);
    
    float2 IntegratedBRDF = IntegrateBRDF(UV.x, UV.y);
    StorageTexturesRG[Pcb.OutputTextureHandle][GlobalInvocationID.xy] = IntegratedBRDF;
}