#define FLIP_Y

#include "Rendering/Scene.h"
#include "Algorithm/PathTracing.h"

scene_t scene_create(uint64_t triangle_count, uint8_t material_count)
{
    scene_t scene = {0};

    scene.triangles = triangle_collection_create(triangle_count);
    scene.materials = material_collection_create(material_count);
    scene.camera = camera_create(
        (vec3s){0.0f, 0.0f, 5.0f},
        (vec3s){0.0f, 0.0f, -1.0f},
        (vec3s){0.0f, 1.0f, 0.0f},
        0.025f,
        0.036f,
        1.777777f
    );

    return scene;
}

static inline vec2s compute_ndc(uint32_t x, uint32_t y, uint32_t width, uint32_t height)
{
    return (vec2s){
        .x = (float)x / (float)width,
#ifdef FLIP_Y
        .y = 1.0f - (float)y / (float)height
#else
        .y = (float)y / (float)height
#endif
    };
}

render_target_t scene_render(scene_t* scene, rendering_config_t config)
{
    if (fabsf((float)config.width / config.height - scene->camera.aspect_ratio) > FLT_EPSILON)
    {
        scene->camera = camera_create(
            scene->camera.position,
            scene->camera.forward,
            scene->camera.up,
            scene->camera.focal_length,
            scene->camera.size_x,
            (float)config.width / (float)config.height
        );
    }

    // The actual float buffer inside the render target is on the heap, copy return shoudl be fine.
    render_target_t img = render_target_create(config.width, config.height);

    uint32_t tile_count_x = (config.width + config.tile_size - 1) / config.tile_size;
    uint32_t tile_count_y = (config.height + config.tile_size - 1) / config.tile_size;
    uint32_t tile_count = tile_count_x * tile_count_y;

    float brightness_per_sample = ((float)M_PI * 2.0f) * (1.0f / config.sample_count);

    int64_t x, y, tile_index; // OpenMP requires these to be declared outside the parallel region. Also, they need to be signed integers for OpenMP to work correctly with the loop bounds. To avoid overflow, we need to use int64_t
#ifndef DEBUG
    #pragma omp parallel for schedule(dynamic, 1) default(none) \
                             shared(tile_count_x, tile_count_y, tile_count) \
                             private(x, y, tile_index)
#endif
    for (tile_index = 0; tile_index < tile_count; tile_index++)
    {
        uint32_t tile_x_0 = tile_index % tile_count_x * config.tile_size;
        uint32_t tile_y_0 = tile_index / tile_count_x * config.tile_size;
        uint32_t tile_x_1 = (uint32_t)fmin(tile_x_0 + config.tile_size, config.width);
        uint32_t tile_y_1 = (uint32_t)fmin(tile_y_0 + config.tile_size, config.height);

        vec3s coord = glms_vec3_add(scene->camera.position, glms_vec3_scale(scene->camera.forward, scene->camera.focal_length));

        for (y = tile_y_0; y < tile_y_1; y++)
        {
            for (x = tile_x_0; x < tile_x_1; x++)
            {
                for (uint16_t k = 0; k < config.sample_count; k++)
                {
                    vec2s position_ndc = compute_ndc(x, y, config.width, config.height);
                    float screen_x = position_ndc.x * 2.0f - 1.0f;
                    float screen_y = position_ndc.y * 2.0f - 1.0f;
                    float sensor_offset_x = screen_x * scene->camera.size_x / 2.0f;
                    float sensor_offset_y = screen_y * scene->camera.size_y / 2.0f;

                    vec3s image_plane_point = coord;
                    image_plane_point = glms_vec3_add(image_plane_point, glms_vec3_scale(scene->camera.right, sensor_offset_x));
                    image_plane_point = glms_vec3_add(image_plane_point, glms_vec3_scale(scene->camera.up, sensor_offset_y));

                    ray_t ray = {
                        .origin = scene->camera.position,
                        .direction = glms_vec3_normalize(glms_vec3_sub(image_plane_point, scene->camera.position))
                    };

                    vec3s out_color = path_trace(&scene->triangles, &scene->materials, ray, 4);
                    out_color = glms_vec3_scale(out_color, brightness_per_sample);
                    if (glms_vec3_eq(out_color, 0.0f))
                    {
                        continue;
                    }

                    //TODO: Handle alpha
                    vec4s color = {
                        .x = out_color.x,
                        .y = out_color.y,
                        .z = out_color.z,
                        .w = 1.0f
                    };
                    vec4s old_color = render_target_get_pixel(&img, x, y);

                    render_target_set_pixel(&img, x, y, glms_vec4_add(old_color, color));
                }
            }
        }
    }

    return img;
}

void scene_free(scene_t* scene)
{
    triangle_collection_free(&scene->triangles);
    material_collection_free(&scene->materials);
}