#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
uniform float4 ViewportSizeInv;

// 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.x
#define ASSAO_INTENSITY AOStrengthParams.y
#define ASSAO_DEPTH_RADIUS AOStrengthParams.w

#define ASSAO_FALLOFF_START DepthFalloffParams.x
#define ASSAO_FALLOFF_RCP DepthFalloffParams.y

#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 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, int subpixelIndex, 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;
    float4 depth4 = GatherDepth(int2(sc));

    zvn = depth4[subpixelIndex];
#else
    zvn = _SampleDepth(uv);
#endif
    posV = ReconstructViewPos(uv, zvn);
}

void GetWorldNormalAndViewPos(float2 uv, int subpixelIndex, 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[subpixelIndex]);

    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 / (max(zvn, AdaptiveOccRadiusParams.w) + AdaptiveOccRadiusParams.y) +
              AdaptiveOccRadiusParams.z;
    float r = clipFactor * y;
    return r;
}

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

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

    float3 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;
}

half GetDepthAttenuation(float zvn) {
    // Fall off by view depth
    half 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;
}

const float2 _Dither2x2PositionMap[4] = float2[4](
    float2(0.0, 0.0),
    float2(0.6, 0.6),
    float2(0.0, 0.6),
    float2(0.6, 0.0)
);
const int _Dither2x2IndexMap[4] = int[4](
    0,
    3,
    2,
    1
);

void DitheringPixelPosition(inout float2 uv, out int subpixelIndex) {
    int2 indexBits = int2(gl_FragCoord.xy);
    indexBits = indexBits % 2;
    int index = indexBits.x + indexBits.y * 2;
    uv += _Dither2x2PositionMap[index] * ViewportSizeInv.zw;
    subpixelIndex = _Dither2x2IndexMap[index];
}

half GetAmbientOcclusion(float2 texCoord) {

    int subpixelIndex;
    DitheringPixelPosition(texCoord, subpixelIndex);

    half3 n;
    float3 pos;
    float zvn;
    GetWorldNormalAndViewPos(texCoord, subpixelIndex, 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;

    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, subpixelIndex, zvn, pos, n);
    }

    half 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() { fs_AO = float4(GetAmbientOcclusion(UV_f)); }
