#include "/lib/Surface/BRDF.glsl"

#define RAYTRACE_SAMPLES 16           
#define RAYTRACE_REFINEMENT
#define RAYTRACE_REFINEMENT_STEPS 4   
#define WATER_REFRACT_IOR 1.333       

#define FLICKER_SUPPRESS_DISTANCE 160.0

vec3 CalculateUnderwaterReflectDir(in vec3 viewDir, in vec3 normal) {
    vec3 reflectDir = reflect(viewDir, normal);
    float NdotV = max(dot(normal, -viewDir), 0.0);
    vec3 upVector = vec3(0.0, 1.0, 0.0);
    float upAlignment = dot(reflectDir, upVector);
    reflectDir.xz = mix(reflectDir.xz, normalize(reflectDir.xz) * 0.9, saturate(upAlignment * 0.5));
    return normalize(reflectDir);
}

bool ValidateDepthContinuity(in ivec2 texel, in float centerDepth, in float threshold) {
    float depthL = texelFetch(depthtex1, texel + ivec2(-1, 0), 0).x;
    float depthR = texelFetch(depthtex1, texel + ivec2(1, 0), 0).x;
    float depthU = texelFetch(depthtex1, texel + ivec2(0, -1), 0).x;
    float depthD = texelFetch(depthtex1, texel + ivec2(0, 1), 0).x;
    
    float avgDepth = (depthL + depthR + depthU + depthD) * 0.25;
    float variance = abs(depthL - avgDepth) + abs(depthR - avgDepth) + 
                     abs(depthU - avgDepth) + abs(depthD - avgDepth);
    
    return variance < threshold;
}

bool ScreenSpaceRayTrace(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                         inout vec3 rayPos, vec3 worldPos, bool isUnderwater) {
    const float maxLength = 1.0 / steps;
    const float minLength = length(screenPixelSize);
    
    vec3 position = ViewToScreenSpace(viewDir * abs(viewPos.z) + viewPos);
    vec3 screenDir = normalize(position - rayPos);
    float stepWeight = 1.0 / abs(screenDir.z);
    float stepLength = minOf((step(0.0, screenDir) - rayPos) / screenDir) * rcp(steps);
    
    screenDir.xy *= screenSize;
    rayPos.xy *= screenSize;
    
    float stableDither = dither;
    float distRatio = 0.0;
    
    if (isUnderwater) {
        distRatio = saturate(length(viewPos) / FLICKER_SUPPRESS_DISTANCE);
        
        stableDither = mix(dither, 0.5, distRatio * 0.9);
    }
    
    float spiralAngle = stableDither * TAU;
    vec2 spiralOffset = vec2(cos(spiralAngle), sin(spiralAngle)) * minLength * 2.0;
    
    vec3 rayStep = screenDir * stepLength;
    rayPos += rayStep * stableDither + screenDir * minLength + vec3(spiralOffset, 0.0);
    
    float depthTolerance = max(abs(rayStep.z) * 3.0, 0.02 / sqr(viewPos.z));
    
    if (isUnderwater) {
        depthTolerance *= 2.0;
        
        depthTolerance *= (1.0 + distRatio * 4.0);
    }
    
    bool hit = false;
    float depth = texelFetch(depthtex1, ivec2(rayPos.xy), 0).x;
    
    for (uint i = 0u; i < steps; ++i) {
        if (clamp(rayPos.xy, vec2(0.0), screenSize) != rayPos.xy) return false;
        if (rayPos.z >= 1.0) break;
        
        stepLength = abs(depth - rayPos.z) * stepWeight;
        
        float adaptiveStep = clamp(stepLength, minLength, maxLength);
        if (isUnderwater && i < 4u) {
            adaptiveStep *= 0.5;
        }
        
        rayPos += screenDir * adaptiveStep;
        depth = texelFetch(depthtex1, ivec2(rayPos.xy), 0).x;
        
        if (depth < rayPos.z && abs(depthTolerance - (rayPos.z - depth)) < depthTolerance) {
            if (isUnderwater) {
                if (distRatio > 0.8) {
                    hit = true;
                    break;
                } else if (ValidateDepthContinuity(ivec2(rayPos.xy), depth, 0.1 * (1.0 + distRatio))) {
                    hit = true;
                    break;
                }
            } else {
                hit = true;
                break;
            }
        }
    }
    
    if (!hit) return false;
    rayStep = screenDir * stepLength;
    
    #ifdef RAYTRACE_REFINEMENT
        for (uint i = 0u; i < RAYTRACE_REFINEMENT_STEPS; ++i) {
            if (clamp(rayPos.xy, vec2(0.0), screenSize) != rayPos.xy) break;
            rayStep *= 0.5;
            depth = texelFetch(depthtex1, ivec2(rayPos.xy), 0).x;
            if (depth < rayPos.z && abs(depthTolerance - (rayPos.z - depth)) < depthTolerance) {
                rayPos -= rayStep;
            } else {
                rayPos += rayStep;
            }
        }
    #endif
    
    return depth < 1.0;
}

bool ScreenSpaceRayTrace(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                         inout vec3 rayPos, vec3 worldPos) {
    return ScreenSpaceRayTrace(viewPos, viewDir, dither, steps, rayPos, worldPos, false);
}

bool ScreenSpaceRayTrace(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                         inout vec3 rayPos) {
    vec3 fallbackWorldPos = vec3(0.0, cameraPosition.y - 100.0, 0.0);
    return ScreenSpaceRayTrace(viewPos, viewDir, dither, steps, rayPos, fallbackWorldPos, false);
}

#if defined DISTANT_HORIZONS
    bool ScreenSpaceRayTraceDH(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                               inout vec3 rayPos, vec3 worldPos, bool isUnderwater) {
        const float maxLength = 1.0 / steps;
        const float minLength = length(screenPixelSize);
        
        vec3 position = ViewToScreenSpaceDH(viewDir * abs(viewPos.z) + viewPos);
        vec3 screenDir = normalize(position - rayPos);
        float stepWeight = 1.0 / abs(screenDir.z);
        float stepLength = minOf((step(0.0, screenDir) - rayPos) / screenDir) * rcp(steps);
        
        screenDir.xy *= screenSize;
        rayPos.xy *= screenSize;
        
        float stableDither = dither;
        float distRatio = 0.0;
        
        if (isUnderwater) {
            distRatio = saturate(length(viewPos) / FLICKER_SUPPRESS_DISTANCE);
            stableDither = mix(dither, 0.5, distRatio * 0.9);
        }
        
        float spiralAngle = stableDither * TAU;
        vec2 spiralOffset = vec2(cos(spiralAngle), sin(spiralAngle)) * minLength * 2.0;
        
        vec3 rayStep = screenDir * stepLength;
        rayPos += rayStep * stableDither + screenDir * minLength + vec3(spiralOffset, 0.0);
        
        float depthTolerance = max(abs(rayStep.z) * 3.0, 0.02 / sqr(viewPos.z));
        
        if (isUnderwater) {
            depthTolerance *= 2.5;
            depthTolerance *= (1.0 + distRatio * 5.0);
        }
        
        bool hit = false;
        float depth = texelFetch(dhDepthTex1, ivec2(rayPos.xy), 0).x;
        
        for (uint i = 0u; i < steps; ++i) {
            if (clamp(rayPos.xy, vec2(0.0), screenSize) != rayPos.xy) return false;
            if (rayPos.z >= 1.0) break;
            
            stepLength = abs(depth - rayPos.z) * stepWeight;
            
            float adaptiveStep = clamp(stepLength, minLength, maxLength);
            if (isUnderwater && i < 4u) {
                adaptiveStep *= 0.5;
            }
            
            rayPos += screenDir * adaptiveStep;
            depth = texelFetch(dhDepthTex1, ivec2(rayPos.xy), 0).x;
            
            if (depth < rayPos.z && abs(depthTolerance - (rayPos.z - depth)) < depthTolerance) {
                if (isUnderwater) {
                    if (distRatio > 0.8) {
                        hit = true;
                        break;
                    }
                    float depthNeighbor = texelFetch(dhDepthTex1, ivec2(rayPos.xy) + ivec2(1, 0), 0).x;
                    if (abs(depth - depthNeighbor) < 0.05 * (1.0 + distRatio * 2.0)) {
                        hit = true;
                        break;
                    }
                } else {
                    hit = true;
                    break;
                }
            }
        }
        
        if (!hit) return false;
        rayStep = screenDir * stepLength;
        
        #ifdef RAYTRACE_REFINEMENT
            for (uint i = 0u; i < RAYTRACE_REFINEMENT_STEPS; ++i) {
                if (clamp(rayPos.xy, vec2(0.0), screenSize) != rayPos.xy) break;
                rayStep *= 0.5;
                depth = texelFetch(dhDepthTex1, ivec2(rayPos.xy), 0).x;
                if (depth < rayPos.z && abs(depthTolerance - (rayPos.z - depth)) < depthTolerance) {
                    rayPos -= rayStep;
                } else {
                    rayPos += rayStep;
                }
            }
        #endif
        
        return depth < 1.0;
    }
    
    bool ScreenSpaceRayTraceDH(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                               inout vec3 rayPos, vec3 worldPos) {
        return ScreenSpaceRayTraceDH(viewPos, viewDir, dither, steps, rayPos, worldPos, false);
    }
    
    bool ScreenSpaceRayTraceDH(in vec3 viewPos, in vec3 viewDir, in float dither, in uint steps, 
                               inout vec3 rayPos) {
        vec3 fallbackWorldPos = vec3(0.0, cameraPosition.y - 100.0, 0.0);
        return ScreenSpaceRayTraceDH(viewPos, viewDir, dither, steps, rayPos, fallbackWorldPos, false);
    }
#endif

vec3 sampleGGXVNDF(in vec3 viewDir, in float roughness, in vec2 xy) {
    #define SPECULAR_TAIL_CLAMP
    #ifdef SPECULAR_TAIL_CLAMP
        xy.y = clamp(xy.y * 0.25, 1e-3, 0.25);
    #endif
    viewDir = normalize(vec3(roughness * viewDir.xy, viewDir.z));
    float phi = TAU * xy.x;
    float cosTheta = oneMinus(xy.y) * (1.0 + viewDir.z) - viewDir.z;
    float sinTheta = sqrt(saturate(1.0 - cosTheta * cosTheta));
    vec3 reflected = vec3(cossin(phi) * sinTheta, cosTheta);
    vec3 halfway = reflected + viewDir;
    return normalize(vec3(roughness * halfway.xy, halfway.z));
}