

vec3 rayTrace(vec3 dir, vec3 position, float dither, float rayQuality)
{
    float quality = 32;

    vec3 clipPosition = toClipSpace3(position);
    vec3 direction = normalize(toClipSpace3(position + dir) - clipPosition);
    vec3 maxLengths = (step(0.0, direction) - clipPosition) / direction;
    float minLength = min(min(maxLengths.x, maxLengths.y), maxLengths.z);
    vec3 stepv = (direction * minLength) / rayQuality;

    float tolerance = max(abs(stepv.z) * 4, 0.02 / (position.z * position.z));

    for (int i = 0; i <= int(rayQuality); ++i)
    {

        vec3 spos = clipPosition + stepv * (float(i) + dither);
        float depthSample = texelFetch2D(depthtex0, ivec2(spos.xy / texelSize), 0).r;
        if (depthSample < spos.z && abs(spos.z - depthSample) < tolerance)
        {
            return vec3(spos.xy, depthSample);
        }
    }

    return vec3(1.1); // Indicates no intersection found
}

float linZ(float depth)
{
    return (2.0 * near) / (far + near - depth * (far - near));
    // l = (2*n)/(f+n-d(f-n))
    // f+n-d(f-n) = 2n/l
    // -d(f-n) = ((2n/l)-f-n)
    // d = -((2n/l)-f-n)/(f-n)
}

float invLinZ(float lindepth)
{
    return -((2.0 * near / lindepth) - far - near) / (far - near);
}

vec3 RT5(vec3 dir, vec3 position, float noise, float roughness)
{

    float ssptbias = 0.05;

    float stepSize = 32;
    // float stepSize = 6;

    int maxSteps = int(16);

    vec3 clipPosition = toClipSpace3(position);

    float rayLength = ((position.z + dir.z * sqrt(3.0) * far) > -sqrt(3.0) * near) ? (-sqrt(3.0) * near - position.z) / dir.z : sqrt(3.0) * far;

    vec3 end = toClipSpace3(position + dir * rayLength);
    vec3 direction = end - clipPosition; // convert to clip space

    float len = max(abs(direction.x) / texelSize.x, abs(direction.y) / texelSize.y) / stepSize;

    // get at which length the ray intersects with the edge of the screen
    vec3 maxLengths = (step(0., direction) - clipPosition) / direction;
    float mult = min(min(maxLengths.x, maxLengths.y), maxLengths.z);
    vec3 stepv = direction / len;
    int iterations = min(int(min(len, mult * len) - 2), maxSteps);

    // Do one iteration for closest texel (good contact shadows)
    vec3 spos = clipPosition + stepv / stepSize * 6;
    spos.xy += offsets[framemod8] * texelSize * 0.5;

    spos += stepv * noise;

    for (int i = 0; i < iterations; i++)
    {
        if (clamp(clipPosition.xy, 0, 1) != clipPosition.xy)
            break;

        float sp = linZ(texelFetch(depthtex0, ivec2(spos.xy / texelSize), 0).r);

        float currZ = linZ(spos.z);
        if (sp < currZ)
        {
            if (spos.x < 0.0 || spos.y < 0.0 || spos.z < 0.0 || spos.x > 1.0 || spos.y > 1.0 || spos.z > 1.0)
                return vec3(1.1);
            float dist = abs(sp - currZ) / currZ;

            if (dist <= ssptbias)
                return vec3(spos.xy, invLinZ(sp));
        }

        spos += stepv;
    }

    return vec3(1.1);
}

vec3 rayTraceDiffuse(vec3 dir, vec3 position, float dither, float roughness)
{
    float quality = 32;

    vec3 clipPosition = toClipSpace3(position);
    vec3 direction = normalize(toClipSpace3(position + dir) - clipPosition);
    vec3 maxLengths = (step(0.0, direction) - clipPosition) / direction;
    float minLength = min(min(maxLengths.x, maxLengths.y), maxLengths.z);
    vec3 stepv = (direction * minLength) / quality;

    float tolerance = max(abs(stepv.z) * 4, 0.02 / (position.z * position.z));

    for (int i = 0; i <= int(quality); ++i)
    {

        vec3 spos = clipPosition + stepv * (float(i) + dither);
        float depthSample = texelFetch2D(depthtex0, ivec2(spos.xy / texelSize), 0).r;
        if (depthSample < spos.z && abs(spos.z - depthSample) < tolerance)
        {
            return vec3(spos.xy, depthSample);
        }
    }

    return vec3(1.1); // Indicates no intersection found
}

// Adapted from "Sampling the GGX Distribution of Visible Normals",
// http://jcgt.org/published/0007/04/01/
vec3 SampleGGXVNDF(vec3 Ve, float alpha_x, float alpha_y, float U1, float U2)
{
    // Transforming the view direction to the hemisphere configuration
    vec3 Vh = normalize(vec3(alpha_x * Ve.x, alpha_y * Ve.y, Ve.z));

    // Orthonormal basis
    vec3 T1 = (Vh.z < 1.0f) ? normalize(cross(vec3(0, 0, 1), Vh)) : vec3(1, 0, 0);
    vec3 T2 = cross(Vh, T1);

    // Parameterization of the projected area
    float r = sqrt(U1);
    float phi = 2.0f * PI * U2;
    float t1 = r * cos(phi);
    float t2 = r * sin(phi);
    float s = 0.5f * (1.0f + Vh.z);
    t2 = (1.0f - s) * sqrt(1.0f - t1 * t1) + s * t2;

    // Reprojection onto hemisphere
    vec3 Nh = t1 * T1 + t2 * T2 + sqrt(max(0.0f, 1.0f - t1 * t1 - t2 * t2)) * Vh;

    // Transform the normal back to the ellipsoid configuration
    vec3 Ne = normalize(vec3(alpha_x * Nh.x, alpha_y * Nh.y, max(0.0f, Nh.z)));

    return Ne;
}
void frisvad(in vec3 n, out vec3 f, out vec3 r)
{
    if (n.z < -0.999999)
    {
        f = vec3(0., -1, 0);
        r = vec3(-1, 0, 0);
    }
    else
    {
        float a = 1. / (1. + n.z);
        float b = -n.x * n.y * a;
        f = vec3(1. - n.x * n.x * a, b, -n.x);
        r = vec3(b, 1. - n.y * n.y * a, -n.y);
    }
}
mat3 CoordBase(vec3 n)
{
    vec3 x, y;
    frisvad(n, x, y);
    return mat3(x, y, n);
}
vec2 R2_samples(int n)
{
    vec2 alpha = vec2(0.75487765, 0.56984026);
    return fract(alpha * n);
}
float square(float x)
{
    return x * x;
}

// GGX distribution function G1
float g(float NdotL, float roughness)
{
    float k = (roughness * roughness) / 2.0;
    return NdotL / (NdotL * (1.0 - k) + k);
}

float vdc(uint bits)
{
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return float(bits) * 2.3283064365386963e-10;
}

// Hammersley point set
vec2 hs(uint i, uint N)
{
    return vec2(float(i) / float(N), vdc(i));
}

vec2 StratifiedSample(int i, int N)
{
    int sqrtN = int(sqrt(float(N)));
    int x = i % sqrtN;
    int y = i / sqrtN;
    return (vec2(x, y) + vec2(fract(sin(float(i)) * 43758.5453123), fract(cos(float(i)) * 43758.5453123))) /
           float(sqrtN);
}
vec3 importanceSampleGGX(vec2 Xi, vec3 N, float roughness, float rot)
{
    float a = roughness * roughness;

    float phi = 2.0 * PI * fract(Xi.x + rot);
    float cosTheta = sqrt((1.0 - Xi.y) / (1.0 + (a * a - 1.0) * Xi.y));
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);

    // From spherical coordinates to cartesian coordinates
    vec3 H;
    H.x = cos(phi) * sinTheta;
    H.y = sin(phi) * sinTheta;
    H.z = cosTheta;

    // From tangent-space vector to world-space sample vector
    vec3 up = abs(N.z) < 0.5 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    vec3 tangent = normalize(cross(up, N));
    vec3 bitangent = cross(N, tangent);

    vec3 sampleVec = tangent * H.x + bitangent * H.y + N * H.z;

    return normalize(sampleVec);
}
// const bool colortex4MipmapEnabled = true;

////

vec3 GetN(int idx)
{
    switch (idx)
    {
    case 230:
        return vec3(2.9114, 2.9497, 2.5845);
    case 231:
        return vec3(0.18299, 0.42108, 1.3734);
    case 232:
        return vec3(1.3456, 0.96521, 0.61722);
    case 233:
        return vec3(3.1071, 3.1812, 2.3230);
    case 234:
        return vec3(0.27105, 0.67693, 1.3164);
    case 235:
        return vec3(1.9100, 1.8300, 1.4400);
    case 236:
        return vec3(2.3757, 2.0847, 1.8453);
    case 237:
        return vec3(0.15943, 0.14512, 0.13547);
    default:
        return vec3(0.0);
    }
}

vec3 GetK(int idx)
{
    switch (idx)
    {
    case 230:
        return vec3(3.0893, 2.9318, 2.7670);
    case 231:
        return vec3(3.4242, 2.3459, 1.7704);
    case 232:
        return vec3(7.4746, 6.3995, 5.3031);
    case 233:
        return vec3(3.3314, 3.3291, 3.1350);
    case 234:
        return vec3(3.6092, 2.6248, 2.2921);
    case 235:
        return vec3(3.5100, 3.4000, 3.1800);
    case 236:
        return vec3(4.2655, 3.7153, 3.1365);
    case 237:
        return vec3(3.9291, 3.1900, 2.3808);
    default:
        return vec3(1.0);
    }
}

float fresnel0(vec3 l, vec3 n, float F0)
{
    // Assume eta1 is 1.0 (air)
    float eta1 = (-1.0 - sqrt(F0)) / (sqrt(F0) - 1.0);

    // Compute eta2 from F0
    float eta2 = 1;

    float cosThetaI = abs(dot(l, n)); // cos(theta_i) - angle of incidence
    float eta = eta2 / eta1;          // relative index of refraction
    float sin2ThetaI = 1.0 - cosThetaI * cosThetaI;
    float sin2ThetaT = eta * eta * sin2ThetaI; // sin^2(theta_t) using Snell's law

    float cosThetaT = sqrt(max(0.0, 1.0 - sin2ThetaT)); // cos(theta_t)

    // Reflectance for s-polarized light
    float Rs = (eta1 * cosThetaI - eta2 * cosThetaT) / (eta1 * cosThetaI + eta2 * cosThetaT);
    Rs = Rs * Rs;

    // Reflectance for p-polarized light
    float Rp = (eta2 * cosThetaI - eta1 * cosThetaT) / (eta2 * cosThetaI + eta1 * cosThetaT);
    Rp = Rp * Rp;

    // Average of the s and p polarization reflectances
    return 0.5 * (Rs + Rp);
}
vec3 ComplexFresnel(vec3 fragpos, vec3 normal, float f0)
{
    int metalidx = int((f0 * 255));

    float fresnel = fresnel0(fragpos, normal, f0);
    if (metalidx > 229.5)
    {
        vec3 k = GetK(metalidx);
        vec3 n = GetN(metalidx);
        float f = 1.0 - fresnel;

        vec3 k2 = k * k;
        vec3 n2 = n * n;
        float f2 = f * f;

        vec3 rs_num = n2 + k2 - 2 * n * f + f2;
        vec3 rs_den = n2 + k2 + 2 * n * f + f2;
        vec3 rs = rs_num / rs_den;

        vec3 rp_num = (n2 + k2) * f2 - 2 * n * f + 1;
        vec3 rp_den = (n2 + k2) * f2 + 2 * n * f + 1;
        vec3 rp = rp_num / rp_den;

        vec3 fresnel3 = clamp(0.5 * (rs + rp), vec3(0.0), vec3(1.0));
        fresnel3 *= fresnel3;
        return fresnel3;
    }
    else
    {
        return vec3(fresnel);
    }
}
///////////

float gSimple(float dp, float roughness)
{
    float k = roughness + 1;
    k *= k / 8.0;
    return dp / (dp * (1.0 - k) + k);
}

vec3 GGX2(vec3 n, vec3 v, vec3 l, float r, vec3 F0)
{
    float alpha = square(r);

    vec3 h = normalize(l + v);

    float dotLH = clamp(dot(h, l), 0., 1.);
    float dotNH = clamp(dot(h, n), 0., 1.);
    float dotNL = clamp(dot(n, l), 0., 1.);
    float dotNV = clamp(dot(n, v), 0., 1.);
    float dotVH = clamp(dot(h, v), 0., 1.);

    float D = alpha / (3.141592653589793 * square(square(dotNH) * (alpha - 1.0) + 1.0));
    float G = gSimple(dotNV, r) * gSimple(dotNL, r);
    vec3 F = F0 + (1. - F0) * exp2((-5.55473 * dotVH - 6.98316) * dotVH);

    return dotNL * F * (G * D / (4 * dotNV * dotNL + 1e-7));
}
vec3 sampleGGXVNDF(vec3 V_, float alpha_x, float alpha_y, float U1, float U2)
{
    // stretch view
    vec3 V = normalize(vec3(alpha_x * V_.x, alpha_y * V_.y, V_.z));
    // orthonormal basis
    vec3 T1 = (V.z < 0.9999) ? normalize(cross(V, vec3(0, 0, 1))) : vec3(1, 0, 0);
    vec3 T2 = cross(T1, V);
    // sample point with polar coordinates (r, phi)
    float a = 1.0 / (1.0 + V.z);
    float r = sqrt(U1);
    float phi = (U2 < a) ? U2 / a * 3.141592653589793 : 3.141592653589793 + (U2 - a) / (1.0 - a) * 3.141592653589793;
    float P1 = r * cos(phi);
    float P2 = r * sin(phi) * ((U2 < a) ? 1.0 : V.z);
    // compute normal
    vec3 N = P1 * T1 + P2 * T2 + sqrt(max(0.0, 1.0 - P1 * P1 - P2 * P2)) * V;
    // unstretch
    N = normalize(vec3(alpha_x * N.x, alpha_y * N.y, max(0.0, N.z)));
    return N;
}

vec3 calculateSpecular(vec3 f0, float roughness, float lightmapY, vec3 normal, vec3 viewDir, vec3 fragPos, float noise, vec3 fresnel, vec3 scene)
{
    vec3 fresnelDiffuse = vec3(0.0);

    vec3 indirectSpecular = vec3(0.0);
    const int numSpecularSamples = REFLECTION_SAMPLES;
    mat3 basis = CoordBase(normal);
    vec3 viewDirInNormalSpace = -viewDir * basis;

    for (int sampleIndex = 0; sampleIndex < numSpecularSamples; sampleIndex++)
    {
        // Seed calculation for random sampling
        int seed = frameCounter * numSpecularSamples + sampleIndex;

        // Generate random samples
        vec2 randomSample = fract(R2_samples(seed) + blueNoise4(gl_FragCoord.xy).rg);

        // Sample GGX VNDF and importance sampling
        vec3 halfVector1 = sampleGGXVNDF(viewDirInNormalSpace, roughness, roughness, randomSample.x, randomSample.y);
        vec3 halfVector2 = importanceSampleGGX(randomSample, vec3(0.0, 0.0, 1.0), roughness, 0.0);

        // Mix between two half-vectors
        vec3 halfVector = mix(halfVector2, halfVector1, 0);

        // Reflect view vector around half-vector
        vec3 reflectedDirInNormalSpace = reflect(-viewDirInNormalSpace, halfVector);
        vec3 reflectedDir = basis * reflectedDirInNormalSpace;

        // Calculate geometric attenuation
        float geomAttenuation = g(clamp(dot(normal, reflectedDir), 0.0, 1.0), roughness);

        // Calculate Fresnel reflection
        vec3 fresnelReflectance = f0 + (1.0 - f0) * pow(clamp(1.0 + dot(-reflectedDirInNormalSpace, halfVector), 0.0, 1.0), 5);
        // vec3 rayContribution = fresnelReflectance * geomAttenuation;
        vec3 rayContribution = fresnel * geomAttenuation;

        // Only consider significant ray contributions
        if (luma(rayContribution) > 0.1)
        {
            vec4 reflection = vec4(0.0);
            float rayQuality = 35.0 * sqrt(luma(rayContribution));

            // Perform ray tracing if ray quality is high enough
            if (rayQuality > 5.0)
            {
                vec3 rayTracePos = RT5(mat3(gbufferModelView) * reflectedDir, fragPos, noise, rayQuality);

                // Check if the ray hit a surface
                if (rayTracePos.z < 1.0)
                {

                    reflection.a = 1.0;
                    reflection.rgb = texture2D(colortex0, rayTracePos.xy).rgb;
                }
            }

            // If no valid reflection, fetch from the sky texture
            if (reflection.a < 1)
            {

                float bounceAmount = float(reflectedDir.y > 0.0) + clamp(-reflectedDir.y * 0.1 + 0.1, 0.0, 1.0);

                vec4 skyColor = skyFromTex2(reflectedDir, colortex6) * bounceAmount;
                reflection.rgb += (skyColor.rgb + skyColor.a * sunLight * lightmapY) * pow(lightmapY, 8.0);
            }

            // Accumulate the indirect specular contribution
            indirectSpecular += reflection.rgb * pow(rayContribution, vec3(1.5));
            fresnelDiffuse += rayContribution;
        }
    }

    // Average the indirect specular contribution over all samples
    return (indirectSpecular / numSpecularSamples);
    // return (indirectSpecular / numSpecularSamples) + (1.0 - fresnelDiffuse / numSpecularSamples) * scene;
}

vec3 computeReflections(float lightmapY, float smoothness, vec3 f0, vec3 normal, vec3 fragpos, float materialsX)
{
    vec3 scene = texture(colortex0, texcoord).rgb;

    vec3 p3 = mat3(gbufferModelViewInverse) * fragpos;
    int metalidx = int((f0 * 255));
    if (metalidx > 229.5)
    {
        // scene = vec3(0);
    }
    vec3 reflections = vec3(0);
    float noise = blueNoise(gl_FragCoord.xy);
    float f0_orig = f0.y;
    float porosity = (materialsX * 255 <= 64.5) ? step(0.0, materialsX) : 0.0;
    float rainMult = sqrt(lightmapY) * wetness * (1.0 - (porosity * porosity));

    vec3 fresnel = ComplexFresnel(normalize(fragpos), worldToView(normal), f0_orig);
    float roughness = clamp(pow(1.0 - smoothness, 2.0), 0.0001, 0.9999) * 1;
    roughness = mix(roughness, 0.01, rainMult);
    reflections = calculateSpecular(f0, roughness, lightmapY, normal, normalize(p3), fragpos, noise, fresnel, scene) * fresnel;

    return reflections;
}
