#include "depth_normal.fsh"
#include "idioms.sh"
#include "ASSAOSamples.fsh"

#ifdef _SHADERC_PRESET
#define ASSAO_USE_DEPTH_GATHER
#endif

// Mutable uniforms
uniform float4 ScreenToViewParams[2];// [0]: Pack ruler: (u, v, z) scales, z linear scale i.e. far - near
                                     // [1]: Pack ruler: (u, v, z) offsets, z linear offset i.e. near
uniform float4 AdaptiveOccRadiusParams;// xyz: Pack ruler: a, b, c of curve: y = a / (z + b) + c
                                       // w: clamped minimum normalized depth

uniform float4 DepthTextureSizeInv;    // Depth texture size and reciprocal

// Immutable uniforms
uniform half4 AOStrengthParams;        // Pack ruler: a, b of curve: y = x^a * b, reciprocal of falloff (normalized) depth in view space
uniform half2 DepthFalloffParams;           // Pack ruler: fall-off depth start, fall-off depth range reciporcal

#define ASSAO_CONTRAST      AOStrengthParams.xxxx
#define ASSAO_INTENSITY     AOStrengthParams.yyyy
#define ASSAO_DEPTH_RADIUS  AOStrengthParams.wwww

#define ASSAO_FALLOFF_START DepthFalloffParams.xxxx
#define ASSAO_FALLOFF_RCP   DepthFalloffParams.yyyy

#ifdef ASSAO_USE_DEPTH_GATHER
uniform mediump sampler2D DepthGatherTexture;
#endif

const float ASSAO_FALLOFF_MIN_RADIUS = 1.0f;

const float SKY_DEPTH_VALUE = 0.9999;

const float TWO_PI         = 6.283185307179586;

const half4 HALF_ONE        = float4(1.f);
const half4 HALF_ZERO       = float4(0.f);
const half4 HALF_POINT_FIVE = float4(0.5f);
const half4 HALF_POINT_TWO  = float4(0.2f);

// Reconstruct position in view space
float3 ReconstructViewPos(float2 uv, float zvn) {
    float3 pos = float3(uv, zvn) * ScreenToViewParams[0].xyz + ScreenToViewParams[1].xyz;
    pos.xy *= pos.z;
    return pos;
}

#ifdef ASSAO_USE_DEPTH_GATHER
float4 GatherDepth(float2 uv) {
    return texture(DepthGatherTexture, uv);
}
float4 GatherDepth(int2 st) {
    return texelFetch(DepthGatherTexture, st, 0);
}
#endif// ASSAO_USE_DEPTH_GATHER

void GetSampleViewPos(float2 uv, out float zvn, out float3 posV) {
#ifdef ASSAO_USE_DEPTH_GATHER
    uv            = clamp(uv, float2(0.0), 1.0 - DepthTextureSizeInv.zw);// keep off texel index UB
    float2 sc     = uv * DepthTextureSizeInv.xy;
    float2 st     = fract(sc);
    sc            = sc - st;
    float4 depth4 = GatherDepth(int2(sc));

// #define COMPARE_DEPTH(a, b) step(a, b)
//     float3 result = float3(0.0, 0.0, depth4.x);
//     result        = lerp(result, float3(1.0, 0.0, depth4.y), COMPARE_DEPTH(depth4.y, result.z));
//     result        = lerp(result, float3(0.0, 1.0, depth4.z), COMPARE_DEPTH(depth4.z, result.z));
//     result        = lerp(result, float3(1.0, 1.0, depth4.w), COMPARE_DEPTH(depth4.w, result.z));
// #undef COMPARE_DEPTH
//     zvn = result.z;
//     uv = (sc + result.xy) * DepthTextureSizeInv.zw;

    float2 idxf   = step(float2(0.5f), st);
    float2 z0     = lerp(depth4.xz, depth4.yw, idxf.xx);
    zvn           = lerp(z0.x, z0.y, idxf.y);
    uv            = (sc + idxf) * DepthTextureSizeInv.zw;
#else
    zvn = _SampleDepth(uv);
#endif
    posV = ReconstructViewPos(uv, zvn);
}

void GetWorldNormalAndViewPos(float2 uv, out float4 zvn, out float4x3 posV, out half4x3 normalV) {
    float4 depth4 = GatherDepth(uv);

    float2 texelOffsetRT = DepthTextureSizeInv.zw;
    float2 uv4[4] = float2[4](
        uv,
        uv + float2(texelOffsetRT.x, 0.0),
        uv + float2(0.0, texelOffsetRT.y),
        uv + texelOffsetRT
    );
    posV[0] = ReconstructViewPos(uv4[0], depth4.x);
    posV[1] = ReconstructViewPos(uv4[1], depth4.y);
    posV[2] = ReconstructViewPos(uv4[2], depth4.z);
    posV[3] = ReconstructViewPos(uv4[3], depth4.w);

    zvn = depth4;

    normalV[0] = _SampleNormal(uv4[0]);
    normalV[1] = _SampleNormal(uv4[1]);
    normalV[2] = _SampleNormal(uv4[2]);
    normalV[3] = _SampleNormal(uv4[3]);
}

float GetAdaptiveOcclusionRadius(float zvn) {
    // For pixels 'infinitely' faraway, discard AO
    float clipFactor = step(zvn, SKY_DEPTH_VALUE);
    float y          = AdaptiveOccRadiusParams.x / (max(zvn, AdaptiveOccRadiusParams.w) + AdaptiveOccRadiusParams.y) +
              AdaptiveOccRadiusParams.z;
    float r = clipFactor * y;
    return r;
}

half4 EvalAmbientOcclusion(float2 uv, float4 zvn, float4x3 p, half4x3 n) {
    float zvn_s;
    float3 p_s;
    GetSampleViewPos(uv, zvn_s, p_s);

    float4 zvn_sX4 = float4(zvn_s);
    float4x3 p_sX4 = float4x3(p_s, p_s, p_s, p_s);

    half4 isInsideRadius = step(abs(zvn_sX4 - zvn), ASSAO_DEPTH_RADIUS) * step(zvn_sX4, float4(SKY_DEPTH_VALUE));

    float4x3 pps = p_sX4 - p; // Note: we cannot use 'half4x3' directly because of data overflow.
#if 0
    const half kBeta    = 0.008;
    const half kEpsilon = 0.001;
    float4 w;
    w.x = dot(pps[0], pps[0]);
    w.y = dot(pps[1], pps[1]);
    w.z = dot(pps[2], pps[2]);
    w.w = dot(pps[3], pps[3]);
    w = rsqrt(w + kEpsilon);

    float4 NdotD;
    NdotD.x = dot(n[0], pps[0]);
    NdotD.y = dot(n[1], pps[1]);
    NdotD.z = dot(n[2], pps[2]);
    NdotD.w = dot(n[3], pps[3]);

    float4 pzc = float4(p[0].z, p[1].z, p[2].z, p[3].z);
    half4 ao = max(NdotD + kBeta * pzc, HALF_ZERO) * w * isInsideRadius;
#else
    half4x3 v;
    v[0]      = normalize(pps[0]);
    v[1]      = normalize(pps[1]);
    v[2]      = normalize(pps[2]);
    v[3]      = normalize(pps[3]);

    half4 ao;
    ao.x    = dot(v[0], n[0]);
    ao.y    = dot(v[1], n[1]);
    ao.z    = dot(v[2], n[2]);
    ao.w    = dot(v[3], n[3]);
    ao      = max(ao - HALF_POINT_TWO, HALF_ZERO) * isInsideRadius;
#endif

    return ao;
}

half4 RemapAO(half4 ao) {
    return pow(ao, ASSAO_CONTRAST) * ASSAO_INTENSITY;
}

half4 GetDepthAttenuation(float4 zvn) {
    // Fall off by view depth
    half4 falloff = HALF_ONE - saturate((zvn - ASSAO_FALLOFF_START) * ASSAO_FALLOFF_RCP);
    falloff = falloff * falloff;
    return falloff;
}

float InterleavedGradientNoise(float2 positionSS) {
    float3 magic = float3(0.06711056f, 4.0f * 0.00583715f, 52.9829189f);
    return fract(magic.z * fract(dot(positionSS, magic.xy)));
}

float2x2 GetSampleRotationMatrix(float2 positionSS) {
    float noise        = InterleavedGradientNoise(positionSS);
    float2 sincosPhi = sincos(TWO_PI * noise);
    return float2x2(sincosPhi.y, -sincosPhi.x, sincosPhi.x, sincosPhi.y);
}

float2 GetSampleOffset(float2x2 matSampleRatate, int sampleIndex) {
    float2 sampleOffset = SamplingDiskOffsets[sampleIndex];
    return matSampleRatate * sampleOffset;
}

half4 GetAmbientOcclusion(float2 texCoord) {

    half4x3 n;
    float4x3 pos;
    float4 zvn;
    GetWorldNormalAndViewPos(texCoord.xy, zvn, pos, n);

    // Attenuation AO radius by view depth
    float radiusSS = GetAdaptiveOcclusionRadius(zvn.x);

    // Early clipped and return
    if (radiusSS <= ASSAO_FALLOFF_MIN_RADIUS) {
        return HALF_ONE;
    }

    float2 depthViewportSizeInv = DepthTextureSizeInv.zw;
    float2 uvRadius             = radiusSS * depthViewportSizeInv;

    half4 ao = HALF_ZERO;

    float2x2 matSampleRotate = GetSampleRotationMatrix(gl_FragCoord.xy);

#pragma unroll
    for (int i = 0; i < ASSAO_SAMPLE_COUNT; ++i) {
        float2 sampleOffset = GetSampleOffset(matSampleRotate, i);
        sampleOffset *= uvRadius;
        ao += EvalAmbientOcclusion(texCoord + sampleOffset, zvn, pos, n);
    }

    half4 falloff = GetDepthAttenuation(zvn);
    ao = ao * falloff * ASSAO_SAMPLE_COUNT_RCP;

    ao = HALF_ONE - RemapAO(ao);
    return ao;
}


in float2 UV_f;
layout(location = 0) out half4 fs_AO;

void main() {
    half4 ao = GetAmbientOcclusion(UV_f);
    fs_AO = ao;
}
