#version 460 core

const uint MAX_FRAG_TEX_UNITS = 32u;
const uint MAX_NUM_LIGHTS = 100u;
const float PI = 3.14159265358979;

in VS_OUT {
    vec3      pos;
    vec3      tan_pos;
    vec2      uv;
    flat uint draw_id;
    mat3      TBN;
}
fs_in;

out vec4 out_colour;

struct point_light {
    vec4  position;
    vec4  colour;
    float intensity;
    float radius;
};

struct cluster {
    vec4 min_point;
    vec4 max_point;
    uint count;
    uint light_indices[MAX_NUM_LIGHTS];
};

layout(std430, binding = 1) readonly buffer cluster_buf {
    cluster clusters[];
};

layout(std430, binding = 2) readonly buffer lights_buf {
    point_light point_lights[];
};

uniform mat4 view_mat;
uniform float z_near;
uniform float z_far;
uniform uvec3 grid_size;
uniform uvec2 screen_dimensions;
uniform vec3 view_pos; // Camera world position
uniform sampler2DArray packed_tex[MAX_FRAG_TEX_UNITS]; // z index = 0: Albedo, 1: Normal, 2: Metallic-Roughness

vec3 get_normal();
float trowbridge_reitz(vec3 normal, vec3 halfway, float roughness);
float schlick(float n_dot_v, float roughness);
float smith(vec3 normal, vec3 view, vec3 light, float roughness);
vec3 fresnel_schlick(float cos_theta, vec3 F0);

void main() {
    vec3 albedo     = texture(packed_tex[fs_in.draw_id % MAX_FRAG_TEX_UNITS], vec3(fs_in.uv, 0.0)).rgb;
    vec3 normal     = get_normal();
    vec3 met_rough  = texture(packed_tex[fs_in.draw_id % MAX_FRAG_TEX_UNITS], vec3(fs_in.uv, 2.0)).rgb;
    float metallic  = met_rough.b;
    float roughness = met_rough.g;

    vec3 tangent_frag_pos = fs_in.tan_pos;
    vec3 tangent_view_pos = fs_in.TBN * view_pos;
    vec3 view_dir = normalize(tangent_view_pos - tangent_frag_pos);
    
    // Normal incidence reflectance 
    vec3 F0 = mix(vec3(0.04), albedo, metallic);

    // Find light cluster index

    // Position of this fragment in view space
    vec3 view_space_pos = vec3(view_mat * vec4(fs_in.pos, 1.0));

    // Locating which cluster this fragment is part of
    uint z_tile = uint((log(abs(view_space_pos.z) / z_near) * grid_size.z) / log(z_far / z_near));
    vec2 tile_size = screen_dimensions / grid_size.xy;
    uvec3 tile = uvec3(gl_FragCoord.xy / tile_size, z_tile);
    uint tile_idx = tile.x + (tile.y * grid_size.x) + (tile.z * grid_size.x * grid_size.y);

    // Shade fragment for each light

    uint n_lights = clusters[tile_idx].count;

    vec3 light_reflected = vec3(0.0, 0.0, 0.0);

    for (uint i = 0u; i < n_lights; ++i) {
        uint light_idx = clusters[tile_idx].light_indices[i];
        point_light light = point_lights[light_idx];

        vec3 tangent_light_pos = fs_in.TBN * vec3(light.position);

        // Per-light radiance
        vec3 light_dir = normalize(tangent_light_pos - tangent_frag_pos);
        vec3 halfway = normalize(view_dir + light_dir);
        float dist = length(tangent_light_pos - tangent_frag_pos);
        float attenuation = light.intensity / (dist * dist);
        vec3 radiance = vec3(light.colour) * attenuation;

        // Cook-Torrance BRDF

        float NDF = trowbridge_reitz(normal, halfway, roughness);   
        float G   = smith(normal, view_dir, light_dir, roughness);      
        vec3  F   = fresnel_schlick(clamp(dot(halfway, view_dir), 0.0, 1.0), F0);
           
        vec3 numerator    = NDF * G * F; 
        float denominator = 4.0 * max(dot(normal, view_dir), 0.0) * max(dot(normal, light_dir), 0.0) + 0.0001; // + 0.0001 to prevent divide by zero
        vec3 specular = numerator / denominator;
        
        // kS is equal to Fresnel
        vec3 kS = F;
        // for energy conservation, the diffuse and specular light can't
        // be above 1.0 (unless the surface emits light); to preserve this
        // relationship the diffuse component (kD) should equal 1.0 - kS.
        vec3 kD = vec3(1.0) - kS;
        // multiply kD by the inverse metalness such that only non-metals 
        // have diffuse lighting, or a linear blend if partly metal (pure metals
        // have no diffuse light).
        kD *= 1.0 - metallic;	  

        // scale light by NdotL
        float n_dot_l = max(dot(normal, light_dir), 0.0);

        // add to outgoing radiance Lo
        light_reflected += (kD * albedo / PI + specular) * radiance * n_dot_l;  // note that we already multiplied the BRDF by the Fresnel (kS) so we won't multiply by kS again
    }
  
    // Don't need to gamma-correct as PNG textures are already gamma corrected (?)
    // const float screen_gamma = 2.2;
    // light_collect = pow(light_collect, vec3(1.0 / screen_gamma));

    vec3 ambient = vec3(0.03) * albedo;

    out_colour = vec4(ambient, 1.0) + vec4(light_reflected, 1.0);
}

vec3 get_normal() {
    vec3 normal_map_sample = texture(packed_tex[fs_in.draw_id % MAX_FRAG_TEX_UNITS], vec3(fs_in.uv, 1.0)).rgb;
    // Transform normal vector to range [-1, 1]
    return normalize(normal_map_sample * 2.0 - 1.0); // Tangent-space
}

vec3 cook_torrance(vec3 albedo) {
    return albedo;
}

// Normal distribution function (a.k.a. GGX)
float trowbridge_reitz(vec3 normal, vec3 halfway, float roughness) {
    float a_squared         = roughness * roughness;
    float n_dot_h           = max(dot(normal, halfway), 0.0); // how closely the normal & halfway align
    float n_dot_h_squared   = n_dot_h * n_dot_h;
	
    float nom    = a_squared;
    float denom  = (n_dot_h_squared * (a_squared - 1.0) + 1.0);
    denom        = PI * denom * denom;
	
    return nom / denom;
}

// Geometry function (self-shadowing)
float schlick(float n_dot_v, float roughness) {
    float nom   = n_dot_v;
    float denom = n_dot_v * (1.0 - roughness) + roughness;
	
    return nom / denom;
}

// Geometry approximation using Schlick-GGX
float smith(vec3 normal, vec3 view, vec3 light, float roughness) {
    float n_dot_v = max(dot(normal, view), 0.0);
    float n_dot_l = max(dot(normal, light), 0.0);
    float ggx_1 = schlick(n_dot_v, roughness);
    float ggx_2 = schlick(n_dot_l, roughness);
	
    return ggx_1 * ggx_2;
}

// Fresnel effect
vec3 fresnel_schlick(float cos_theta, vec3 F0) {
    return F0 + (1.0 - F0) * pow(clamp(1.0 - cos_theta, 0.0, 1.0), 5.0);
}