<!
//  AmbientOcclusion.slh
//  libraries/render-utils/src
//
//  Created by Sam Gateau on 1/1/16.
//  Copyright 2016 High Fidelity, Inc.
//
//  Distributed under the Apache License, Version 2.0.
//  See the accompanying file LICENSE or http://www.apache.org/licenses/LICENSE-2.0.html
!>
<@if not SSAO_SLH@>
<@def SSAO_SLH@>

<@include render-utils/ShaderConstants.h@>
<@include ssao_shared.h@>

<@func declarePackOcclusionDepth()@>

float CSZToDepthKey(float z) {
    return clamp(z * (-1.0 / SSAO_DEPTH_KEY_SCALE), 0.0, 1.0);
}

vec4 packOcclusionOutput(float occlusion, float depth, vec3 eyeNormal) {
    depth = CSZToDepthKey(depth);
#if SSAO_BILATERAL_BLUR_USE_NORMAL    
    return vec4(occlusion, depth, eyeNormal.xy / eyeNormal.z);
#else
    // Round to the nearest 1/256.0
    depth *= 256;
    float temp = floor(depth);
    return vec4(occlusion, temp * (1.0 / 256.0), depth - temp, 0.0);
#endif
}

struct UnpackedOcclusion {
    vec3 normal;
    float depth;
    float occlusion;
};

void unpackOcclusionOutput(vec4 raw, out UnpackedOcclusion result) {
    result.occlusion = raw.x;
#if SSAO_BILATERAL_BLUR_USE_NORMAL    
    result.depth = raw.y;
    result.normal = normalize(vec3(raw.zw, 1.0));
#else
    result.depth = (raw.y + raw.z / 256.0);
    result.normal = vec3(0,0,1);
#endif
}

float unpackOcclusion(vec4 raw) {
    return raw.x;
}

<@endfunc@>

<@func declareAmbientOcclusion()@>
<@include DeferredTransform.slh@>
<$declareDeferredFrameTransform()$>

LAYOUT(binding=RENDER_UTILS_BUFFER_SSAO_PARAMS) uniform ambientOcclusionParamsBuffer {
    AmbientOcclusionParams params;
};

LAYOUT(binding=RENDER_UTILS_BUFFER_SSAO_FRAME_PARAMS) uniform ambientOcclusionFrameParamsBuffer {
    AmbientOcclusionFrameParams frameParams;
};

float getPerspectiveScale() {
    return (params._resolutionInfo.z);
}
int getResolutionLevel() { 
    return int(params._resolutionInfo.x);
}

bool isHorizonBased() {
    return params._resolutionInfo.y!=0.0;
}

vec2 getNormalsSideSize() {
    return params._sideSizes[0].xy;
}
int getNormalsResolutionLevel() {
    return int(params._sideSizes[0].z);
}
int getDepthResolutionLevel() {
    return int(params._sideSizes[0].w);
}
vec2 getOcclusionSideSize() {
    return params._sideSizes[1].xy;
}
vec2 getOcclusionSplitSideSize() {
    return params._sideSizes[1].zw;
}

ivec2 getWidthHeightRoundUp(int resolutionLevel) {
    ivec2 fullRes = ivec2(getWidthHeight(0));
    int resolutionDivisor = 1 << resolutionLevel;
    return (fullRes + resolutionDivisor - 1) / resolutionDivisor;
}

float getRadius() {
    return params._radiusInfo.x;
}
float getRadius2() {
    return params._radiusInfo.y;
}
float getInvRadius6() {
    return mix(params._radiusInfo.z, 1.0, isHorizonBased());
}
float getInvRadius2() {
    return params._radiusInfo.z;
}

float getObscuranceScaling() {
    return getInvRadius6() * params._radiusInfo.w;
}

float isDitheringEnabled() {
    return params._ditheringInfo.x;
}
float isBorderingEnabled() {
    return params._ditheringInfo.w;
}

float getFalloffCosAngle() {
    return params._falloffInfo.x;
}
float getFalloffCosAngleScale() {
    return params._falloffInfo.y;
}

float getFalloffSinAngle() {
    return params._falloffInfo.z;
}
float getFalloffSinAngleScale() {
    return params._falloffInfo.w;
}

float getNumSamples() {
    return params._sampleInfo.x;
}
float getInvNumSamples() {
    return params._sampleInfo.y;
}
float getNumSpiralTurns() {
    return params._sampleInfo.z;
}

int doFetchMips() {
    return int(params._sampleInfo.w);
}

<@endfunc@>

<@func declareSamplingDisk()@>

float getAngleDitheringWorldPos(in vec3 pixelWorldPos) {
    vec3 worldPosFract = fract(pixelWorldPos * 1.0);

    ivec3 pixelPos = ivec3(worldPosFract * 256.0);

    return isDitheringEnabled() * float(((3 * pixelPos.x ^ pixelPos.y + pixelPos.x * pixelPos.y) + (3 * pixelPos.y ^ pixelPos.z + pixelPos.x * pixelPos.z)) * 10);
}

float getAngleDitheringSplit() {
    return isDitheringEnabled() * frameParams._angleInfo.x;
}

float getAngleDithering(in ivec2 pixelPos) {
#if SSAO_USE_QUAD_SPLIT
    return getAngleDitheringSplit();
#else
    // Hash function used in the AlchemyAO paper
    return getAngleDitheringPixelPos(pixelPos);
#endif
}

float getAngleDitheringPixelPos(in ivec2 pixelPos) {
    // Hash function used in the AlchemyAO paper
    return isDitheringEnabled() *  float((3 * pixelPos.x ^ pixelPos.y + pixelPos.x * pixelPos.y) * 10);
}

float evalDiskRadius(float Zeye, vec2 sideImageSize) {
    // Choose the screen-space sample radius
    // proportional to the projected area of the sphere
    float diskPixelRadius = -( getProjScale(getResolutionLevel()) * getRadius() / Zeye ) * getPerspectiveScale();

    // clamp the disk to fit in the image otherwise too many unknown
    diskPixelRadius = min(diskPixelRadius, sideImageSize.y * 0.5);

    return diskPixelRadius;
}

const float PI = 3.1415926;
const float TWO_PI = 6.2831852;

vec3 getUnitTapLocation(int sampleNumber, float spiralTurns, float spinAngle, float angleRange){
    // Radius relative to ssR
    float alpha = float(sampleNumber) * getInvNumSamples();
    float angle = alpha * (spiralTurns * angleRange) + spinAngle;
    return vec3(cos(angle), sin(angle), alpha);
}

vec3 getTapLocationSSAO(int sampleNumber, float spinAngle, float outerRadius) {
    vec3 tap = getUnitTapLocation(sampleNumber, getNumSpiralTurns(), spinAngle, TWO_PI);
    tap.xy *= tap.z;
    tap *= outerRadius;
    return tap;
}

vec3 getTapLocationClampedSSAO(int sampleNumber, float spinAngle, float outerRadius, vec2 pixelPos, vec2 sideImageSize) {
    vec3 tap = getTapLocationSSAO(sampleNumber, spinAngle, outerRadius);
    vec2 tapPos = pixelPos + tap.xy;

    if (!(isBorderingEnabled() > 0.0)) {
        return tap;
    }
    bool redoTap = false;

    if ((tapPos.x < 0.5)) {
        tapPos.x = -tapPos.x;
        redoTap = true;
    } else if ((tapPos.x > sideImageSize.x - 0.5)) {
        tapPos.x -= (sideImageSize.x - tapPos.x);
        redoTap = true;
    }

    if ((tapPos.y < 0.5)) {
        tapPos.y = -tapPos.y;
        redoTap = true;
    } else if ((tapPos.y > sideImageSize.y - 0.5)) {
        tapPos.y -= (sideImageSize.y - tapPos.y);
        redoTap = true;
    }

    if (redoTap) {
        tap.xy = tapPos - pixelPos;
        tap.z = length(tap.xy);
        tap.z = 0.0;  
    }
    
    return tap;
}

<@endfunc@>


<@func declareFetchDepthPyramidMap()@>

// the depth pyramid texture
LAYOUT(binding=RENDER_UTILS_TEXTURE_SSAO_DEPTH) uniform sampler2D depthPyramidTex;
LAYOUT(binding=RENDER_UTILS_TEXTURE_SSAO_NORMAL) uniform sampler2D normalTex;

vec2 getFramebufferUVFromSideUV(ivec4 side, vec2 uv) {
    return isStereo() ? vec2((uv.x + float(getStereoSide(side))) * 0.5, uv.y) : uv;
}

vec2 getSideUVFromFramebufferUV(ivec4 side, vec2 uv) {
    return isStereo() ? vec2(uv.x * 2.0 - float(getStereoSide(side)), uv.y) : uv;
}

ivec2 getDepthTextureSize(int level) {
    return textureSize(depthPyramidTex, level);
}
ivec2 getDepthTextureSideSize(int level) {
    ivec2 size = getDepthTextureSize(level);
    size.x >>= int(isStereo()) & 1;
    return size;
}

ivec2 getNormalTextureSize(int level) {
    return textureSize(normalTex, level);
}
ivec2 getNormalTextureSideSize(int level) {
    ivec2 size = getNormalTextureSize(level);
    size.x >>= int(isStereo()) & 1;
    return size;
}

vec2 getStereoSideSizeRoundUp(int resolutionLevel) {
    ivec2 fullRes = ivec2(getStereoSideSize(0));
    int resolutionDivisor = 1 << resolutionLevel;
    return vec2((fullRes + resolutionDivisor - 1) / resolutionDivisor);
}

float getZEyeAtUV(vec2 texCoord, float level) {
    return -textureLod(depthPyramidTex, texCoord, level).x;
}

<@func getZEyeAtUVOffset(texCoord, level, texelOffset)@>
-textureLodOffset(depthPyramidTex, <$texCoord$>, <$level$>, <$texelOffset$>).x;
<@endfunc@>

float getZEyeAtUV(ivec4 side, vec2 texCoord, float level) {
    texCoord = getFramebufferUVFromSideUV(side, texCoord);
    return getZEyeAtUV(texCoord, level);
}

vec3 packNormal(vec3 normal) {
    vec3 absNormal = abs(normal);
    return 0.5 + normal * 0.5 / max(absNormal.x, max(absNormal.y, absNormal.z));
}

vec3 unpackNormal(vec3 packedNormal) {
    return normalize(packedNormal*2.0 - 1.0);
}

vec3 getNormalEyeAtUV(vec2 texCoord, float level) {
    return unpackNormal(textureLod(normalTex, texCoord, level).xyz);
}

vec3 getNormalEyeAtUV(ivec4 side, vec2 texCoord, float level) {
    texCoord = getFramebufferUVFromSideUV(side, texCoord);
    return getNormalEyeAtUV(texCoord, level);
}

vec2 snapToTexel(vec2 uv, vec2 pixelSize) {
	return (floor(uv * pixelSize - 0.5) + 0.5) / pixelSize;
}

int evalMipFromRadius(float radius) {
    const int LOG_MAX_OFFSET = 2;
    const int MAX_MIP_LEVEL = 5;
    return clamp(findMSB(int(radius)) - LOG_MAX_OFFSET, 0, MAX_MIP_LEVEL);
}

vec2 fetchTap(ivec4 side, vec2 tapUV, float tapRadius) {
    int mipLevel = evalMipFromRadius(tapRadius * float(doFetchMips()));

    vec2 fetchUV = clamp(tapUV, vec2(0), vec2(1));
    fetchUV = getFramebufferUVFromSideUV(side, fetchUV);

    vec2 P;
    P.x = float(mipLevel);
    P.y = -textureLod(depthPyramidTex, fetchUV, P.x).x;
    return P;
}

vec3 buildPosition(ivec4 side, vec2 fragUVPos) {
    float Zeye = getZEyeAtUV(side, fragUVPos, 0.0);
    return evalEyePositionFromZeye(side.x, Zeye, fragUVPos);
}

<@func buildPositionOffset(side, fragUVPos, sideFragUVPos, texelOffset, deltaUV, position)@>
{
    float Zeye = <$getZEyeAtUVOffset($sideFragUVPos$, 0.0, $texelOffset$)$>
    <$position$> = evalEyePositionFromZeye(<$side$>.x, Zeye, <$fragUVPos$> + vec2(<$texelOffset$>)*<$deltaUV$>);
}
<@endfunc@>

vec3 getMinDelta(vec3 centralPoint, vec3 offsetPointPos, vec3 offsetPointNeg) {
    vec3 delta0 = offsetPointPos - centralPoint;
    vec3 delta1 = centralPoint - offsetPointNeg;
    float sqrLength0 = dot(delta0, delta0);
    float sqrLength1 = dot(delta1, delta1);

    return sqrLength0 < sqrLength1 ? delta0 : delta1;
}

const ivec2 UV_RIGHT = ivec2(1,0);
const ivec2 UV_LEFT = ivec2(-1,0);
const ivec2 UV_TOP = ivec2(0,1);
const ivec2 UV_BOTTOM = ivec2(0,-1);

vec3 buildNormal(ivec4 side, vec2 fragUVPos, vec3 fragPosition, vec2 deltaDepthUV) {
    vec2 fullUVPos = getFramebufferUVFromSideUV(side, fragUVPos);

    vec3 fragPositionDxPos;
    vec3 fragPositionDxNeg;
    vec3 fragPositionDyPos;
    vec3 fragPositionDyNeg;

    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_RIGHT, deltaDepthUV, fragPositionDxPos)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_LEFT, deltaDepthUV, fragPositionDxNeg)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_TOP, deltaDepthUV, fragPositionDyPos)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_BOTTOM, deltaDepthUV, fragPositionDyNeg)$>

    vec3 fragDeltaDx = getMinDelta(fragPosition, fragPositionDxPos, fragPositionDxNeg);
    vec3 fragDeltaDy = getMinDelta(fragPosition, fragPositionDyPos, fragPositionDyNeg);

    return normalize( cross(fragDeltaDx, fragDeltaDy) );
}

void buildTangentBinormal(ivec4 side, vec2 fragUVPos, vec3 fragPosition, vec3 fragNormal, vec2 deltaDepthUV,
                          out vec3 fragTangent, out vec3 fragBinormal) {
    vec2 fullUVPos = getFramebufferUVFromSideUV(side, fragUVPos);

    vec3 fragPositionDxPos;
    vec3 fragPositionDxNeg;
    vec3 fragPositionDyPos;
    vec3 fragPositionDyNeg;

    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_RIGHT, deltaDepthUV, fragPositionDxPos)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_LEFT, deltaDepthUV, fragPositionDxNeg)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_TOP, deltaDepthUV, fragPositionDyPos)$>
    <$buildPositionOffset(side, fragUVPos, fullUVPos, UV_BOTTOM, deltaDepthUV, fragPositionDyNeg)$>

    vec3 fragDeltaDx = getMinDelta(fragPosition, fragPositionDxPos, fragPositionDxNeg);
    vec3 fragDeltaDy = getMinDelta(fragPosition, fragPositionDyPos, fragPositionDyNeg);

    //fragTangent = normalize( cross(fragDeltaDy, fragNormal) );
    //fragBinormal = normalize( cross(fragNormal, fragDeltaDx) );

    fragTangent = fragDeltaDx;
    fragBinormal = fragDeltaDy;
}

<@endfunc@>


<@func declareEvalObscurance()@>

struct TBNFrame {
    vec3 tangent;
    vec3 binormal;
    vec3 normal;
};

vec3 fastAcos(vec3 x) {
	// [Eberly2014] GPGPU Programming for Games and Science
	vec3 absX = abs(x);
	vec3 res = absX * (-0.156583) + vec3(PI / 2.0);
	res *= sqrt(vec3(1.0) - absX);
	return mix(res, vec3(PI) - res, greaterThanEqual(x, vec3(0)));
}

float evalVisibilitySSAO(in vec3 centerPosition, in vec3 centerNormal, in vec3 tapPosition) {
    vec3 v = tapPosition - centerPosition;
    float vv = dot(v, v);
    float vn = dot(v, centerNormal);

    // Falloff function as recommended in SSAO paper
    const float epsilon = 0.01;
    float f = max(getRadius2() - vv, 0.0);
    return f * f * f * max((vn - getFalloffCosAngle()) / (epsilon + vv), 0.0);
}

#define HBAO_USE_COS_ANGLE 1
#define HBAO_USE_OVERHANG_HACK 0

float computeWeightForHorizon(float horizonLimit, float distanceSquared) {
    return max(0.0, 1.0 - distanceSquared * getInvRadius2());
}

float computeWeightedHorizon(float horizonLimit, float distanceSquared) {
    float radiusFalloff = computeWeightForHorizon(horizonLimit, distanceSquared);

#if !HBAO_USE_COS_ANGLE
    horizonLimit = getFalloffSinAngle() - horizonLimit;
#endif
    horizonLimit *= radiusFalloff;
#if !HBAO_USE_COS_ANGLE
    horizonLimit = getFalloffSinAngle() - horizonLimit;
#endif

    return horizonLimit;
}

<@func computeHorizon()@>
        if (tapUVPos.x<0 || tapUVPos.y<0 || tapUVPos.x>=1.0 || tapUVPos.y>=1.0) {
            // Early exit because we've hit the borders of the frame
            break;
        }
        vec2 tapMipZ = fetchTap(side, tapUVPos, radius);
        vec3 tapPositionES = evalEyePositionFromZeye(side.x, tapMipZ.y, tapUVPos);
        vec3 deltaVec = tapPositionES - fragPositionES;
        float distanceSquared = dot(deltaVec, deltaVec);
        float deltaDotNormal = dot(deltaVec, fragFrameES.normal);
#if HBAO_USE_COS_ANGLE
        float tapHorizonLimit = deltaDotNormal;
#else
        float tapHorizonLimit = dot(deltaVec, fragFrameES.tangent);
#endif
        tapHorizonLimit *= inversesqrt(distanceSquared);

        if (distanceSquared < getRadius2() && deltaDotNormal>0.0) {
#if HBAO_USE_COS_ANGLE
            float weight = computeWeightForHorizon(tapHorizonLimit, distanceSquared);
            if (tapHorizonLimit > horizonLimit) {
                occlusion += weight * (tapHorizonLimit - horizonLimit);
                horizonLimit = tapHorizonLimit;
            }
#if HBAO_USE_OVERHANG_HACK 
            else if (dot(deltaVec, fragFrameES.tangent) < 0.0) {
                // This is a hack to try to handle the case where the occlusion angle is
                // greater than 90°
                occlusion = mix(occlusion, (occlusion+1.0) * 0.5, weight);
            }
#endif
#else   
            if (tapHorizonLimit < horizonLimit) {
                tapHorizonLimit = computeWeightedHorizon(tapHorizonLimit, distanceSquared);
                horizonLimit = min(horizonLimit, tapHorizonLimit);
            }
#endif
        }
<@endfunc@>

#define HBAO_HORIZON_SEARCH_CONSTANT_STEP 0

float computeOcclusion(ivec4 side, vec2 fragUVPos, vec3 fragPositionES, TBNFrame fragFrameES, vec2 searchDir, float searchRadius, int stepCount) {
    float occlusion = 0.0;
#if HBAO_USE_COS_ANGLE
    float horizonLimit = getFalloffCosAngle();
#else
    float horizonLimit = getFalloffSinAngle();
#endif

    if (stepCount>0) {
        vec2 deltaTapUV = searchDir / float(stepCount);
        vec2 tapUVPos;
        float deltaRadius = searchRadius / float(stepCount);
        vec2 sideDepthSize = getDepthTextureSideSize(0);

#if HBAO_HORIZON_SEARCH_CONSTANT_STEP
        float radius = 0.0;
        int stepIndex;

        for (stepIndex=0 ; stepIndex<stepCount ; stepIndex++) {
            fragUVPos += deltaTapUV;
            radius += deltaRadius;
            tapUVPos = snapToTexel(fragUVPos, sideDepthSize);

            <$computeHorizon()$>
        }
#else
        // Step is adapted to Mip level
        float radius = deltaRadius;
        float mipLevel = evalMipFromRadius(radius * float(doFetchMips()));

        while (radius<=searchRadius) {
            fragUVPos += deltaTapUV;
            tapUVPos = snapToTexel(fragUVPos, sideDepthSize);

            <$computeHorizon()$>

            if (tapMipZ.x != mipLevel) {
                mipLevel = tapMipZ.x;
                deltaRadius *= 2;
                deltaTapUV *= 2;
                sideDepthSize = getDepthTextureSideSize(int(mipLevel));
            }
            radius += deltaRadius;
        }
#endif
    }

#if HBAO_USE_COS_ANGLE
    occlusion = min(occlusion * getFalloffCosAngleScale(), 1.0);
#else
    occlusion = horizonLimit > 0.0 ? horizonLimit * getFalloffSinAngleScale() : horizonLimit;
#endif

    return occlusion;
}

float evalVisibilityHBAO(ivec4 side, vec2 fragUVPos, vec2 invSideImageSize, vec2 deltaTap, float diskPixelRadius, 
                         vec3 fragPositionES, vec3 fragNormalES) {
    vec2 pixelSearchVec = deltaTap * diskPixelRadius;
	vec2 searchDir = pixelSearchVec * invSideImageSize;
    vec2 deltaTapUV = deltaTap * invSideImageSize;
	float obscuranceH1 = 0.0;
	float obscuranceH2 = 0.0;
    pixelSearchVec = abs(pixelSearchVec);
    int stepCount = int(ceil(max(pixelSearchVec.x, pixelSearchVec.y)));
    TBNFrame fragFrameES;

    fragFrameES.tangent = vec3(0.0);
    fragFrameES.binormal = vec3(0.0);
    fragFrameES.normal = fragNormalES;

#if HBAO_USE_OVERHANG_HACK || !HBAO_USE_COS_ANGLE
    vec3 positionPos = buildPosition(side, fragUVPos + deltaTapUV);
    vec3 positionNeg = buildPosition(side, fragUVPos - deltaTapUV);

    fragFrameES.tangent = getMinDelta(fragPositionES, positionPos, positionNeg);
    fragFrameES.tangent -= dot(fragNormalES, fragFrameES.tangent) * fragNormalES;
    fragFrameES.tangent = normalize(fragFrameES.tangent);
#endif
    // Forward search for h1
    obscuranceH1 = computeOcclusion(side, fragUVPos, fragPositionES, fragFrameES, searchDir, diskPixelRadius, stepCount);

    // Backward search for h2
#if HBAO_USE_OVERHANG_HACK || !HBAO_USE_COS_ANGLE
    fragFrameES.tangent = -fragFrameES.tangent;
#endif
    obscuranceH2 = computeOcclusion(side, fragUVPos, fragPositionES, fragFrameES, -searchDir, diskPixelRadius, stepCount);

	return obscuranceH1 + obscuranceH2;
}

<@endfunc@>


<@endif@>