#include "Algorithm/PathTracing.h"
#include "Algorithm/RayIntersection.h"
#include "Algorithm/BSDF.h"
#include "Lighting/LightEvaluation.h"

// TODO: Implement a faster methods like BVH, KD-Tree or uniform grid acceleration
// TODO: Split the diffuse and specular into different Monte Carlo, so we can decide the sample count for each one
vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth)
{
    const triangle_collection_t* triangles = &scene->triangles;
    const bvh_tree_t* bvh_tree = &scene->bvh_tree;
    const material_collection_t* materials = &scene->materials;
    const light_collection_t* lights = &scene->lights;

    vec4s accumulated_color = glms_vec4_zero();
    vec3s throughput = glms_vec3_one();

    ray_t active_ray = ray;
    vec3s prev_normal = glms_vec3_zero();
    float pdf_bsdf = 1.0f;

    uint16_t depth = 0;
    while (depth < max_depth)
    {
        hit_result_t closest_hit = ray_intersect_scene(active_ray, scene);

        if (!closest_hit.hit)
        {
            vec3s sky_light = evaluate_bsdf_sky(scene, NULL, throughput, sample_index);
            if (depth > 0)
            {
                // Have to multiply the weight since we evaluate the sky at each bounce
                float pdf_nee = pdf_cosine_weighted_hemisphere(prev_normal, active_ray.direction);
                float weight = power_heuristic(pdf_bsdf, pdf_nee);
                sky_light = glms_vec3_scale(sky_light, weight);
            }
            accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(sky_light, 1.0f)); // TODO: Physical Skybox
            break;
        }

        // Add the emission of the hit material to the accumulated color
        material_t* hit_material = &materials->buffer[triangles->buffer[closest_hit.triangle_id].material_id];
        vec3s emission = hit_material->emission;
        accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(glms_vec3_mul(throughput, emission), 1.0f));

        light_shading_context_t light_context = {
            .normal = closest_hit.normal,
            .hit_point = closest_hit.point,
            .wo = active_ray.direction,

            .bounce_depth = depth,

            .bvh_tree = bvh_tree,
            .material = hit_material
        };

        // Running the light loop.
        // TODO: Implementing other light types.
        for (uint32_t i = 0; i < lights->directional_light_count; i++)
        {
            vec3s l = evaluate_bsdf_directional(lights->directional_lights[i], &light_context, throughput, sample_index);
            accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(l, 1.0f));
        }

        vec3s sky_light = evaluate_bsdf_sky(scene, &light_context, throughput, sample_index);
        accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(sky_light, 1.0f));

        // Bounce and prepare for the next iteration
        vec3s wo = glms_vec3_negate(active_ray.direction); // We need to negate the direction of the incoming ray
        vec3s wi = sample_material_bsdf(hit_material, closest_hit.normal, wo, sample_index, depth, &pdf_bsdf);
        if (pdf_bsdf <= 0.0f)
        {
            break;
        }

        shading_context_t shading_context = {
            .normal = closest_hit.normal,
            .wi = wi,
            .wo = wo
        };
        vec3s bsdf = evaluate_material_bsdf(hit_material, &shading_context);
        float cos_theta = fmaxf(0.0f, glms_vec3_dot(wi, closest_hit.normal));

        throughput = glms_vec3_mul(throughput, glms_vec3_scale(bsdf, cos_theta / pdf_bsdf));

         // We do Russian roulette to decide whether to continue tracing or terminate the path
        if (depth > 2)
        {
            float q = fminf(glms_vec3_max(throughput), 0.95f);
            float rr_sample = sobol_sample(sample_index, sobol_get_dimension(depth, PRNG_TERMINATE));
            // float rr_sample = random_float();
            if (rr_sample > q)
            {
                 break; // Terminate the path
            }
            // Keep the energy of the path by scaling the throughput
            throughput = glms_vec3_scale(throughput, 1.0f / q);
        }

        active_ray.origin = glms_vec3_add(closest_hit.point, glms_vec3_scale(closest_hit.normal, 1.192092896e-04F));
        active_ray.direction = wi;
        prev_normal = closest_hit.normal;

        depth++;
    }

    return accumulated_color;
}