#version 450

layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;

#include<light_structs.h>
#include<poisson.h>

layout(set = 0, binding = 0) uniform sampler2D in_depth; 
layout(set = 0, binding = 1) uniform sampler2D in_norm_rough;
layout(set = 0, binding = 2) writeonly uniform image2D out_diffuse;
layout(set = 0, binding = 3) writeonly uniform image2D out_specular;

layout(std430, set = 0, binding = 4) buffer DirectionalLightData 
{
    int n;
    DirectionalLight directional_lights[10];
}directional_light_data;

layout(std430, set = 0, binding = 5) buffer PointLightData 
{
    PointLight point_lights[10];
}point_light_data;

layout(std430, set = 0, binding = 6) buffer SpotLightData 
{
    SpotLight spot_lights[10];
}spot_light_data;

layout(std430, set = 0, binding = 7) buffer AreaLightData
{
    AreaLight area_lights[10];
}area_light_data;

layout(std430, set = 0, binding = 8) buffer ReflectionProbeData 
{
    ReflectionProbe probes[10];
}reflection_probe_data;

/*layout(constant_id = 0) const int tiles_count_x = 64;
layout(constant_id = 1) const int tiles_count_y = 64;*/
const int ntiles_x = 120,  ntiles_y = 68;
layout(set = 0, binding = 9) buffer TileData 
{
    Tile tiles[ntiles_x][ntiles_y];
    Item items[81600];
}tile_data;

layout(set = 0, binding = 10) buffer Res
{
    uvec2 res;
}res;

layout(set = 0, binding = 11) uniform sampler2DShadow shadow_atlas;
layout(set = 0, binding = 12) uniform sampler2D shadow_atlas_raw;
layout(set = 0, binding = 13) buffer ShadowData 
{
    mat4 shadow_transforms[16];
    vec4 shadow_maps[16];
} shadow_data;

layout(push_constant) uniform ProjectionData
{
    mat4 mv;
    mat4 p;
    //vec3 view_pos;
} projection_data;

#include<common.h>

#include<pbr.h>

void pbr_light(out vec3 specular, out vec3 diffuse, vec3 point_to_light, vec3 color, float attenuation, vec3 norm, vec3 point, float roughness)
{
    //diffuse = vec3(sin(3.1415*point.yyy));return;
    vec3 vdir = normalize(get_view_pos() - point);
    vec3 halfway = normalize(point_to_light + vdir); 
    float cos_theta = max(dot(norm, point_to_light), 0.0);

    vec3 radiance = color.xyz * attenuation;
    float F0 = 0.04;
    float fresnel = fresnel_schlick(max(dot(halfway, vdir), 0.0), F0);
    
    float ndf = normal_distribution_ggx(norm, halfway, roughness);
    float g = geometry_smith(norm, vdir, point_to_light, k_direct(roughness));
    //cook-torrance
    float specular_factor = ndf * g * fresnel / max(4.0 * max(dot(norm, vdir) * dot(norm, point_to_light), 0.0), 0.001);

    float ks = fresnel,
          kd = 1.0-ks;

    diffuse = (kd / PI) * radiance * cos_theta;
    specular = specular_factor * radiance * cos_theta;
}

/*****************************************light processing functions***************************************************/

#define CUTOFF_MARGIN 5

float cutoff_function(float dist, float influence_radius)
{
    return square(clamp((CUTOFF_MARGIN-CUTOFF_MARGIN*dist/influence_radius), 0.0, 1.0));
}

void process_directional(out vec3 specular, out vec3 diffuse, DirectionalLight light, vec3 norm, vec3 point, float roughness)
{
    vec3 point_to_light = normalize(-light.direction.xyz);
    vec3 vdir = normalize(get_view_pos() - point);
    vec3 halfway = normalize(vdir+norm); 
    point_to_light = normalize(point_to_light);
    pbr_light(specular, diffuse, point_to_light, light.color.xyz, 1.0, norm, point, roughness);
}

void process_point(out vec3 specular, out vec3 diffuse, PointLight light, vec3 norm, vec3 point, float roughness)
{
    vec3 point_to_light = light.position.xyz-point;
    float dist = length(point_to_light);
    vec3 vdir = normalize(get_view_pos() - point);
    vec3 halfway = normalize(vdir+norm); 
    point_to_light = normalize(point_to_light);
    float attenuation = cutoff_function(dist, light.influence_radius) / square(dist);
    pbr_light(specular, diffuse, point_to_light, light.color.xyz, attenuation, norm, point, roughness);
}

void process_spot(out vec3 specular, out vec3 diffuse, SpotLight light, vec3 norm, vec3 point, float roughness)
{
    vec3 point_to_light = light.position.xyz-point;
    float dist = length(point_to_light);
    float attenuation = light.size*light.size/square(tan(light.cone_angle)*dist);
    vec3 vdir = normalize(get_view_pos() - point);
    vec3 halfway = normalize(vdir+norm); 
    point_to_light = normalize(point_to_light);
    attenuation *= clamp((dot(-light.direction, normalize(point_to_light))-cos(light.cone_angle))/cos(light.cone_angle*0.1), 0.0, 1.0);
    pbr_light(specular, diffuse, point_to_light, light.color.xyz, attenuation, norm, point, roughness);
}

void process_area(out vec3 specular, out vec3 diffuse, AreaLight light, vec3 norm, vec3 point, float roughness)
{
    specular = vec3(1.0);
    diffuse = vec3(1.0);
}
void process_probe(out vec3 specular, out vec3 diffuse, ReflectionProbe light, vec3 norm, vec3 point, float roughness)
{
    specular = vec3(1.0);
    diffuse = vec3(1.0);
}
/*****************************************light processing functions***************************************************/

/**********************************************shadow mapping**********************************************************/
#define CASCADE_MARGIN 0.01
#define PCF_SAMPLES 32

vec3 get_shadow_coords(vec4 wp_coord, mat4 mat, vec4 shadow_map, vec3 norm, vec3 ldir)
{
    //vec3 bias = ldir * max(0.005 * (1 - dot(norm, ldir)), shadow_map.w);
    //vec4 r = mat * (wp_coord + vec4(bias, 0.0));
    vec4 r = mat * wp_coord;
    r.xyz /= r.w;
    r.xy = (r.xy+1)/2.0;
    r.z -=  max(0.005 * (1 - dot(norm, ldir)), shadow_map.w);
    r.xy *= shadow_map.z;
    r.xy += shadow_map.xy;
    return clamp(r.xyz, vec3(0.0), vec3(1.0));
}

vec3 get_shadow_uv(vec4 wp_coord, uint shadow_index) 
{
    mat4 mat = shadow_data.shadow_transforms[shadow_index];
    vec4 shadow_map = shadow_data.shadow_maps[shadow_index];
    vec4 r = mat * wp_coord;
    r.xyz /= r.w;
    r.xy = (r.xy + 1) / 2.0;
    r.z -=  shadow_map.w;
    return r.xyz;
}

vec2 get_atlas_coord(vec2 uv, uint shadow_index)
{
    vec4 shadow_map = shadow_data.shadow_maps[shadow_index];
    uv *= shadow_map.z;
    uv += shadow_map.xy;
    return uv;
}

bool check_shadow_range(vec4 wp_coord, mat4 mat) 
{
    vec4 r = mat * wp_coord;
    r.xyz /= r.w;

    return r.x < (1.0-CASCADE_MARGIN) && r.x > -(1.0-CASCADE_MARGIN) &&
           r.y < (1.0-CASCADE_MARGIN) && r.y > -(1.0-CASCADE_MARGIN) &&
           r.z < (1.0-CASCADE_MARGIN) && r.z > -(1.0-CASCADE_MARGIN);
}

uint find_shadow_cascade(vec3 point, uint shadow_index, uint n_cascade) 
{
    for (int cascade = 1; 

        cascade <= n_cascade && 
                !check_shadow_range(vec4(point, 1.0), shadow_data.shadow_transforms[shadow_index]); 

        cascade++, shadow_index++
    );

    return shadow_index;
}

void find_blocker(vec3 uv, float light_angle, out float avg_depth, out uint n_blockers, uint shadow_index)
{
    float search_width = light_angle * 0.2;

    float blocker_sum = 0.0;
    n_blockers = 0;

    for(int i = 0; i < 16; i++) 
    {
        vec2 coords = uv.xy + poisson_disk_border[i] * search_width;
        vec2 atlas_coords = get_atlas_coord(coords, shadow_index);
        float shadow_map_depth = texture(shadow_atlas_raw, atlas_coords).r;

        if(shadow_map_depth < uv.z)
        {
            blocker_sum += shadow_map_depth;
            n_blockers++;
        }
    }
    avg_depth = blocker_sum / n_blockers;
}

float pcf(vec3 uv, float radius, uint shadow_index) 
{
    float shadow_fact = 0.0;

    for(int i = 0;i < PCF_SAMPLES; i++)
    {
        vec3 atlas_coords = vec3(get_atlas_coord(uv.xy + radius * poisson_disk[i], shadow_index), uv.z);
        shadow_fact += texture(shadow_atlas, 
                               atlas_coords);
    }
    return shadow_fact / PCF_SAMPLES;
}

float directional_shadows(vec3 point, uint light_index) 
{
    float light_angle = directional_light_data.directional_lights[light_index].angle;
    uint shadow_index = directional_light_data.directional_lights[light_index].shadows;
    if(shadow_index == 0xffffffff) {
        return 1.0;
    }
    uint n_cascades = directional_light_data.directional_lights[light_index].n_cascades;
    shadow_index = find_shadow_cascade(point, shadow_index, n_cascades);
    vec3 uv = get_shadow_uv(vec4(point, 1.0), shadow_index);

    float avg_depth = 0;
    uint n_blockers = 0;

    find_blocker(uv, light_angle, avg_depth, n_blockers, shadow_index);

    if(n_blockers < 1)
        return 1.0;

    float radius = (uv.z - avg_depth) * light_angle;
    return pcf(uv, radius, shadow_index);
}
/**********************************************shadow mapping**********************************************************/

void main()
{
    if(gl_GlobalInvocationID.x > res.res.x || gl_GlobalInvocationID.y > res.res.y)
        return;

    vec2 uv = vec2(gl_GlobalInvocationID.xy) / vec2(res.res);

    ivec2 ntiles = ivec2(ceil(res.res.x/16), ceil(res.res.y/16));
    Tile tile = tile_data.tiles[int(uv.x*ntiles.x)][int(uv.y*ntiles.y)];
    int point_count, spot_count, area_count, probe_count;
    point_count = tile.point_spot_area_probe_count & 0xff;
    spot_count = tile.point_spot_area_probe_count >> 8 & 0xff;
    area_count = tile.point_spot_area_probe_count >> 16 & 0xff;
    probe_count = tile.point_spot_area_probe_count >> 24 & 0xff;

    vec3 norm = texelFetch(in_norm_rough, ivec2(gl_GlobalInvocationID.xy), 0).xyz*2.0-1.0;
    float roughness = texture(in_norm_rough, uv).w/4.0; 
    float depth = texelFetch(in_depth, ivec2(gl_GlobalInvocationID.xy), 0).x;
    vec3 point = world_space_location(uv, depth);

    if(depth >= 1.0) return;

    /****************light processing****************/
    vec3 specular = vec3(0.0), diffuse = vec3(0.0);

    //processes directional lights
    for(int i = 0; i < directional_light_data.n; i++)
    {
        float shadow_fact = directional_shadows(point, i);
        vec3 tspec, tdiff;
        process_directional(tspec, tdiff, directional_light_data.directional_lights[i], norm, point, roughness);
        specular += tspec * shadow_fact;
        diffuse += tdiff * shadow_fact;
    }
    
    //processes point lights
    for(int i = 0;i < point_count;i++)
    {
        int idx = tile_data.items[i+tile.offset].point_spot_offset >> 16 & 0xffff;
        vec3 tspec, tdiff;
        process_point(tspec, tdiff, point_light_data.point_lights[idx], norm, point, roughness);
        specular += tspec;
        diffuse += tdiff;
    }

    //processes spot lights
    for(int i = 0;i < spot_count;i++)
    {
        int idx = tile_data.items[i+tile.offset].point_spot_offset & 0xffff;
        float shadow_fact = 0.0;
        uint shadows_index = spot_light_data.spot_lights[idx].shadows;
        if(shadows_index != 0xffffffff)
        {
            for(int i = 0;i < 32; i++)
            {
                shadow_fact += texture(shadow_atlas, 
                                       get_shadow_coords(vec4(point, 1.0), 
                                                         shadow_data.shadow_transforms[shadows_index], 
                                                         shadow_data.shadow_maps[shadows_index] ,
                                                         norm, 
                                                         normalize(spot_light_data.spot_lights[idx].position-point))
                                       + vec3(poisson_disk[i]/3000.0, 0.0) 
                                       );
            }
            shadow_fact /= 32;
        }
        else
        {
            shadow_fact = 1.0;
        }
        vec3 tspec, tdiff;
        process_spot(tspec, tdiff, spot_light_data.spot_lights[idx], norm, point, roughness);
        specular += tspec * shadow_fact;
        diffuse += tdiff * shadow_fact;
    }

    //processes area lights
    for(int i = 0;i < area_count;i++)
    {
        int idx = i+tile.offset;
        vec3 tspec, tdiff;
        process_area(tspec, tdiff, area_light_data.area_lights[idx], norm, point, roughness);
        specular += tspec;
        diffuse += tdiff;
    }

    //processes probes
    for(int i = 0;i < probe_count;i++)
    {
        int idx = i+tile.offset;
        vec3 tspec, tdiff;
        //process_probe(tspec, tdiff, reflection_probe_data.probes[idx], norm, point, roughness);
        specular += tspec;
        diffuse += tdiff;
    }

    imageStore(out_diffuse,  ivec2(gl_GlobalInvocationID.xy), diffuse.xyzz);
    imageStore(out_specular, ivec2(gl_GlobalInvocationID.xy), specular.xyzz);
}