-- glslfx version 0.1

//
// Copyright 2019 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache License.
//

-- configuration
{
    "techniques": {
        "default": {
            "domeLightIrradiance": {
                "source": [ "DomeLight.Common",
                            "DomeLight.Irradiance" ]
            },
            "domeLightPrefilter": {
                "source": [ "DomeLight.Common",
                            "DomeLight.CommonSampling",
                            "DomeLight.Prefilter" ]
            },
            "domeLightBRDF": {
                "source": [ "DomeLight.Common",
                            "DomeLight.CommonSampling",
                            "DomeLight.BRDF" ]
            }
        }
    }
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.Common

const float PI = 3.1415926536;

layout(binding = 0) uniform sampler2D inTexture;
layout(rgba16f, binding = 1) uniform image2D outTexture;

// compute texture coords based on the size of the output image/texture
vec2 GetTexCoords(ivec2 outCoords)
{
    vec2 outDims = imageSize(outTexture);
    vec2 texCoords = outCoords / outDims;
    return texCoords;
}

// sample lat/long env map input texture
vec3 SampleEnvMapLod(vec3 sampleVec, float sampleLod) {
    vec2 coord = vec2((atan(sampleVec.z, sampleVec.x) + PI) / (2.0 * PI),
                      acos(sampleVec.y) / PI);
    return textureLod(inTexture, coord, sampleLod).rgb;
}

// compute world position from texture coords
vec3 GetWorldPos(vec2 textureCoord)
{
    // have theta range from [-PI, PI] so the origin is in the center
    // of the image
    float theta = (textureCoord.x * 2.0 * PI) - PI;
    float phi = (textureCoord.y * PI);
    float x = cos(theta) * sin(phi);
    float y = cos(phi);
    float z = sin(theta) * sin(phi);
    return vec3(x, y, z);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.CommonSampling

float RadicalInverse(uint a)
{
    return float(bitfieldReverse(a)) * 2.3283064365386963e-10; // 0x1p-32
}

vec2 Hammersley2d(uint a, uint N)
{
    return vec2(float(a) / float(N), RadicalInverse(a));
}

vec3 ImportanceSample_GGX(vec2 Xi, float roughness, vec3 normal)
{
    // Maps a 2D point to a hemisphere with spread based on roughness
    float alpha = roughness * roughness;
    float phi = 2.0 * PI * Xi.x;
    float cosTheta = sqrt((1.0 - Xi.y) / (1.0 + (alpha*alpha - 1.0) * Xi.y));
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);
    vec3 H = vec3(sinTheta * cos(phi), sinTheta * sin(phi), cosTheta);

    // Tangent space
    vec3 up = abs(normal.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    vec3 tangentX = normalize(cross(up, normal));
    vec3 tangentY = normalize(cross(normal, tangentX));

    // Convert to world Space
    return normalize(tangentX * H.x + tangentY * H.y + normal * H.z);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.Irradiance

layout(local_size_x = 32, local_size_y = 32) in;

const float deltaPhi = (2.0f * float(PI)) / 180.0f;
const float deltaTheta = (0.5f * float(PI)) / 64.0f;

vec3 SampleEnvMap(vec3 sampleVec)
{
    // sample from a mipmap level of the environment map determined by the
    // size of the environment map and the number of samples we are taking
    ivec2 inDims = textureSize(inTexture, 0);
    float mipLevel = ceil(log2(inDims.x * deltaPhi/(2.0 * PI)) * 2.0f);

    return SampleEnvMapLod(sampleVec, mipLevel);
}

vec3 ComputeIrradiance(vec3 inPos)
{
    vec3 N = normalize(inPos);
    vec3 up = vec3(0.0, 1.0, 0.0);
    vec3 right = normalize(cross(up, N));
    up = cross(N, right);

    const float TWO_PI = PI * 2.0;
    const float HALF_PI = PI * 0.5;

    vec3 color = vec3(0.0);
    uint sampleCount = 0u;
    for (float phi = 0.0; phi < TWO_PI; phi += deltaPhi) {
        for (float theta = 0.0; theta < HALF_PI; theta += deltaTheta) {
            vec3 tempVec = cos(phi) * right + sin(phi) * up;
            vec3 sampleVector = cos(theta) * N + sin(theta) * tempVec;
            color += SampleEnvMap(sampleVector).rgb * cos(theta) * sin(theta);
            sampleCount++;
        }
    }
    return PI * color / float(sampleCount);
}

void main(void)
{
    ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);

    vec2 texCoords = GetTexCoords(outCoords);
    vec3 pos3D = GetWorldPos(texCoords);
    vec4 outColor = vec4(ComputeIrradiance(pos3D), 1.0);

    imageStore(outTexture, outCoords, outColor);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.Prefilter

layout(local_size_x = 32, local_size_y = 32) in;
uniform float roughness;

// Normal Distribution function
float Distribution_GGX(float dotNH, float roughness)
{
    float alpha = roughness * roughness;
    float alpha2 = alpha * alpha;
    float denom = dotNH * dotNH * (alpha2 - 1.0) + 1.0;
    return (alpha2)/(PI * denom*denom);
}

vec3 PrefilterEnvMap(vec3 R, float roughness)
{
    vec3 N = R;
    vec3 V = R;
    vec3 color = vec3(0.0);
    float totalWeight = 0.0;
    float envMapDim = float(textureSize(inTexture, 0).s);
    const uint numSamples = 1024u;
    for (uint i = 0u; i < numSamples; i++) {
        vec2 Xi = Hammersley2d(i, numSamples);
        vec3 H = ImportanceSample_GGX(Xi, roughness, N);
        vec3 L = 2.0 * dot(V, H) * H - V;
        float dotNL = clamp(dot(N, L), 0.0, 1.0);
        if (dotNL > 0.0) {

            float dotNH = clamp(dot(N, H), 0.0, 1.0);
            float dotVH = clamp(dot(V, H), 0.0, 1.0);

            // Probability Distribution Function
            float pdf = Distribution_GGX(dotNH, roughness) * dotNH
                                                / (4.0 * dotVH) + 0.0001;
            // Solid angle of current sample
            float omegaS = 1.0 / (float(numSamples) * pdf);
            // Solid angle of 1 pixel across all cube faces
            float omegaP = 4.0 * PI / (6.0 * envMapDim * envMapDim);
            // Biased (+1.0) mip level for better result
            float mipLevel = roughness == 0.0
                           ? 0.0 : max(0.5 * log2(omegaS / omegaP) + 1.0, 0.0f);
            color += SampleEnvMapLod(L, mipLevel).rgb * dotNL;
            totalWeight += dotNL;

        }
    }
    return (color / totalWeight);
}

void main(void)
{
    ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);

    vec2 texCoords = GetTexCoords(outCoords);
    vec3 pos3D = GetWorldPos(texCoords);
    vec3 R = normalize(pos3D);
    vec4 outColor = vec4(PrefilterEnvMap(R, roughness), 1.0);

    imageStore(outTexture, outCoords, outColor);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.BRDF

layout(local_size_x = 32, local_size_y = 32) in;
uniform int sampleLevel = 0;

float Geometry_SchlicksmithGGX(float dotNL, float dotNV, float roughness)
{
    float k = (roughness * roughness) / 2.0;
    float GL = dotNL / (dotNL * (1.0 - k) + k);
    float GV = dotNV / (dotNV * (1.0 - k) + k);
    return GL * GV;
}

vec2 ComputeBRDF(float NoV, float roughness)
{
    // make sure NoV doesn't go exactly to 0 to avoid NaN
    NoV = max(NoV, 0.001);

    // Normal always points along z-axis for the 2D lookup
    const vec3 N = vec3(0.0, 0.0, 1.0);
    vec3 V = vec3(sqrt(1.0 - NoV*NoV), 0.0, NoV);

    vec2 LUT = vec2(0.0);
    const uint NUM_SAMPLES = 1024u;
    for (uint i = 0u; i < NUM_SAMPLES; i++) {
        vec2 Xi = Hammersley2d(i, NUM_SAMPLES);
        vec3 H = ImportanceSample_GGX(Xi, roughness, N);
        vec3 L = 2.0 * dot(V, H) * H - V;

        float dotNL = max(dot(N, L), 0.0);
        float dotNV = max(dot(N, V), 0.0);
        float dotVH = max(dot(V, H), 0.0);
        float dotNH = max(dot(H, N), 0.0);

        if (dotNL > 0.0) {
            float G = Geometry_SchlicksmithGGX(dotNL, dotNV, roughness);
            float G_Vis = (G * dotVH) / (dotNH * dotNV);
            float Fc = pow(1.0 - dotVH, 5.0);
            LUT += vec2((1.0 - Fc) * G_Vis, Fc * G_Vis);
        }
    }
    return LUT / float(NUM_SAMPLES);
}

void main(void)
{
    ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);

    vec2 texCoords = GetTexCoords(outCoords);
    // texCoords.x represents N dot E and texCoords.y represents roughness
    vec4 outColor = vec4(ComputeBRDF(texCoords.x, texCoords.y), 0.0, 1.0);

    imageStore(outTexture, outCoords, outColor);
}