/*
====================================================================================================

    Copyright (C) 2020 RRe36

    All Rights Reserved unless otherwise explicitly stated.


    By downloading this you have agreed to the license and terms of use.
    These can be found inside the included license-file
    or here: https://rre36.com/copyright-license

    Violating these terms may be penalized with actions according to the Digital Millennium
    Copyright Act (DMCA), the Information Society Directive and/or similar laws
    depending on your country.

====================================================================================================
*/

#include "/lib/head.glsl"
#include "/lib/util/encoders.glsl"

const bool colortex7Clear   = false;

/*DRAWBUFFERS:37*/
layout(location = 0) out vec4 lightmapData;
layout(location = 1) out vec3 temporalData;

in vec2 coord;

uniform sampler2D colortex1;
uniform sampler2D colortex3;
uniform sampler2D colortex6;
uniform sampler2D colortex7;

uniform sampler2D depthtex0;
uniform sampler2D depthtex2;

uniform sampler2D noisetex;

uniform int frameCounter;

uniform float far, near;

uniform vec2 taaOffset;
uniform vec2 viewSize;

uniform vec3 cameraPosition, previousCameraPosition;
uniform vec3 upvecView;

uniform mat4 gbufferModelView, gbufferModelViewInverse;
uniform mat4 gbufferProjection, gbufferProjectionInverse;
uniform mat4 gbufferPreviousProjection, gbufferPreviousModelView;

#ifndef ATROUS_SIZE
    #define ATROUS_SIZE 1
#endif


/* ------ includes ------*/
#include "/lib/util/transforms.glsl"

#define FUTIL_LINDEPTH
#include "/lib/fUtil.glsl"

#include "/lib/util/kernel.glsl"
#include "/lib/light/atrousLightmap.glsl"

/* ------ reprojection ----- */
vec3 reproject(vec3 sceneSpace) {
    vec3 prevScreenPos = cameraPosition - previousCameraPosition;
    prevScreenPos = sceneSpace + prevScreenPos;
    prevScreenPos = viewMAD(gbufferPreviousModelView, prevScreenPos);
    prevScreenPos = viewMAD(gbufferPreviousProjection, prevScreenPos) * (0.5 / -prevScreenPos.z) + 0.5;

    return prevScreenPos;
}

#include "/lib/frag/bluenoise.glsl"
#include "/lib/frag/gradnoise.glsl"

/* ------ RTAO ------ */
//Fermats Spiral and Spheremap functions from raspberry by rutherin https://www.patreon.com/user?u=12731730
vec2 fermatsSpiralGoldenN(float index, float total) {
	float theta = index * goldenAngle;
	return vec2(sin(theta), cos(theta)) * sqrt(index / total);
}
vec3 sphereMap(vec2 a) {
    float phi = a.y * 2.0 * pi;
    float cosTheta = 1.0 - a.x;
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);

    return vec3(cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta);
}

float ditherBluenoise2() {
    ivec2 coord = ivec2(fract(gl_FragCoord.xy/256.0)*256.0);
    float noise = texelFetch(noisetex, coord, 0).a;

        noise   = fract(noise + (frameCounter * 7.0 * rcp(255.0)));

    return noise;
}

float bayer2(vec2 c) { 
    c = 0.5 * floor(c);
    return fract(1.5 * fract(c.y) + c.x);
}
float bayer4(vec2 c) {
    return 0.25 * bayer2 (0.5 * c) + bayer2(c);
}
float bayer8(vec2 c) {
    return 0.25 * bayer4 (0.5 * c) + bayer2(c);
}
float bayer16(vec2 c) {
    return 0.25 * bayer8 (0.5 * c) + bayer2(c);
}
float bayer32(vec2 c) {
    return 0.25 * bayer16(0.5 * c) + bayer2(c);
}
float bayer32() {
    return fract(bayer32(gl_FragCoord.xy) + float(frameCounter) * 7.0 / 255.0);
}

vec3 screenspaceRT_long(vec3 position, vec3 direction, float noise) {
    const uint maxSteps     = rtaoSteps;
    const float stepSize    = (halfPi * rtaoLength) / float(maxSteps);

    vec3 stepVector         = direction * stepSize;

    vec3 endPosition        = position + stepVector * maxSteps;
    vec3 endScreenPosition  = viewToScreenSpace(endPosition);

    vec2 maxPosXY           = max(abs(endScreenPosition.xy * 2.0 - 1.0), vec2(1.0));
    float stepMult          = min2(vec2(1.0) / maxPosXY);
        stepVector         *= stepMult;

    vec3 samplePos          = position + stepVector / maxSteps / 2.0;
        samplePos          += stepVector * noise;

    for (uint i = 0; i < maxSteps; ++i) {
        vec3 screenPos      = viewToScreenSpace(samplePos);
            samplePos      += stepVector;
        if (saturate(screenPos.xy) != screenPos.xy) break;

        float depthSample   = texelFetch(depthtex0, ivec2(screenPos.xy * viewSize), 0).x;
        float linearSample  = depthLinear(depthSample);
        float currentDepth  = depthLinear(screenPos.z);

        if (linearSample < currentDepth) {
            float dist      = abs(linearSample - currentDepth) / currentDepth;
            if (dist <= 0.2) return vec3(screenPos.xy, depthSample);
        }
    }

    return vec3(1.1);
}

#define HASHSCALE1 .1031
float hash12(vec2 x) {
    vec3 x3  = fract(vec3(x.xyx) * HASHSCALE1);
        x3  += dot(x3, x3.yzx + 19.19);
    return fract((x3.x + x3.y) * x3.z);
}

float ditherHashNoise() {
    float noise     = hash12(gl_FragCoord.xy);
        noise   = fract(noise + (frameCounter * 7.0 * rcp(255.0)));

    return noise;
}

float getRTAO(vec3 viewPos, vec3 viewNormal, float dither, bool hand) {
    //if (hand) return 1.0;

    const uint samples  = rtaoSPP;
    const uint steps    = rtaoSteps;

    vec3 viewDir        = normalize(viewPos);

    float ambient       = 0.0;

    const float maxRayDist  = rtaoLength;
    const float maxStepLength  = maxRayDist / float(steps);
    float distCoeff     = mix(0.125, 1.0, linStep(length(viewPos), maxStepLength * 0.125, maxRayDist));
    float startStepLength = maxStepLength * distCoeff;
    float stepLength    = startStepLength;

    vec3 rStart         = viewPos;
    vec3 rPos           = rStart;

    bool hit    = false;
    bool hitBehind = false;
    float hitErrorScale = 1.0;

    vec3 hitPos = vec3(0.0);

    const float refineMult  = pow(0.5, float(steps));

    float noise     = bayer32();
    float gnoise    = ditherHashNoise();

    vec3 upDir  = mat3(gbufferModelView) * vec3(0.0, 1.0, 0.0);

    for (uint i = 0; i < samples; ++i) {
        vec3 rDir           = sphereMap(fermatsSpiralGoldenN(float(i + 1) * 64.0 + noise * 64.0, float(samples) * 64.0));
            rDir            = mat3(gbufferModelView) * rDir;

        #ifndef DIM
            rDir            = mix(rDir, upDir, 0.33);
        #endif

        if (dot(rDir, viewNormal) < 0.0) rDir = -rDir;

        vec3 hitPosition    = screenspaceRT_long(viewPos, rDir, gnoise);

        if (hitPosition.z < 1.0) continue;

        ambient    += saturate(dot(rDir, viewNormal));
        //ambient += 1.0;
    }
    ambient    *= halfPi;
    ambient    /= float(samples);
    //if (ambient > 1.0) return 1.0;
    //else return 0.0;
    ambient     = saturate(ambient);

    return (ambient);
}

#define maxFrames 30.0

void main() {
    float sceneDepth    = stex(depthtex0).x;

    vec4 data           = lightmap3x3(colortex3, depthtex0, colortex6, coord, ATROUS_SIZE);

    temporalData        = vec3(1.0);

    #ifdef rtaoEnabled
    if (landMask(sceneDepth)) {
        #ifdef lightmapSmoothingEnabled
        vec3 sceneNormal    = stex(colortex6).xyz * 2.0 - 1.0;
        #else
        vec3 sceneNormal    = decodeNormal(stex(colortex1).xy);
        #endif
        vec3 viewNormal     = mat3(gbufferModelView) * sceneNormal;
        vec3 viewPos        = screenToViewSpace(vec3(coord, sceneDepth));

        bool hand           = sceneDepth < stex(depthtex2).x;

        vec3 scenePos       = viewToSceneSpace(viewPos);

        vec3 reprojection   = reproject(scenePos);
        bool offscreen      = saturate(reprojection) != reprojection;

        vec3 history        = saturate(texture(colortex7, reprojection.xy).xyz);

        vec3 cameraMovement = mat3(gbufferModelView) * (cameraPosition - previousCameraPosition);

        float currentDistance   = length(scenePos);
        float distanceDelta     = distance(history.y * far, currentDistance) - abs(cameraMovement.z);
        //float depthRejection    = (offscreen) ? 0.0 : exp(-max(distanceDelta - 0.01, 0.0) * 8.0);
        float depthRejection    = (offscreen || distanceDelta >= 1.0) ? 0.0 : exp(-max(distanceDelta - 0.01, 0.0) * 8.0);

        float accumulated   = 1.0;
        float pixelAge      = history.z;

        float rtao          = getRTAO(viewPos, viewNormal, ditherBluenoise2(), hand);

        if (offscreen || hand) {
            accumulated = rtao;
            pixelAge    = 0.0;
        } else {
            accumulated = mix(history.x, rtao, max(1.0 - depthRejection, 0.1));
            pixelAge    = saturate(pixelAge + rcp(maxFrames));
            //pixelAge    = mix(saturate(pixelAge + rcp(maxFrames)), 0.0, max(1.0 - depthRejection, 0.0));
        }

        temporalData        = saturate(vec3(accumulated, currentDistance / far, pixelAge));

        data.a  = accumulated;
        if (hand) data.a = 1.0;
    }
    #endif
    
    lightmapData = clampDrawbuffer(data);
}