#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.xx
#define ASSAO_INTENSITY AOStrengthParams.yy
#define ASSAO_DEPTH_RADIUS AOStrengthParams.ww

#define ASSAO_FALLOFF_START DepthFalloffParams.xx
#define ASSAO_FALLOFF_RCP DepthFalloffParams.yy

#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 half2 HALF_ONE        = float2(1.f);
const half2 HALF_ZERO       = float2(0.f);
const half2 HALF_POINT_FIVE = float2(0.5f);
const half2 HALF_POINT_TWO  = float2(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(float4 uvs, int2 subpixelIndex, out float2 zvn, out float2x3 posV, out half2x3 normalV) {
    float4 depth4 = GatherDepth(uvs.xy);

    posV[0] = ReconstructViewPos(uvs.xy, depth4[subpixelIndex.x]);
    posV[1] = ReconstructViewPos(uvs.zw, depth4[subpixelIndex.y]);

    zvn.x = depth4[subpixelIndex.x];
    zvn.y = depth4[subpixelIndex.y];

    normalV[0] = _SampleNormal(uvs.xy);
    normalV[1] = _SampleNormal(uvs.zw);
}

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

half2 EvalAmbientOcclusion(float4 uvs, int2 subpixelIndex, float2 zvn, float2x3 p, half2x3 n) {
    float zvn_s;
    float3 p_s;
    GetSampleViewPos(uvs.xy, subpixelIndex.x, zvn_s, p_s);

    float2 zvn_sX2 = float2(zvn_s);
    float2x3 p_sX2 = float2x3(p_s, p_s);

    half2 isInsideRadius = step(abs(zvn_sX2 - zvn), ASSAO_DEPTH_RADIUS) * step(zvn_sX2, float2(SKY_DEPTH_VALUE));

    float2x3 pps = p_sX2 - p;
    half2x3 v;
    v[0]      = normalize(pps[0]);
    v[1]      = normalize(pps[1]);
    half2 ao;
    ao.x = dot(v[0], n[0]);
    ao.y = dot(v[1], n[1]);
    ao   = max(ao - HALF_POINT_TWO, HALF_ZERO) * isInsideRadius;

    return ao;
}

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

half2 GetDepthAttenuation(float2 zvn) {
    // Fall off by view depth
    half2 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 float4 _Dither2x2PositionMap[6] = float4[6](
    float4(0.0, 0.0, 0.6, 0.0),
    float4(0.6, 0.0, 0.6, 0.6),
    float4(0.0, 0.0, 0.6, 0.6),
    float4(0.6, 0.0, 0.0, 0.6),
    float4(0.0, 0.0, 0.0, 0.6),
    float4(0.0, 0.6, 0.6, 0.6)
 );
 const int2 _Dither2x2SubpixelIndexMap[6] = int2[6](
    int2(0, 1),
    int2(1, 3),
    int2(0, 3),
    int2(1, 2),
    int2(0, 2),
    int2(2, 3)
 );

void DitheringPixelPosition(float2 uv, out float4 subpixelUVs, out int2 subpixelIndex) {
    int2 indexBits = int2(gl_FragCoord.xy);
    indexBits = indexBits % int2(3, 2);
    int idx = indexBits.x + indexBits.y * 3;
    subpixelUVs = uv.xyxy + _Dither2x2PositionMap[idx] * ViewportSizeInv.zwzw;
    subpixelIndex = _Dither2x2SubpixelIndexMap[idx];
}

half2 GetAmbientOcclusion(float2 texCoord) {

    float4 texCoords;
    int2 subpixelIndex;
    DitheringPixelPosition(texCoord, texCoords, subpixelIndex);

    half2x3 n;
    float2x3 pos;
    float2 zvn;
    GetWorldNormalAndViewPos(texCoords, subpixelIndex, 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;

    half2 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(texCoords + sampleOffset.xyxy, subpixelIndex, zvn, pos, n);
    }

    half2 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 = GetAmbientOcclusion(UV_f).rgrg; }
