// ----------------------------------------------------------------------
// Most of the BRDF stuff in this code comes from github.com/boksajak/referencePT
// This repo is used it Ray Tracing gems 2, reference path tracer chaper
// ----------------------------------------------------------------------

#extension GL_EXT_buffer_reference : require
#extension GL_EXT_nonuniform_qualifier : require

#include "math.h"
#include "color_utils.h"
#include "rand.h"
#include "brdf.h"

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

struct tex {
    uint texture_id;
    uint sampler_id;
};

struct material {
    vec4 base_color;
    // bindless indices
    tex base_color_texture;
    tex metallic_roughness_texture;
    float metalness;
    float roughness;
    float[2] padding;
};

struct sphere {
    vec4 position;
    material mat;
    float radius;
    float[3] padding;
};

struct triangle {
    vec3 positions[3];
    vec3 normals[3];
    vec2 uvs[3];
    material mat;
};

struct camera {
    vec4 position;

    vec4 forward;
    vec4 up;
    vec4 right;

    vec4 horizontal;
    vec4 vertical;
    vec4 first_pixel;

    float lens_radius;
    float[3] padding;
};

struct bvh_node {
    vec3 min;
    int next_id;
    vec3 max;
    int primitive_id;
};

layout(buffer_reference) readonly buffer scene_metadata {
    camera cam;

    uint max_bounce;
    uint min_bounce;
    uint width;
    uint height;
    uint sample_index;

    uint enable_dof;
    uint debug_bvh;
    int downscale_factor;
};

layout(buffer_reference) readonly buffer indices_array {
    uint indices[];
};

layout(buffer_reference) readonly buffer positions_array {
    float positions[];
};

layout(buffer_reference) readonly buffer normals_array {
    float normals[];
};

layout(buffer_reference) readonly buffer uvs_array {
    float uvs[];
};

layout(buffer_reference) readonly buffer materials_array {
    material materials[];
};

layout(buffer_reference) readonly buffer nodes_array {
    bvh_node[] nodes;
};

layout(set = 0, binding = 0) uniform sampler samplers[];
layout(set = 0, binding = 1) uniform texture2D textures[];
layout(set = 0, binding = 2, rgba32f) uniform image2D images[];

layout(push_constant) uniform buffers {
    scene_metadata scene;
    nodes_array bvh;
    indices_array indices_arr;
    positions_array positions_arr;
    normals_array normals_arr;
    uvs_array uvs_arr;
    materials_array materials_arr;
    uint output_image_index;
    uint accumulation_image_index;
    uint padding[2];
} bufs;

#include "ray.h"

vec3 ray_color(ray r, uint seed) {
    hit_info info;
    vec3 radiance = vec3(0.0);
    vec3 throughput = vec3(1.0);

    for (uint bounce = 0; bounce < bufs.scene.max_bounce; bounce++) {
        if (!hit_node(r, info))
            return throughput * vec3(1.0);

        vec3 v = -r.direction;



        triangle tri = get_triangle(info.primitive_id);

        vec2 uv = interpolate_attribute(tri.uvs, info.barycentrics);
        vec3 shading_normal = normalize(interpolate_attribute(tri.normals, info.barycentrics));

        vec3 diffuse_color = srgb_to_linear(texture(
            sampler2D(
                textures[nonuniformEXT(tri.mat.base_color_texture.texture_id)],
                samplers[nonuniformEXT(tri.mat.base_color_texture.sampler_id)]
            ), uv
        ).xyz) * tri.mat.base_color.xyz;
        vec2 metalness_roughness = texture(
            sampler2D(
                textures[nonuniformEXT(tri.mat.metallic_roughness_texture.texture_id)],
                samplers[nonuniformEXT(tri.mat.metallic_roughness_texture.sampler_id)]
            ), uv
        ).xy * vec2(tri.mat.metalness, tri.mat.roughness);




        if (dot(info.geometry_normal, v) < 0.f) info.geometry_normal = -info.geometry_normal;
        if (dot(info.geometry_normal, shading_normal) < 0.0f) shading_normal = -shading_normal;

        uint brdf_type;
        if (metalness_roughness.x == 1.f && metalness_roughness.y == 0.f) {
            brdf_type = SPECULAR;
        } else {
            float brdf_probability = get_brdf_probability(diffuse_color, metalness_roughness.x, -r.direction, shading_normal);
            if (rand(seed) < brdf_probability) {
                brdf_type = SPECULAR;
                throughput /= brdf_probability;
            } else {
                brdf_type = DIFFUSE;
                throughput /= (1.0 - brdf_probability);
            }
        }

        // Ignore ray coming from below the hemisphere
        if (dot(shading_normal, v) <= 0.f) break;

        // Move to tangent space
        vec4 rotation_to_z = rotation_to_z_axis(shading_normal);
        vec3 view_local = rotate_point(rotation_to_z, -r.direction);
        vec3 normal_local = vec3(0.0, 0.0, 1.0);
        vec3 ray_dir_local;
        vec3 sample_weight;

        if (brdf_type == DIFFUSE) { // Lambertian diffuse
            vec2 r = vec2(rand(seed), rand(seed));
            ray_dir_local = sample_hemisphere(r);
            sample_weight = base_color_to_diffuse_reflectance(diffuse_color, metalness_roughness.x);
        } else if (brdf_type == SPECULAR) {
            float alpha = metalness_roughness.y * metalness_roughness.y;
            float alpha_squared = alpha * alpha;

            vec3 half_local;
            if (alpha == 0.0f) {
                    // Fast path for zero roughness (perfect reflection), also prevents NaNs appearing due to divisions by zeroes
                    half_local = vec3(0.0f, 0.0f, 1.0f);
            } else {
                    // For non-zero roughness, this calls VNDF sampling for GG-X distribution or Walter's sampling for Beckmann distribution
                    half_local = sample_GGXVNDF(view_local, vec2(alpha), seed);
            }

            // Reflect view direction to obtain light vector
            vec3 light_local = reflect(-view_local, half_local);

            vec3 specular_F0 = base_color_to_diffuse_reflectance(diffuse_color, metalness_roughness.x);

            // Note: HdotL is same as HdotV here
            // Clamp dot products here to small value to prevent numerical instability.
            float HdotL = clamp(dot(half_local, light_local), 0.00001f, 1.0f);
            float NdotL = clamp(dot(normal_local, light_local), 0.00001f, 1.0f);
            float NdotV = clamp(dot(normal_local, view_local), 0.00001f, 1.0f);
            float NdotH = clamp(dot(normal_local, half_local), 0.00001f, 1.0f);
            vec3 f = eval_fresnel(specular_F0, shadowed_F90(specular_F0), HdotL);

            sample_weight = f * specular_sample_weight_GGXVNDF(alpha, alpha_squared, NdotL, NdotV, HdotL, NdotH);

            ray_dir_local = light_local;
        }

        // Prevent tracing direction with no contribution
        if (luminance(sample_weight) == 0.0f) break;

        // Move to global space
        vec4 rotation_from_z = invert_rotation(rotation_to_z);
        vec3 ray_dir = rotate_point(rotation_from_z, ray_dir_local);

        // Prevent tracing direction "under" the hemisphere (behind the triangle)
        if (dot(info.geometry_normal, ray_dir) <= 0.0f) break;

        throughput *= sample_weight;

        r = ray(info.point, ray_dir, 0.001f, 1e15);

        // Russian Roulette
        // As the throughput gets smaller, the ray is more likely to get terminated early.
        // Survivors have their value boosted to make up for fewer samples being in the average.
        if (bounce > bufs.scene.min_bounce) {
            float probability = min(0.95f, luminance(throughput));
            if (probability < rand(seed)) break;
            throughput /= probability;
        }
    }

    return radiance;
}

ray generate_camera_ray(ivec2 coords, uint seed) {
    vec2 scene_size = vec2(bufs.scene.width - 1, bufs.scene.height - 1);
    vec2 uv = vec2(coords.x / scene_size.x, 1.0 - coords.y / scene_size.y);
    vec2 rand_disk = disk_vec(vec2(rand(seed), rand(seed)));
    vec2 jittered_uvs = uv + vec2(rand(seed), rand(seed)) / scene_size;
    vec2 lens_disk = vec2(0.0);

    if (bufs.scene.enable_dof == 1) {
        lens_disk = rand_disk * bufs.scene.cam.lens_radius;
    }

    vec3 offset = lens_disk.x * bufs.scene.cam.right.xyz + lens_disk.y * bufs.scene.cam.up.xyz;
    vec3 proj_plane_pos = bufs.scene.cam.first_pixel.xyz + jittered_uvs.x * bufs.scene.cam.horizontal.xyz + jittered_uvs.y * bufs.scene.cam.vertical.xyz;

    return ray(bufs.scene.cam.position.xyz + offset, proj_plane_pos.xyz - bufs.scene.cam.position.xyz - offset, 0.001f, 1e15);
}

void main() {
    if (any(greaterThanEqual(gl_GlobalInvocationID.xy, uvec2(bufs.scene.width, bufs.scene.height))))
        return;

    ivec2 coords = ivec2(gl_GlobalInvocationID.xy) * bufs.scene.downscale_factor;

    uint seed = uint(coords.x * uint(1973) + coords.y * uint(9277) + bufs.scene.sample_index * uint(26699)) | uint(1);
    ray r = generate_camera_ray(coords, seed);

    vec3 out_color = vec3(0.0);
    if (bufs.scene.debug_bvh == 1) {
        out_color = hit_aabbs(r) * vec3(0.001, 0.0, 0.0);
    } else {
        out_color = ray_color(r, seed);
    }

    float spp_scale = (1.0 / bufs.scene.sample_index);
    vec3 acc_color = srgb_to_linear(imageLoad(images[nonuniformEXT(bufs.accumulation_image_index)], coords).xyz);
    out_color = linear_to_srgb(mix(acc_color, out_color, spp_scale));

    for (uint i = 0; i < bufs.scene.downscale_factor; i++) {
        for (uint j = 0; j < bufs.scene.downscale_factor; j++) {
            imageStore(images[nonuniformEXT(bufs.output_image_index)], coords + ivec2(i, j), vec4(out_color, 1.0));
        }
    }
}