-- glslfx version 0.1

//
// Copyright 2018 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": {
            "displacementShader": {
                "source": [ "Preview.Displacement" ]
            },
            "surfaceShader": {
                "source": [ "Preview.LightStructures",
                            "Preview.Lighting",
                            "Preview.LightIntegration",
                            "Preview.NormalMapping",
                            "Preview.Surface" ]
            }
        }
    }
}


-- glsl Preview.Displacement

vec4
displacementShader(int index, vec4 Peye, vec3 Neye, vec4 patchCoord)
{
    // Calculate scalar displacement.
    float texDisplacement = HdGetScalar_displacement(index);
    return Peye + vec4(Neye*texDisplacement, 0);
}


-- glsl Preview.Surface

vec4
surfaceShader(vec4 Peye, vec3 Neye, vec4 color, vec4 patchCoord)
{
    float clearcoatAmount    = HdGetScalar_clearcoat();
    float clearcoatRoughness = HdGetScalar_clearcoatRoughness();
    vec3 diffuseColor        = HdGet_diffuseColor().xyz;
    vec3 emissiveColor       = HdGet_emissiveColor().xyz;
    float ior                = HdGetScalar_ior();
    float metallic           = HdGetScalar_metallic();
    float occlusion          = HdGetScalar_occlusion();
    float opacity            = HdGetScalar_opacity();
    float opacityThreshold   = HdGetScalar_opacityThreshold();
    float roughness          = HdGetScalar_roughness();
    vec3 specularColor       = HdGet_specularColor().xyz;
    bool useSpecularWorkflow = (HdGetScalar_useSpecularWorkflow() == 1);

    vec3 clearcoatColor      = vec3(1.0);
    float specularAmount     = 1.0;

#ifdef HD_MATERIAL_TAG_MASKED   
    if (opacity < opacityThreshold) {
        discard;
    } 
    opacity = 1.0;
#endif            

    // Selection highlighting.
    vec4 colorAndOpacity = vec4(diffuseColor, opacity);
    diffuseColor = ApplyColorOverrides(colorAndOpacity).rgb;

    // Pre-multiply diffuse color by opacity if not done so already
#ifndef diffuseColor_IS_PREMULTIPLIED
    diffuseColor *= opacity;
#endif    

    // Evaluate all lights.
    vec3 c = evaluateLights(
        emissiveColor,
        diffuseColor,
        useSpecularWorkflow,
        ior,
        metallic,
        specularAmount,
        specularColor,
        roughness,
        clearcoatAmount,
        clearcoatColor,
        clearcoatRoughness,
        occlusion,
        Peye,
        perturbFragmentNormal(Peye.xyz, Neye));

    return vec4(c, colorAndOpacity.a);
}


-- glsl Preview.LightStructures

struct LightingContributions
{
    vec3 diffuse;
    vec3 specular;
};


-- glsl Preview.Lighting

#define PI 3.1415
#define EPSILON 0.001


float
SchlickFresnel(float EdotH)
{
    return pow(max(0.0, 1.0 - EdotH), 5.0);
}

float
NormalDistribution(float specularRoughness, float NdotH)
{
    float alpha = specularRoughness * specularRoughness;
    float alpha2 = alpha * alpha;
    float NdotH2 = NdotH * NdotH;
    float DDenom = (NdotH2 * (alpha2 - 1.0)) + 1.0;
    DDenom *= DDenom;
    DDenom *= PI;
    float D = (alpha2 + EPSILON) / DDenom;
    return D;
}

float
Geometric(
        float specularRoughness,
        float NdotL,
        float NdotE,
        float NdotH,
        float EdotH)
{
    float alpha = specularRoughness * specularRoughness;
    float k = alpha * 0.5;
    float G = NdotE / (NdotE * (1.0 - k) + k);
    G *= NdotL / (NdotL * (1.0 - k) + k);
    return G;
}

vec3
evaluateDirectDiffuse()
{
    return vec3(1.0 / PI);
}

vec3
evaluateDirectSpecular(
        vec3 specularColorF0,
        vec3 specularColorF90,
        float specularRoughness,
        float fresnel,
        float NdotL,
        float NdotE,
        float NdotH,
        float EdotH)
{
    vec3 F = mix(specularColorF0, specularColorF90, fresnel);
    float D = NormalDistribution(specularRoughness, NdotH);
    float G = Geometric(specularRoughness, NdotL, NdotE, NdotH, EdotH);
    vec3 RNum = F * G * D;
    float RDenom = 4.0f * NdotL * NdotE + EPSILON;
    return RNum / RDenom;
}

LightingContributions
evaluateLight(
        vec3 diffuseColor,
        bool useSpecularWorkflow,
        float ior,
        float metallic,
        float specularAmount,
        vec3 specularColor,
        float specularRoughness,
        float clearcoatAmount,
        vec3 clearcoatColor,
        float clearcoatRoughness,
        float occlusion,
        float NdotL,
        float NdotE,
        float NdotH,
        float EdotH,
        vec3 lightDiffuseIrradiance,
        vec3 lightSpecularIrradiance)
{
    specularRoughness = max(0.001, specularRoughness);
    clearcoatRoughness = max(0.001, clearcoatRoughness);

    float fresnel = SchlickFresnel(EdotH);

    // Evaluate diffuse
    vec3 d = diffuseColor * evaluateDirectDiffuse();

    // Evaluate specular first lobe
    vec3 s1 = vec3(0.0);
    const float R = (1.0 - ior) / (1.0 + ior);
    if (specularAmount > 0.0) {
        vec3 F0 = specularColor;
        vec3 F90 = vec3(1.0);

        if (!useSpecularWorkflow) {
            vec3 specColor = mix(vec3(1.0), diffuseColor, metallic);
            F0 = mix(R * R * specColor, specColor, metallic);
            F90 = specColor;

            // For metallic workflows, pure metals have no diffuse
            d *= 1.0 - metallic;
        }

        s1 = specularAmount * evaluateDirectSpecular(
             F0,                          // Specular color 0
             F90,                         // Specular color 90
             specularRoughness,           // Roughness
             fresnel,                     // Fresnel
             NdotL, NdotE, NdotH, EdotH); // Dot products needed for lights

        // Adjust the diffuse so glazing angles have less diffuse
        d *= (1.0 - mix(F0, F90, fresnel));
    }

    // Evaluate clearcoat
    vec3 s2 = vec3(0.0);
    if (clearcoatAmount > 0.0) {
        s2 = clearcoatAmount * evaluateDirectSpecular(
             R * R * clearcoatColor,      // Clearcoat color 0
             clearcoatColor,              // Clearcoat color 90
             clearcoatRoughness,          // Roughness
             fresnel,                     // Fresnel
             NdotL, NdotE, NdotH, EdotH); // Dot products needed for lights
    }

    LightingContributions lightingContrib;

    lightingContrib.diffuse =
        occlusion * NdotL * d * lightDiffuseIrradiance;

    lightingContrib.specular =
        occlusion * NdotL * (s1 + s2) * lightSpecularIrradiance;

    return lightingContrib;
}


-- glsl Preview.LightIntegration

vec2 projectToLatLong(vec3 sample3D)
{
    // project spherical coord onto latitude-longitude map with
    // latitude: +y == pi/2 and longitude: +z == 0, +x == pi/2
    vec2 coord = vec2((atan(sample3D.z, sample3D.x) + 0.5 * PI) / (2.0 * PI),
                      acos(sample3D.y) / PI);
    return coord;
}

mat4 GetDomeLightTransform(mat4 worldToLightTransform)
{
    // transform from view space to light space
    mat4 worldToViewInverse = GetWorldToViewInverseMatrix();
    return worldToLightTransform * worldToViewInverse;
}

LightingContributions
evaluateIndirectLighting(
        vec3 diffuseColor,
        vec3 specularColor,
        vec3 Neye,
        vec3 Reye,
        float NdotE,
        float EdotH,
        float ior,
        float metallic,
        float occlusion,
        float roughness,
        bool useSpecularWorkflow,
        float clearcoatAmount, 
        vec3 clearcoatColor, 
        float clearcoatRoughness,
        mat4 worldToLightTransform)
{
    LightingContributions indirect;
    indirect.diffuse = vec3(0.0);
    indirect.specular = vec3(0.0);

#ifdef HD_HAS_domeLightIrradiance
    vec3 F0 = specularColor;
    vec3 F90 = vec3(1.0);
    vec3 d = diffuseColor;
    const float R = (1.0 - ior) / (1.0 + ior);
    if (!useSpecularWorkflow) {
        vec3 specColor = mix(vec3(1.0), diffuseColor, metallic);
        F0 = mix(R * R * specColor, specColor, metallic);
        F90 = specColor;

        // For metallic workflows, pure metals have no diffuse
        d *= 1.0 - metallic;
    }
    // Adjust the diffuse so glazing angles have less diffuse
    float fresnel = SchlickFresnel(EdotH);
    vec3 F = mix(F0, F90, fresnel);
    d *= (1.0 - F);

    mat4 transformationMatrix = GetDomeLightTransform(worldToLightTransform);

    // Diffuse Component
    vec3 dir = normalize((transformationMatrix * vec4(Neye,0.0)).xyz);
    vec2 coord = projectToLatLong(dir);
    vec3 diffuse = HdGet_domeLightIrradiance(coord).rgb;

    // Specular Component 
    const float MAX_REFLECTION_LOD = 
        textureQueryLevels(HdGetSampler_domeLightPrefilter());
    float lod = roughness * MAX_REFLECTION_LOD;
    vec3 Rdir = normalize((transformationMatrix * vec4(Reye,0.0)).xyz);
    vec2 Rcoord = projectToLatLong(Rdir);
    vec3 prefilter = HdTextureLod_domeLightPrefilter(Rcoord, lod).rgb;

    vec2 brdf = HdGet_domeLightBRDF(vec2(NdotE, roughness)).rg;

    vec3 specular = prefilter * (F * brdf.x + brdf.y);

    // Clearcoat Component
    vec3 clearcoat = vec3(0.0);
    if (clearcoatAmount > 0.0) {
        const vec3 clearcoatF = clearcoatAmount * mix(
            R * R * clearcoatColor, // Clearcoat F0
            clearcoatColor,         // Clearcoat F90
            fresnel);
        lod = clearcoatRoughness * MAX_REFLECTION_LOD;
        prefilter = HdTextureLod_domeLightPrefilter(Rcoord, lod).rgb;
        clearcoat = prefilter * (clearcoatF * brdf.x + brdf.y);
    }

    // Indirect Lighting
    indirect.diffuse = (d * diffuse) * occlusion;
    indirect.specular = (specular + clearcoat) * occlusion;

#endif

    return indirect;
}

vec3
evaluateLights(
        vec3 emissiveColor,
        vec3 diffuseColor,
        bool useSpecularWorkflow,
        float ior,
        float metallic,
        float specularAmount,
        vec3 specularColor,
        float specularRoughness,
        float clearcoatAmount,
        vec3 clearcoatColor,
        float clearcoatRoughness,
        float occlusion,
        vec4 Peye,
        vec3 Neye)
{
    vec3 n = Neye;
    vec3 e = normalize(-Peye.xyz);
    float NdotE = max(0.0, dot(n, e));

    vec3 Reye = reflect(-e, n);

    vec3 directLight = vec3(0.0);
    vec3 indirectLight = vec3(0.0);

#if NUM_LIGHTS > 0
    for (int i = 0; i < NUM_LIGHTS; ++i) {
        LightSource light = GetLightSource(i);

        // Calculate necessary vector information for lighting
        vec4 Plight = light.isIndirectLight
                        ? vec4(0,0,0,1)
                        : light.position;
        vec3 l = (Plight.w == 0.0)
                    ? normalize(Plight.xyz)
                    : normalize(Plight - Peye).xyz;
        vec3 h = normalize(e + l);
        float NdotL = max(0.0, dot(n, l));
        float NdotH = max(0.0, dot(n, h));
        float EdotH = max(0.0, dot(e, h));

        // Calculate light intensity
        float atten = lightDistanceAttenuation(Peye, i);
        float spot = lightSpotAttenuation(l, i);

        // Calculate the shadow factor
        float shadow = 1.0;
    #if USE_SHADOWS
        shadow = light.hasShadow ?
            shadowing(/*lightIndex=*/i, Peye) : 1.0;
    #endif

        float intensity = atten * spot * shadow;

        vec3 lightDiffuseIrradiance = intensity * light.diffuse.rgb;
        vec3 lightSpecularIrradiance = intensity * light.specular.rgb;

        LightingContributions lightingContrib = evaluateLight(
            diffuseColor,
            useSpecularWorkflow,
            ior,
            metallic,
            specularAmount,
            specularColor,
            specularRoughness,
            clearcoatAmount,
            clearcoatColor,
            clearcoatRoughness,
            occlusion,
            NdotL,
            NdotE,
            NdotH,
            EdotH,
            lightDiffuseIrradiance,
            lightSpecularIrradiance);

        // calculate the indirect light (DomeLight)
        if (light.isIndirectLight) {

            LightingContributions indirectLightContrib = 
                evaluateIndirectLighting(diffuseColor, specularColor,
                    Neye, Reye, NdotE, EdotH, ior, metallic, occlusion,
                    specularRoughness, useSpecularWorkflow,
                    clearcoatAmount, clearcoatColor, clearcoatRoughness,
                    light.worldToLightTransform);
            indirectLight = (indirectLightContrib.diffuse * light.diffuse.rgb 
                           + indirectLightContrib.specular * light.specular.rgb);
        }
        // all other light sources contribute to the direct lighting
        else {
            directLight += (lightingContrib.diffuse + lightingContrib.specular);
        }
    }
#endif

    return (emissiveColor + directLight + indirectLight);
}


-- glsl Preview.NormalMapping

vec3
perturbFragmentNormal(vec3 P, vec3 N)
{
#ifdef HD_HAS_COORD_normal
    // Calculation of TBN matrix and terminology based on "Surface 
    // Gradient-Based Bump Mapping Framework" (2020)
    vec3 Nt = normalize(HdGet_normal().xyz);
    vec2 st = HdGetCoord_normal().xy;

    // Get screen space derivatives of position
    vec3 dPdx = dFdx(P);
    vec3 dPdy = dFdy(P);

    // Ensure position derivatives are perpendicular to N
    vec3 sigmaX = dPdx - dot(dPdx, N) * N;
    vec3 sigmaY = dPdy - dot(dPdy, N) * N;

    float flipSign = dot(dPdy, cross(N, dPdx)) < 0 ? -1 : 1;

    // Get screen space derivatives of st
    vec2 dSTdx = dFdx(st);
    vec2 dSTdy = dFdy(st);

    // Get determinant and determinant sign of st matrix
    float det = dot(dSTdx, vec2(dSTdy.y, -dSTdy.x));
    float signDet = det < 0 ? -1 : 1;

    // Get first column of inv st matrix
    // Don't divide by det, but scale by its sign
    vec2 invC0 = signDet * vec2(dSTdy.y, -dSTdx.y);

    vec3 T = sigmaX * invC0.x + sigmaY * invC0.y;

    if (abs(det) > 0) {
        T = normalize(T);
    }

    vec3 B = (signDet * flipSign) * cross(N, T);

    mat3 TBN = mat3(T, B, N);
    return normalize(TBN * Nt);
#endif
    return N;
}


