/* Prisma - water.glsl
Water reflection and refraction related functions.
*/

vec3 fastRaymarch(vec3 direction, vec3 hitPoint, inout float infinite, float dither) {
    vec3 pathIncrement;
    vec3 currentMarchPoint = hitPoint;
    vec3 oldMarchPoint;
    float screenDepth;
    float depthDifference = 1.0;
    vec3 screenMarchPos = camera_to_screen(hitPoint);
    float prevScreenDepth = screenMarchPos.z;
    float hitPointDepth = screenMarchPos.z;
    bool searchFlag = false;
    bool hiddenFlag = false;
    bool firstHidden = true;
    bool outOfEyeFlag = false;
    bool toFar = false;
    vec3 lastScreenMarchPos;
    
    int no_hidden_steps = 0;
    bool hiddens = false;

    // Ray marching
    for (int i = 0; i < RAYMARCH_STEPS; i++) {
        if (searchFlag) {
            pathIncrement *= 0.5;
            currentMarchPoint += pathIncrement * sign(depthDifference);
        } else {
            oldMarchPoint = currentMarchPoint;
            currentMarchPoint = hitPoint + ((direction * exp2(i + dither)) - direction);
            pathIncrement = currentMarchPoint - oldMarchPoint;
        }

        lastScreenMarchPos = screenMarchPos;
        screenMarchPos = camera_to_screen(currentMarchPoint);

        if ( // Is outside screen space
            screenMarchPos.x < 0.0 ||
            screenMarchPos.x > 1.0 ||
            screenMarchPos.y < 0.0 ||
            screenMarchPos.y > 1.0 ||
            screenMarchPos.z < 0.0
        ) {
            outOfEyeFlag = true;
        }

        if (screenMarchPos.z > 0.9999) {
            toFar = true;
        }

        screenDepth = texture2D(depthtex1, screenMarchPos.xy).x;
        depthDifference = screenDepth - screenMarchPos.z;

        if (depthDifference < 0.0 && abs(screenDepth - prevScreenDepth) > abs(screenMarchPos.z - lastScreenMarchPos.z)) {
            hiddenFlag = true;
            hiddens = true;
            if (firstHidden) {
                firstHidden = false;
            }
        } else if (depthDifference > 0.0) {
            hiddenFlag = false;
            if (!hiddens) {
                no_hidden_steps++;
            }
        }

        if (searchFlag == false && depthDifference < 0.0 && hiddenFlag == false) {
            searchFlag = true;
        }

        prevScreenDepth = screenDepth;
    }

    infinite = float(screenDepth > 0.9999);

    if (outOfEyeFlag) {
        infinite = 1.0;
        return screenMarchPos;
    } else if (toFar) {
        if (screenDepth > 0.9999) {
            infinite = 1.0;
            return screenMarchPos;
        } else if (no_hidden_steps < 3 || screenDepth > hitPointDepth) {
            return screenMarchPos;
        } else {
            infinite = 1.0;
            return vec3(1.0);
        }
    } else {
        return screenMarchPos;
    }
}

#if SUN_REFLECTION == 1
    #if !defined NETHER && !defined THE_END
        float sun_reflection(vec3 fragpos) {
            vec3 astroLightPos = worldTime > 12900 ? moonPosition : sunPosition;
            float astroAlignment =
                max(dot(normalize(fragpos), normalize(astroLightPos)), 0.0);

            return smoothstep(0.995, 1.0, astroAlignment) *
                clamp(lmcoord.y, 0.0, 1.0) *
                (1.0 - rainStrength) * 3.0;
        }
    #endif
#endif

vec3 prismaReconstructViewPos(vec2 uv, float depth) {
    vec4 clip = vec4(uv * 2.0 - 1.0, depth * 2.0 - 1.0, 1.0);
    vec4 view = gbufferProjectionInverse * clip;
    return view.xyz / max(view.w, 1e-6);
}

vec3 prismaReconstructWorldPos(vec2 uv, float depth) {
    vec3 viewPos = prismaReconstructViewPos(uv, depth);
    vec4 world = gbufferModelViewInverse * vec4(viewPos, 1.0);
    return world.xyz;
}

float prismaHash12(vec2 p) {
    vec3 p3 = fract(vec3(p.xyx) * 0.1031);
    p3 += dot(p3, p3.yzx + 33.33);
    return fract((p3.x + p3.y) * p3.z);
}

vec2 prismaRainRippleNormal(vec2 worldXZ, float visibility) {
    vec2 grid = worldXZ * 0.42;
    vec2 cellBase = floor(grid);
    vec2 ripple = vec2(0.0);
    float time = frameTimeCounter * 0.72;

    for (int y = -1; y <= 1; y++) {
        for (int x = -1; x <= 1; x++) {
            vec2 cell = cellBase + vec2(float(x), float(y));
            vec2 rand = vec2(prismaHash12(cell + 11.17), prismaHash12(cell + 37.91));
            vec2 center = cell + rand;
            float age = fract(time + prismaHash12(cell + 5.73));
            vec2 d = grid - center;
            float dist = length(d) + 1e-4;
            float radius = mix(0.04, 1.18, age);
            float ring = exp(-24.0 * abs(dist - radius));
            float fade = (1.0 - age) * smoothstep(1.35, 0.08, dist);
            ripple += (d / dist) * ring * fade;
        }
    }

    ripple *= rainStrength * visibility * 0.44;
    return ripple;
}

float prismaWaterCaustics(vec3 worldPos, vec3 sunDir) {
    float sunY = max(abs(sunDir.y), 0.18);
    vec2 pos = (worldPos.xz - (sunDir.xz / sunY) * worldPos.y) * 0.018;
    float movement = frameTimeCounter * 0.028;

    mat2 rot = mat2(vec2(0.7374, -0.6755), vec2(0.6755, 0.7374));
    vec2 uv = pos;
    float large = max(texture2D(noisetex, uv * 0.18).b, 0.1);
    float h = 0.0;

    h += pow(abs(abs(texture2D(noisetex, uv / vec2(1.40, 0.42) + vec2(movement,  movement * 0.7) + large * 0.5).b * 2.0 - 1.0) * 2.0 - 1.0), 2.0);
    uv = rot * uv;
    h += pow(abs(abs(texture2D(noisetex, uv / vec2(0.46, 1.55) + vec2(-movement * 0.8, movement * 0.55) + large * 0.4).b * 2.0 - 1.0) * 2.0 - 1.0), 2.0);
    uv = rot * uv;
    h += pow(abs(abs(texture2D(noisetex, uv / vec2(0.95, 0.95) + vec2(movement * 0.55, -movement * 0.75) + large * 0.35).b * 2.0 - 1.0) * 2.0 - 1.0), 2.0);

    float c = exp((1.2 + 4.0 * sqrt(large)) * (h / 3.0 - 0.52));
    return clamp((c - 0.55) * 0.55, 0.0, 1.4);
}

vec3 normal_waves(vec3 pos) {
    float speed = frameTimeCounter * .025;
    vec2 wave_1 =
        texture2D(noisetex, ((pos.xy - pos.z * 0.2) * 0.05) + vec2(speed, speed)).rg;
    wave_1 = wave_1 - .5;
    vec2 wave_2 =
        texture2D(noisetex, ((pos.xy - pos.z * 0.2) * 0.03125) - speed).rg;
    wave_2 = wave_2 - .5;
    vec2 wave_3 =
        texture2D(noisetex, ((pos.xy - pos.z * 0.2) * 0.125) + vec2(speed, -speed)).rg;
    wave_3 = wave_3 - .5;
    wave_3 *= 0.66;

    vec2 partialWave = wave_1 + wave_2 + wave_3;

    // Derivative-inspired rain interaction: real world-space ripples on water during rain.
    vec2 rainRipple = prismaRainRippleNormal(pos.xy - pos.z * 0.16, visibleSky);
    partialWave += rainRipple;

    float waveHeight = WATER_TURBULENCE * mix(1.0, 0.72, rainStrength * visibleSky);
    vec3 finalWave = vec3(partialWave, waveHeight);

    return normalize(finalWave);
}

vec3 refraction(vec3 fragpos, vec3 color, vec3 refraction) {
    vec2 pos = gl_FragCoord.xy * vec2(pixelSizeX, pixelSizeY);

    #if REFRACTION == 1
        pos = pos + refraction.xy * (0.075 / (1.0 + length(fragpos) * 0.4));
    #endif

    float water_absortion;
    float water_thickness = 0.0;
    if (isEyeInWater == 0) {
        float water_distance =
        2.0 * near * far / (far + near - (2.0 * gl_FragCoord.z - 1.0) * (far - near));

        float earth_distance = texture2D(depthtex1, pos.xy).r;
        earth_distance =
            2.0 * near * far / (far + near - (2.0 * earth_distance - 1.0) * (far - near));

        #if defined DISTANT_HORIZONS
            float earth_distance_dh = texture2D(dhDepthTex1, pos.xy).r;
            earth_distance_dh =
                2.0 * dhNearPlane * dhFarPlane / (dhFarPlane + dhNearPlane - (2.0 * earth_distance_dh - 1.0) * (dhFarPlane - dhNearPlane));
            earth_distance = min(earth_distance, earth_distance_dh);
        #endif

        water_thickness = max(earth_distance - water_distance, 0.0);
        float thicknessSq = water_thickness * water_thickness;
        water_absortion = (1.0 / -((thicknessSq * WATER_ABSORPTION) + 1.125)) + 1.0;
    } else {
        water_absortion = 0.0;
    }

    vec3 sceneColor = texture2D(gaux1, pos.xy).rgb;

    float dayFactorSurface = smoothstep(0.18, 0.62, dayNightMix);
    float weatherMute = smoothstep(0.08, 0.95, rainStrength);
    float depthBlend = 1.0 - exp(-water_thickness * (0.13 + WATER_ABSORPTION * 0.62));
    float shallowFactor = 1.0 - clamp(water_thickness * 0.18, 0.0, 1.0);
    float ultraShallow = 1.0 - smoothstep(0.08, 1.55, water_thickness);

    // Three-stage shallow-to-deep transition so the shore does not look like a flat paint pass.
    float lagoonBand = 1.0 - smoothstep(0.55, 1.85, water_thickness);
    float shelfBand = smoothstep(0.65, 2.40, water_thickness) * (1.0 - smoothstep(2.8, 6.8, water_thickness));
    float deepBand = smoothstep(2.6, 7.6, water_thickness);

    float bottomLuma = dot(sceneColor, vec3(0.2126, 0.7152, 0.0722));
    float bottomGreen = clamp(sceneColor.g - sceneColor.b * 0.35, 0.0, 1.0);

    vec3 bottomTint = mix(sceneColor * vec3(0.66, 0.76, 0.72), sceneColor * vec3(0.58, 0.72, 0.76) + vec3(0.0, 0.010, 0.020), 0.42);
    vec3 lagoonTint = mix(bottomTint + vec3(0.030, 0.060, 0.056), vec3(0.092, 0.162, 0.148), 0.62);
    lagoonTint = mix(lagoonTint, vec3(0.078, 0.144, 0.126), bottomGreen * 0.42);
    lagoonTint *= vec3(0.94, 0.99, 0.98);
    lagoonTint = mix(lagoonTint, vec3(0.106, 0.176, 0.156), ultraShallow * 0.44);

    vec3 shelfTint = mix(sceneColor * vec3(0.46, 0.58, 0.60) + vec3(0.010, 0.042, 0.064), vec3(0.050, 0.124, 0.166), 0.52);
    shelfTint = mix(shelfTint, vec3(0.058, 0.124, 0.132), bottomGreen * 0.22);
    shelfTint *= vec3(0.92, 0.98, 1.02) + vec3(0.000, 0.004, 0.012);

    vec3 deepTint = mix(WATER_COLOR, vec3(0.026, 0.082, 0.272), 0.40);
    deepTint *= vec3(0.86, 0.91, 1.08) + vec3(0.000, 0.001, 0.014);

    vec3 rainyLagoonTint = mix(vec3(0.070, 0.116, 0.110), vec3(0.064, 0.110, 0.102), bottomGreen * 0.26);
    vec3 rainyShelfTint = mix(vec3(0.050, 0.096, 0.126), vec3(0.044, 0.082, 0.108), bottomGreen * 0.22);
    lagoonTint = mix(lagoonTint, rainyLagoonTint, weatherMute * 0.92);
    shelfTint = mix(shelfTint, rainyShelfTint, weatherMute * 0.82);
    deepTint = mix(deepTint, vec3(0.032, 0.072, 0.165), weatherMute * 0.42);

    vec3 stagedTint = lagoonTint;
    stagedTint = mix(stagedTint, shelfTint, clamp(shelfBand + depthBlend * 0.25, 0.0, 1.0));
    stagedTint = mix(stagedTint, deepTint, clamp(deepBand, 0.0, 1.0));

    // Preserve bottom detail and make the water feel dense, not painted.
    vec3 absorptionFilter = exp(-water_thickness * vec3(0.26, 0.128, 0.066) * (0.88 + WATER_ABSORPTION * 1.04));
    absorptionFilter = mix(absorptionFilter, exp(-water_thickness * vec3(0.20, 0.118, 0.086) * 1.02), shallowFactor * 0.82);
    vec3 transmittedScene = sceneColor * absorptionFilter;

    float bodyAmount = 1.0 - exp(-water_thickness * (0.132 + WATER_ABSORPTION * 0.62));
    bodyAmount += shallowFactor * mix(0.205, 0.245, dayFactorSurface);
    bodyAmount += ultraShallow * 0.270;
    bodyAmount *= mix(1.0, 0.84, weatherMute * (0.36 + 0.64 * shallowFactor));

    float preserveBottom = mix(0.58, 0.36, bottomLuma);
    preserveBottom *= mix(0.52, 0.90, 1.0 - ultraShallow);
    vec3 refractedColor = mix(transmittedScene, stagedTint, clamp(bodyAmount * preserveBottom, 0.0, 0.94));
    refractedColor = mix(refractedColor, stagedTint, ultraShallow * 0.325 * (1.0 - weatherMute * 0.18));

    float siltAmount = (shallowFactor * 0.94 + ultraShallow * 0.40) * (1.0 - clamp(depthBlend * 1.38, 0.0, 1.0));
    vec3 siltTint = mix(vec3(0.16, 0.18, 0.13), vec3(0.08, 0.16, 0.18), clamp(water_thickness * 0.32, 0.0, 1.0));
    refractedColor = mix(refractedColor, refractedColor * (0.88 + 0.12 * siltTint) + siltTint * 0.028, siltAmount * (0.13 + 0.07 * (1.0 - weatherMute)));

    // Keep a little of the authored water color, but do not let it become a solid paint pass.
    refractedColor = mix(refractedColor, color, clamp(water_absortion * 0.38, 0.0, 0.28));

    // Gentle shoreline caustics on the bottom, only in clear shallow water.
    #if !defined NETHER && !defined THE_END
        float groundDepth = texture2D(depthtex1, pos.xy).r;
        if (groundDepth < 0.9999 && isEyeInWater == 0) {
            vec3 astroLightPos = worldTime > 12900 ? moonPosition : sunPosition;
            vec3 sunDirWorld = normalize((gbufferModelViewInverse * vec4(astroLightPos, 0.0)).xyz);
            vec3 bottomWorldPos = prismaReconstructWorldPos(pos.xy, groundDepth);
            float sunAmount = clamp(sunDirWorld.y * 0.9 + 0.1, 0.0, 1.0) * (1.0 - rainStrength);
            float causticMask = sunAmount * dayFactorSurface * shallowFactor * (1.0 - clamp(depthBlend * 1.15, 0.0, 1.0));
            if (causticMask > 0.001) {
                float caustics = prismaWaterCaustics(bottomWorldPos, sunDirWorld);
                vec3 causticColor = mix(vec3(0.76, 0.86, 0.78), vec3(1.0, 0.96, 0.88), sunAmount);
                refractedColor += causticColor * caustics * causticMask * 0.128;
            }
        }
    #endif

    // Keep nighttime water calm instead of fluorescent.
    refractedColor *= mix(0.23, 1.0, dayFactorSurface);
    refractedColor *= mix(1.0, 0.80, weatherMute * (0.46 + 0.54 * shallowFactor));
    return refractedColor;
}

vec3 get_normals(vec3 bump, vec3 fragpos) {
    float NdotE = abs(dot(waterNormal, normalize(fragpos)));

    bump *= vec3(NdotE) + vec3(0.0, 0.0, 1.0 - NdotE);

    mat3 tbn_matrix = mat3(
        tangent.x, binormal.x, waterNormal.x,
        tangent.y, binormal.y, waterNormal.y,
        tangent.z, binormal.z, waterNormal.z
    );

    return normalize(bump * tbn_matrix);
}

vec4 reflection_calc(vec3 fragpos, vec3 normal, vec3 reflected, inout float infinite, float dither) {
    #if SSR_TYPE == 0  // Flipped image
        #if defined DISTANT_HORIZONS
            vec3 pos = camera_to_screen(fragpos + reflected * 768.0);
        #else
            vec3 pos = camera_to_screen(fragpos + reflected * 76.0);
        #endif
    #else  // Raymarch
        vec3 pos = fastRaymarch(reflected, fragpos, infinite, dither);
    #endif

    float border =
        clamp((1.0 - (max(0.0, abs(pos.y - 0.5)) * 2.0)) * 50.0, 0.0, 1.0);

    border = clamp(border - pow(pos.y, 10.0), 0.0, 1.0);

    pos.x = abs(pos.x);
    if (pos.x > 1.0) {
        pos.x = 1.0 - (pos.x - 1.0);
    }

    return vec4(texture2D(gaux1, pos.xy).rgb, border);
}

vec3 water_shader(
    vec3 fragpos,
    vec3 normal,
    vec3 color,
    vec3 sky_reflect,
    vec3 reflected,
    float fresnel,
    float visibleSky,
    float dither,
    vec3 lightColor
) {
    vec4 reflection = vec4(0.0);
    float infinite = 1.0;

    #if REFLECTION == 1
        reflection =
            reflection_calc(fragpos, normal, reflected, infinite, dither);
    #endif

    reflection.rgb = mix(
        sky_reflect * visibleSky,
        reflection.rgb,
        reflection.a
    );
    reflection.rgb *= mix(vec3(1.0), vec3(0.82, 0.88, 0.94), rainStrength * 0.45);

    #ifdef VANILLA_WATER
        fresnel *= 0.8;
    #endif

    #if SUN_REFLECTION == 1
        #ifndef NETHER
            #ifndef THE_END
                return mix(color, reflection.rgb, fresnel * (REFLEX_INDEX + 0.05)) +
                    vec3(sun_reflection(reflect(normalize(fragpos), normal))) * lightColor * infinite * visibleSky * (1.0 - rainStrength * 0.65);          
            #else
                return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
            #endif
        #else
            return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
        #endif
    #else
        return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
    #endif
}

//  GLASS

vec4 cristal_reflection_calc(vec3 fragpos, vec3 normal, inout float infinite, float dither) {
    #if SSR_TYPE == 0
        #if defined DISTANT_HORIZONS
            vec3 reflectedVector = reflect(normalize(fragpos), normal) * 768.0;
        #else
            vec3 reflectedVector = reflect(normalize(fragpos), normal) * 76.0;
        #endif
            vec3 pos = camera_to_screen(fragpos + reflectedVector);
    #else
        vec3 reflectedVector = reflect(normalize(fragpos), normal);
        vec3 pos = fastRaymarch(reflectedVector, fragpos, infinite, dither);

        if (pos.x > 99.0) { // Fallback
            #if defined DISTANT_HORIZONS
                pos = camera_to_screen(fragpos + reflectedVector * 768.0);
            #else
                pos = camera_to_screen(fragpos + reflectedVector * 76.0);
            #endif
        }
    #endif

    float border_x = max(-fourthPow(abs(2.0 * pos.x - 1.0)) + 1.0, 0.0);
    float border_y = max(-fourthPow(abs(2.0 * pos.y - 1.0)) + 1.0, 0.0);
    float border = min(border_x, border_y);

    return vec4(texture2D(gaux1, pos.xy, 0.0).rgb, border);
}

vec4 cristal_shader(
    vec3 fragpos,
    vec3 normal,
    vec4 color,
    vec3 skyReflectionColor,
    float fresnel,
    float visibleSky,
    float dither,
    vec3 lightColor
) {
    vec4 reflection = vec4(0.0);
    float infinite = 0.0;

    #if REFLECTION == 1
        reflection = cristal_reflection_calc(fragpos, normal, infinite, dither);
    #endif

    skyReflectionColor = mix(color.rgb, skyReflectionColor, visibleSky * visibleSky);

    reflection.rgb = mix(
        skyReflectionColor,
        reflection.rgb,
        reflection.a
    );

    color.rgb = mix(color.rgb, skyReflectionColor, fresnel);
    color.rgb = mix(color.rgb, reflection.rgb, fresnel);

    color.a = mix(color.a, 1.0, fresnel * .9);

    #if SUN_REFLECTION == 1
        #ifndef NETHER
        #ifndef THE_END
            return color + vec4(
                mix(
                    vec3(sun_reflection(reflect(normalize(fragpos), normal)) * lightColor * infinite * visibleSky),
                    vec3(0.0),
                    reflection.a
                ),
                0.0
            );
        #else
            return color;
        #endif
        #else
            return color;
        #endif
    #else
        return color;
    #endif
}
