#include "utils.glsl"

in vec4 vpeye;
in vec4 vneye;
uniform bool u_shadow_receiver;
uniform float u_cascade_distances[MAX_LIGHTS * NUM_SHADOW_CASCADES];
uniform samplerCube shadowMap[MAX_LIGHTS];
uniform sampler2D shadowMap2D[MAX_LIGHTS * NUM_SHADOW_CASCADES];

const float bias_modifier[NUM_SHADOW_CASCADES] = float[NUM_SHADOW_CASCADES](0.95, 0.35, 0.20, 0.15, 0.15, 0.15);

//// From http://fabiensanglard.net/shadowmappingVSM/index.php
float shadow_vsm(float distance, vec3 dir, int light_index) {
    distance = distance/20;
    vec2 moments = texture(shadowMap[light_index], dir).rg;

    // If the shadow map is sampled outside of its bounds, return 1.0
    if (moments.x == 1.0 && moments.y == 1.0) {
        return 1.0;
    }
    
    // Surface is fully lit. as the current fragment is before the light occluder
    if (distance <= moments.x) {
        return 1.0;
    }
    
    // The fragment is either in shadow or penumbra. We now use chebyshev's upperBound to check
    // How likely this pixel is to be lit (p_max)
    float variance = moments.y - (moments.x*moments.x);
    //variance = max(variance, 0.000002);
    // variance = max(variance, 0.00002);
    variance = max(variance, 0.002);
    
    float d = distance - moments.x;
    float p_max = variance / (variance + d*d);

    return p_max;
}

float shadow_pcf(float distance, vec3 lightDir, int light_index) {
    // Determine which cascade to use
    int cascade_index = -1;
    int min_cascades_range = light_index * NUM_SHADOW_CASCADES;
    int max_cascades_range = min_cascades_range + NUM_SHADOW_CASCADES;
    for (int i = min_cascades_range; i < max_cascades_range; i++) {
        if (distance < u_cascade_distances[i]) {
            cascade_index = i;
            break;
        }
    }
    if (cascade_index == -1) {
        cascade_index = max_cascades_range - 1;
    }

    int matrix_index = cascade_index - min_cascades_range;

    vec4 fragPosLightSpace = u_lights[light_index].shadow_matrix[matrix_index] * vec4(v_position_ws, 1.0);

    // Perform perspective divide
    vec3 projCoords = fragPosLightSpace.xyz / fragPosLightSpace.w;
    
    // Transform to [0,1] range
    projCoords = projCoords * 0.5 + 0.5;

    float currentDepth = projCoords.z;

    if (currentDepth > 1.0) {
        return 1.0;
    }

    // Calculate bias
    vec3 normal = normalize(vneye.xyz);
    float bias = max(0.05 * (1.0 - dot(normal, lightDir)), 0.005);
    bias *= 1 / (u_cascade_distances[cascade_index] * bias_modifier[matrix_index]);
 
    // PCF
    float shadow = 0.0;
    vec2 texelSize = 1.0 / textureSize(shadowMap2D[cascade_index], 0);

#if 1
    for(int x = -3; x <= 3; ++x)
    {
        for(int y = -3; y <= 3; ++y)
        {
            float pcfDepth = texture(shadowMap2D[cascade_index], projCoords.xy + vec2(x, y) * texelSize * (rand(vec2(projCoords.x + x, projCoords.y + y))*0.75f + 0.25f)).r;
            shadow += currentDepth - bias > pcfDepth ? 1.0 : 0.0;        
        }    
    }
    shadow /= 36.0;
#else
    for(int x = -1; x <= 1; ++x)
    {
        for(int y = -1; y <= 1; ++y)
        {
            float pcfDepth = texture(shadowMap2D[cascade_index], projCoords.xy + vec2(x, y) * texelSize).r;
            shadow += currentDepth - bias > pcfDepth ? 1.0 : 0.0;        
        }    
    }
    shadow /= 9.0;
#endif
    return 1.0 - shadow;
}

vec4 shadowmap(in vec4 peye, in vec4 neye) {
    vec3 fragment = vec3(peye);

    float shadowFactor = 1.0;
    float totalWeight = 0.0;
    for (int i = 0; i < u_num_lights; i++) {
        light_t light = u_lights[i];
        float factor = 1.0;
        float weight = 1.0;

        if (light.processed_shadows) {
            if (light.type == LIGHT_DIRECTIONAL) {
                factor = shadow_pcf(-peye.z, light.dir, i);
                weight = 0.7;
            } else if (light.type == LIGHT_POINT || light.type == LIGHT_SPOT) {
                vec3 light_pos = (view * vec4(light.pos, 1.0)).xyz;
                vec3 dir = light_pos - fragment;
                vec4 sc = inv_view * vec4(dir, 0.0);
                factor = shadow_vsm(length(dir), -sc.xyz, i);
                weight = 1.0;
            }
            shadowFactor *= mix(1.0, factor, weight);
            totalWeight += weight;
        }
    }

    // Normalize the shadow factor based on total weight
    if (totalWeight > 0.0) {
        shadowFactor = pow(shadowFactor, 1.0 / totalWeight);
    }

    return vec4(vec3(shadowFactor), 1.0);
}

vec4 shadowing() {
    if (u_shadow_receiver) {
        return shadowmap(vpeye, vneye);
    } else {
        return vec4(1.0);
    }
}