#version 450

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

layout(set = 0, binding = 0) uniform sampler2D norm_rough;
layout(set = 0, binding = 1) uniform sampler2D in_depth;
layout(set = 0, binding = 2) uniform sampler2D hdr_color;
layout(set = 0, binding = 3) uniform sampler2D velocity_buffer;
layout(set = 0, binding = 4) writeonly uniform image2D half_res_specular;

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

#include<common.h>
#include<pbr.h>

vec2 clip_to_uv(vec2 clip)
{
    return clip / 2 - 0.5;
}

bool ray_march(ivec2 res, vec3 o, vec3 d, mat4 pv, float near_plane, float thickness, float stride, float jitter, const float max_steps, float max_distance, out vec2 hit) 
{
    float ray_length = ((o.z + d.z * max_distance) < near_plane) ? 
        (near_plane - o.z) / d.z: 
        max_distance;

    vec3 target = o + d * ray_length;
    hit = vec2(-1, -1);

    //project into screen space
    vec4 H0 = pv * vec4(o, 1.0),
         H1 = pv * vec4(target, 1.0);
    float k0 = 1.0 / H0.w,
          k1 = 1.0 / H1.w;
    vec3 Q0 = o * k0,
         Q1 = target * k1;

    //screen space endpoints
    vec2 P0 = H0.xy * k0,
         P1 = H1.xy * k1;

    float x_max = res.x + 0.5, 
          x_min = 0.5,
          y_max = res.y + 0.5,
          y_min = 0.5,
          alpha = 0;

    if (P1.y > y_max || P1.y < y_min)
        alpha = (P1.y - ((P1.y > y_max) ? y_max : y_min)) / (P1.y - P0.y);
    if ((P1.x > x_max) || (P1.x < x_min))
        alpha = max(alpha, (P1.x - ((P1.x > x_max) ? x_max : x_min)) / (P1.x - P0.x));
    P1 = mix(P1, P0, alpha);
    k1 = mix(k1, k0, alpha);
    Q1 = mix(Q1, Q0, alpha);

    vec2 delta = P1 - P0;
    P1 += vec2((dot(delta, delta) < 0.0001) ? 0.01: 0.0);
    delta = P1 - P0;

    bool permute = false;
    if(abs(delta.x) < abs(delta.y))
    {
        permute = true;
        delta = delta.yx;
        P0 = P0.yx;
        P1 = P1.yx;
    }

    float step_dir = sign(delta.x),
          invdx    = step_dir / delta.x;

    //trach the derivatives of Q and k
    vec3 dQ = (Q1 - Q0) * invdx;
    float dk = (k1 - k0) * invdx;
    vec2 dP = vec2(step_dir, delta.y * invdx);

    dP *= stride; dQ *= stride; dk *= stride;

    P0 += dP * jitter; Q0 += dQ * jitter; k0 += dk * jitter;
    float prev_z_max_estimate = o.z;

    //slide ...
    vec3 Q = Q0;
    float k = k0,
          step_count = 0.0,
          end = P1.x * step_dir;
    for(vec2 P = P0;
        ((P.x * step_dir) <= end) && (step_count < max_steps);
        P += dP, Q.z += dQ.z, k += dk, step_count += 1.0)
    {
        //project from homogeneuous to camera space
        hit = permute ? P.yx : P;
        float ray_z_min = prev_z_max_estimate;
        float ray_z_max = (dQ.z * 0.5 + Q.z) / (dk * 0.5 + k);
        prev_z_max_estimate = ray_z_max;
        if(ray_z_min > ray_z_max) swap(ray_z_min, ray_z_max);

        float scene_z_max = texelFetch(in_depth, ivec2(hit * 2), 0).x;
        float ssr_bias = scene_z_max; ssr_bias *= ssr_bias;
        vec2 scene_z_max_h = (inverse(pv) * vec4(0, 0, scene_z_max, 1.0)).zw; //todo optimize
        scene_z_max = scene_z_max_h.x / scene_z_max_h.y;

        float scene_z_min = scene_z_max - thickness;

        float d = ray_z_max - scene_z_max;
        //if(step_count > 150 && d > 0 && d < thickness)
        if((ray_z_max + ssr_bias >= scene_z_min) && (ray_z_min + ssr_bias <= scene_z_max))
            return all(lessThanEqual(abs(hit - (res * 0.5)), (res * 0.5)));
    }
    return false;
    
}

float fresnel(vec3 normal, vec3 vdir)
{
    float F0 = 0.04;
    return fresnel_schlick(max(dot(normal, vdir), 0.0), F0);
}

void main()
{
    ivec2 res = imageSize(half_res_specular);

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

    vec2 uv = vec2(gl_GlobalInvocationID.xy) / vec2(res);
    float depth = texture(in_depth, uv).x;
    vec3 point = world_space_location(uv, depth);

    vec3 normal = texture(norm_rough, uv).xyz * 2.0 - 1.0;

    mat4 p = projection_data.p;
    p[0][3] = 0.0;
    p[1][3] = 0.0;
    p[2][3] = -1.0;
    //ssr
    /*vec4 camera = inverse(p*projection_data.mv)*vec4(0.0, 0.0, 0.0, 1.0);
    camera.xyz /= camera.w; //not necessary*/
    vec3 vdir = normalize(point - get_view_pos());
    vec3 reflected = reflect(vdir, normal);
    /*vec4 clip_target = (p*projection_data.mv) * vec4(point + reflected, 1.0);
    clip_target.xyz /= clip_target.w;
    float step_ratio = 1 / length(clip_target.xyz - point);
    //vpr /= vpr.w; //useless since we normalize
    vec3 cpos = vec3(uv * 2 - 1, depth);
    vec2 unit_dir = normalize(clip_target.xy - cpos.xy);
    float start_reciprocal = 1 / point.z;
    float target_reciprocal = 1 / clip_target.z;
    //  ray march
    float step_len = 0.01;//1/ res.x * 1.41;
    bool hit = false;
    cpos.xy += step_len * unit_dir * 3;
    for(int i = 3; i < 100; i++) 
    {
        if(texture(in_depth, clip_to_uv(cpos.xy)).x > 1/mix(start_reciprocal, target_reciprocal, i * step_len) 
                && cpos.x < 1.0 && cpos.y < 1.0 && cpos.x > -1.0 && cpos.y > -1.0) 
        {
            cpos.xy += step_len * unit_dir;
        }
        else
        {
            hit = cpos.x < 1.0 && cpos.y < 1.0 && cpos.x > -1.0 && cpos.y > -1.0;
        }
    }*/
    mat4 clip_ro_screen = mat4(
        vec4(res.x / 2,     0,      0, 0),
        vec4(0,         res.y / 2,  0, 0),
        vec4(0,             0,      1, 0),
        vec4(res.x / 2, res.y /2,   0, 1)
    );
    vec2 hitp;
    vec4 point_vs = (projection_data.mv * vec4(point, 1.0)),
         refl_vs  = (projection_data.mv * vec4(point + reflected, 1.0));

    point_vs.xyz /= point_vs.w;
    refl_vs.xyz /= refl_vs.w;

    float near_plane = 0.1;
    /*vec2 cnp = (inverse(p) * vec4(0, 0, near_plane, 1.0)).zw;
    near_plane = cnp.x / cnp.y;*/

    bool hit = ray_march(res, point_vs.xyz, normalize(point_vs.xyz - refl_vs.xyz),
        clip_ro_screen*p, near_plane, 1.5, 4.0, 0, 64, 500, hitp);
    
    /*vec3 o = point;
    mat4 mvp = p*projection_data.mv;
    bool hit = false;
    vec2 hitp;
    for(int i = 0; i < 64; i++)
    {
        o += reflected*0.08;
        vec4 o_screen = mvp * vec4(o, 1.0);
        o_screen.xyz /= o_screen.w;
        float s_depth = texture(in_depth, o_screen.xy / 2 + .5).x;
        if(o_screen.z > s_depth + 0.01)
        {
            hit = true;
            hitp = res * (o_screen.xy / 2 + .5);
            break;
        }
    }*/

    vec2 cpos = hitp/res;
    vec2 velocity_vector = texture(velocity_buffer, cpos).xy * 2 - 1;
    vec2 puv = cpos - velocity_vector;
    puv = clamp(puv, 0, 1);
    vec3 reflected_color;
    if(hit && cpos.x > 0 && cpos.x < 1.0 && cpos.y > 0 && cpos.x < 1.0)
    {
        reflected_color = texture(hdr_color, puv).xyz *
            fresnel(normal, -vdir);
    }   
    else 
    {
        reflected_color = vec3(0);
    }

    imageStore(half_res_specular,  ivec2(gl_GlobalInvocationID.xy), reflected_color.xyzz);
}