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

uniform vec4 screenWidthHeight;
uniform float4 screenToViewParams; // Pack ruler: u scale, v scale, zndc scale, zndc offset
uniform vec4 ssaoRadiusStrengthPower;

in vec2 UV_f;

highp vec3 GetWorldNormalDefer(vec2 uv) {
    half3 n = _SampleNormal(uv);
    return normalize(n);
}
highp vec3 GetViewPos(vec2 uv) {

    float z;
    z = _SampleDepth(uv);
    float zViewSpace = z * screenToViewParams.z + screenToViewParams.w;
    float4 posV = float4((2.0 * uv - 1.0) * screenToViewParams.xy, -1.0, 1.0);
    posV.xyz *= zViewSpace;
    return posV.xyz;
}
highp float DoAmbientOcclusion(vec2 tcoord, vec2 uv, vec3 cnorm, vec3 cpos) {
    highp float result       = 0.0;
    highp vec3 offsetPos     = GetViewPos(tcoord + uv);
    highp vec3 deltaDistance = offsetPos - cpos;
    highp vec3 deltaNormal   = normalize(deltaDistance);
    result                   = dot(cnorm, deltaNormal) - 0.2;
    result                   = max(0.0, result);
    return result;
}
highp float rand(vec2 uv) {
    highp float a  = 12.9898;
    highp float b  = 78.233;
    highp float c  = 43758.5453;
    highp float dt = dot(uv.xy, vec2(a, b));
    highp float sn = mod(dt, 3.14);
    return fract(sin(sn) * c);
}
vec2 uniform_disk[4] = vec2[4](vec2(-0.6378324920834658, -0.5891375843299099), vec2(-0.13154583551882734, 0.10526998475744163), vec2(0.4011190068923373, 0.8992707985184016), vec2(0.22993594400661366, -0.14608947189178342));
float GetOcclution() {
    float2 vTexCoord = UV_f;
    float rand_angle   = rand(vTexCoord.xy) * 6.2831852;
    float rand_sin     = sin(rand_angle);
    float randCos      = cos(rand_angle);
    int sampleCount    = 4;
    int stepCount      = 3;
    highp vec3 n       = GetWorldNormalDefer(vTexCoord.xy);
    highp vec3 pos     = GetViewPos(vTexCoord.xy);
    highp float ao     = 0.0f;
    vec2 invViewportZW = screenWidthHeight.zw;
    float rad          = ssaoRadiusStrengthPower.x;
    for (int i = 0; i < sampleCount; ++i) {
        vec2 uvOriginal = uniform_disk[i];
        vec2 uvRotate   = vec2(uvOriginal.x * randCos + uvOriginal.y * rand_sin, -uvOriginal.x * rand_sin + uvOriginal.y * randCos);
        vec2 uv_offset  = uvRotate * rad * invViewportZW;
        for (int j = 0; j < stepCount; ++j) {
            float fade   = float(j + 1) / float(stepCount);
            vec2 uvFinal = uv_offset * fade;
            ao += DoAmbientOcclusion(vTexCoord.xy, uvFinal, n, pos);
        }
    }
    float aoStrength = ssaoRadiusStrengthPower.y;
    ao               = ao / (float(sampleCount) * float(stepCount)) * aoStrength;
    float aoPower    = ssaoRadiusStrengthPower.z;
    ao               = pow(ao, aoPower);
    ao *= ssaoRadiusStrengthPower.w;
    return 1.0 - ao;
}

out half oFragColor;

void main() {
    oFragColor   = GetOcclution();
}


#if 0

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

#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: AO radius of 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

#define ASSAO_DEPTH_RADIUS AdaptiveOccRadiusParams.w
#define ASSAO_CONTRAST AOStrengthParams.x
#define ASSAO_INTENSITY AOStrengthParams.y
#define ASSAO_FALLOFF_RCP AOStrengthParams.z

#ifdef ASSAO_USE_DEPTH_GATHER
uniform mediump sampler2D DepthGatherTexture;
#endif

const float ASSAO_FALLOFF_MIN_RADIUS = 3.0f;

const float SKY_DEPTH_VALUE = 0.9999;

const float TWO_PI = 6.283185307179586;

const half HALF_ONE        = 1.f;
const half HALF_ZERO       = 0.f;
const half HALF_POINT_FIVE = 0.5f;
const half HALF_POINT_TWO  = 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
    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 float zvn, out float3 posV, out half3 normalV) {

#ifdef ASSAO_USE_DEPTH_GATHER
    float4 depth4 = GatherDepth(uv);

    float2 texelOffsetRT = DepthTextureSizeInv.zw;

    float3 result = float3(0.0, 0.0, depth4.x);
    // Compute farthest point
#define COMPARE_DEPTH2(a, b) step(b, a)
    result = lerp(result, float3(1.0, 0.0, depth4.y), COMPARE_DEPTH2(depth4.y, result.z));
    result = lerp(result, float3(0.0, 1.0, depth4.z), COMPARE_DEPTH2(depth4.z, result.z));
    result = lerp(result, float3(1.0, 1.0, depth4.w), COMPARE_DEPTH2(depth4.w, result.z));
#undef COMPARE_DEPTH2

    zvn = result.z;
    uv += texelOffsetRT.xy * result.xy;
    posV    = ReconstructViewPos(uv, zvn);
    normalV = _SampleNormal(uv);
#else
    _SampleDepthNormal(uv, zvn, normalV);
    posV = ReconstructViewPos(uv, zvn);
#endif
}

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

half EvalAmbientOcclusion(float2 uv, float zvn, float3 p, half3 n) {
    float zvn_s;
    float3 p_s;
    GetSampleViewPos(uv, zvn_s, p_s);

    half isInsideRadius = step(abs(zvn_s - zvn), ASSAO_DEPTH_RADIUS) * step(zvn_s, SKY_DEPTH_VALUE);

    half3 v = p_s - p;
    v       = normalize(v);
    half ao = dot(v, n) - HALF_POINT_TWO;
    ao      = max(ao, HALF_ZERO) * isInsideRadius;

    return ao;
}

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

// Random samples
float CanonicalRand(float2 uv) {
    float a  = 12.9898;
    float b  = 78.233;
    float c  = 43758.5453;
    float dt = dot(uv.xy, float2(a, b));
    float sn = mod(dt, 3.14);
    return fract(sin(sn) * c);
}

const float2 UniformDiskSamples[4] =
        float2[4](float2(-0.6378324920834658, -0.5891375843299099), float2(-0.13154583551882734, 0.10526998475744163),
                  float2(0.4011190068923373, 0.8992707985184016), float2(0.22993594400661366, -0.14608947189178342));

float2x2 GetSampleRotationMatrix(float2 uv) {
    float randAngle  = CanonicalRand(uv) * 6.2831852;
    float2 sincosPhi = sincos(randAngle);
    return float2x2(sincosPhi.y, -sincosPhi.x, sincosPhi.x, sincosPhi.y);
}

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

half GetAmbientOcclusion(float2 texCoord) {

    half3 n;
    float3 pos;
    float zvn;
    GetWorldNormalAndViewPos(texCoord.xy, zvn, pos, n);

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

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

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

    half ao = HALF_ZERO;

    const int sampleCount             = 4;
    const int stepCount               = 3;
    const float fadeScales[stepCount] = float[stepCount](1.f / 3.f, 2.f / 3.f, 1.f);
    const float AOAverageFactor       = 1.f / float(sampleCount * stepCount);

    float2x2 matSampleRotate = GetSampleRotationMatrix(texCoord.xy);

#pragma unroll
    for (int i = 0; i < sampleCount; ++i) {

        float2 sampleOffset0 = GetSampleOffset(matSampleRotate, i);
        sampleOffset0 *= uvRadius;

#pragma unroll
        for (int j = 0; j < stepCount; ++j) {
            float2 sampleOffset = sampleOffset0 * fadeScales[j];
            ao += EvalAmbientOcclusion(texCoord + sampleOffset, zvn, pos, n);
        }
    }

    ao = ao * AOAverageFactor;

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

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

void main() { fs_AO = float4(GetAmbientOcclusion(UV_f)); }

#endif