#include "version.sh"

#define SAMPLE_COUNT 42
#define PI 3.1415926

uniform sampler2D texUnit;

uniform float rcpAspect;
uniform float maxRadius;
uniform highp vec2 perTexelDis;
uniform vec2 bokehKernel[SAMPLE_COUNT];
uniform bool postFilter;

in vec2 UV_f;

layout(location = 0) out vec4 fragColor;

vec4 getDownSampler(vec2 uv)
{
    vec3 duv = perTexelDis.xyx * vec3(0.5, 0.5, -0.5);
    vec2 uv0 = uv - duv.xy;
    vec2 uv1 = uv - duv.zy;
    vec2 uv2 = uv + duv.zy;
    vec2 uv3 = uv + duv.xy;
    
    vec4 color0 = texture(texUnit, uv0);
    vec4 color1 = texture(texUnit, uv1);
    vec4 color2 = texture(texUnit, uv2);
    vec4 color3 = texture(texUnit, uv3);
    
    float coc0 = color0.a * 2.0 - 1.0;
    float coc1 = color1.a * 2.0 - 1.0;
    float coc2 = color2.a * 2.0 - 1.0;
    float coc3 = color3.a * 2.0 - 1.0;
    
    float cocMin = min(coc0, min(min(coc1, coc2), coc3));
    float cocMax = max(coc0, max(max(coc1, coc2), coc3));
    
    float coc = (-cocMin > cocMax ? cocMin : cocMax) * maxRadius;
    
    vec3 avg = (color0.rgb + color1.rgb + color2.rgb + color3.rgb) * 0.25;
    
    float w0 = abs(coc0) / (max(color0.x, max(color0.y, color0.z)) + 1.0);
    float w1 = abs(coc1) / (max(color1.x, max(color1.y, color1.z)) + 1.0);
    float w2 = abs(coc2) / (max(color2.x, max(color2.y, color2.z)) + 1.0);
    float w3 = abs(coc3) / (max(color3.x, max(color3.y, color3.z) )+ 1.0);
    
    avg = color0.rgb * w0 + color1.rgb * w1 + color2.rgb * w2 + color3.rgb * w3;
    avg /= max(w0 + w1 + w2 + w3, 1e-5);
    
    avg *= smoothstep(0.0, perTexelDis.y * 2.0, abs(coc));
   
    return vec4(avg, coc);
}

void Accumulate(vec4 samp0, vec2 uv, vec2 disp, inout vec4 farAcc, inout vec4 nearAcc)
{
    float dist = length(disp);

    vec2 duv = vec2(disp.x * rcpAspect, disp.y);
    
    vec4 samp = getDownSampler(uv + duv);
    
    float farCoC = max(min(samp0.a, samp.a), 0.0);

    float margin = perTexelDis.y * 4.0;
    
    float farWeight = clamp((farCoC - dist + margin) / margin, 0.0, 1.0);
    float nearWeight = clamp((-samp.a - dist + margin) / margin, 0.0, 1.0);

    nearWeight *= step(perTexelDis.y * 2.0, -samp.a);

    farAcc += vec4(samp.rgb, 1.0) * farWeight;
    nearAcc += vec4(samp.rgb, 1.0) * nearWeight;
}

vec4 blur()
{
    vec4 samp0 = getDownSampler(UV_f);
    
    vec4 farAcc = vec4(0.0);
    vec4 nearAcc = vec4(0.0);
    
    Accumulate(samp0, UV_f, vec2(0.0), farAcc, nearAcc);
    
    for (int si = 0; si < SAMPLE_COUNT; si++)
    {
        vec2 disp = bokehKernel[si] * maxRadius;
        Accumulate(samp0, UV_f, disp, farAcc, nearAcc);
    }

    farAcc.rgb /= farAcc.a + ((farAcc.a == 0.0) ? 1.0: 0.0);
    nearAcc.rgb /= nearAcc.a + ((nearAcc.a == 0.0) ? 1.0: 0.0);

    nearAcc.a *= PI / float(SAMPLE_COUNT + 1);

    float alpha = clamp(nearAcc.a, 0.0, 1.0);
    vec3 rgb = mix(farAcc.rgb, nearAcc.rgb, alpha);
    return vec4(rgb, clamp(alpha * 0.5 + 0.5, 0.0, 1.0));
}

vec4 post()
{
    const float downSampler = 2.0f;
    vec4 param = vec4(0.5, 0.5, -0.5, 0.0);
    vec4 duv = perTexelDis.xyxy * downSampler * param;
    vec4 acc = texture(texUnit, UV_f - duv.xy);
    acc += texture(texUnit, UV_f - duv.zy);
    acc += texture(texUnit, UV_f + duv.zy);
    acc += texture(texUnit, UV_f + duv.xy);
    return acc * 0.25;
}

void main()
{
    if(postFilter)
        fragColor = post();
    else
        fragColor = blur();
}
