import common_structs;

float3 GetCubeDir(uint3 GlobalInvocationID, float SideLength)
{
    float2 FaceUV = float2(GlobalInvocationID.xy) / SideLength; // (0, SideLength) -> (0, 1)
    FaceUV = 2.0f * FaceUV - 1.0f; // (0, 1) -> (-1, 1)
    
    switch (GlobalInvocationID.z)
    {
        case 0:
            return normalize(float3(1.0f, -FaceUV.y, -FaceUV.x)); // Face +X; x = 1, y = -v, z = -u
        case 1:
            return normalize(float3(-1.0f, -FaceUV.y, FaceUV.x)); // Face -X; x = -1, y = -v, z = u
        case 2:
            return normalize(float3(FaceUV.x, 1.0f, FaceUV.y)); // Face +Y; x = u, y = 1, z = v
        case 3:
            return normalize(float3(FaceUV.x, -1.0f, -FaceUV.y)); // Face -Y; x=u, y=-1, z=-v
        case 4:
            return normalize(float3(FaceUV.x, -FaceUV.y, 1.0f)); // Face +Z; x=u,y=-v, z=1
        case 5:
            return normalize(float3(-FaceUV.x, -FaceUV.y, -1.0f)); // Face -Z; x=u,y=-v, z=-1
        default:
            // Never reach here.
            return 0.0f.xxx;
    }
}

float RadicalInverse_VdC(uint Bits)
{
    Bits = (Bits << 16u) | (Bits >> 16u);
    Bits = ((Bits & 0x55555555u) << 1u) | ((Bits & 0xAAAAAAAAu) >> 1u);
    Bits = ((Bits & 0x33333333u) << 2u) | ((Bits & 0xCCCCCCCCu) >> 2u);
    Bits = ((Bits & 0x0F0F0F0Fu) << 4u) | ((Bits & 0xF0F0F0F0u) >> 4u);
    Bits = ((Bits & 0x00FF00FFu) << 8u) | ((Bits & 0xFF00FF00u) >> 8u);
    return float(Bits) * 2.3283064365386963e-10; // / 0x100000000
}

float2 Hammersley(uint SampleIndex, uint SampleCount)
{
    return float2(float(SampleIndex) / float(SampleCount), RadicalInverse_VdC(SampleIndex));
}

float3 ImportanceSampleGGX(float2 Xi, float3 Normal, float Roughness)
{
    float A = Roughness * Roughness;
	
    float Phi = 2.0f * PI * Xi.x;
    float CosTheta = sqrt((1.0f - Xi.y) / (1.0f + (A * A - 1.0f) * Xi.y));
    float SinTheta = sqrt(1.0f - CosTheta * CosTheta);
	
	// from spherical coordinates to cartesian coordinates
    float3 H;
    H.x = cos(Phi) * SinTheta;
    H.y = sin(Phi) * SinTheta;
    H.z = CosTheta;
	
	// from tangent-space vector to world-space sample vector
    float3 Up = abs(Normal.z) < 0.999f ? float3(0.0f, 0.0f, 1.0f) : float3(1.0f, 0.0f, 0.0f);
    float3 Tangent = normalize(cross(Up, Normal));
    float3 Bitangent = cross(Normal, Tangent);
	
    float3 SampleVec = Tangent * H.x + Bitangent * H.y + Normal * H.z;
    return normalize(SampleVec);
}

float TrowbridgeReitzGGX(float NdotH, float Roughness)
{
    float NdotH = max(NdotH, 0.0f);

    float Coeff = Roughness * Roughness;
    float Alpha2 = Coeff * Coeff;
    float NdotH2 = NdotH * NdotH;

    float Numerator = Alpha2;
    float Denominator = NdotH2 * (Alpha2 - 1.0f) + 1.0f;
    Denominator = PI * Denominator * Denominator;

    return Numerator / Denominator;
}

float GetSampleMipLevel(float NdotH, float HdotV, float SampleCount, float roughness, float envRes)
{
    float D = TrowbridgeReitzGGX(NdotH, roughness);
    float Pdf = (D * NdotH / (4.0f * HdotV)) + 0.0001f;

    float SurfAreaTexel = 4.0f * PI / (6.0f * envRes * envRes);
    float SurfAreaSample = 1.0f / (SampleCount * Pdf + 0.0001f);

    return roughness == 0.0f ? 0.0f : 0.5f * log2(SurfAreaSample / SurfAreaTexel);
}

float IBLGeometrySchlickGGX(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 PunctualGeometrySchlickGGX(float NdotV, float Roughness)
{
    float R = Roughness + 1;
    // (Rough + 1)^2 / 8 for Punctual Lights
    // Rough^2 / 2 for IBL
    float K = (R * R) / 8.0;

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

    return Numerator / Denominator;
}

float IBLGeometrySmith(float NdotV, float NdotL, float Roughness)
{
    float GGX1 = IBLGeometrySchlickGGX(NdotV, Roughness);
    float GGX2 = IBLGeometrySchlickGGX(NdotL, Roughness);

    return GGX1 * GGX2;
}

float PunctualGeometrySmith(float NdotV, float NdotL, float Roughness)
{
    float GGX1 = PunctualGeometrySchlickGGX(NdotV, Roughness);
    float GGX2 = PunctualGeometrySchlickGGX(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 = IBLGeometrySmith(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);
}