#include <defocus/camera.h>

#include <string.h>
#include <stdlib.h>
#include <assert.h>

#ifdef _MSC_VER
#include <malloc.h>
#endif

/* ********************************************************
 *
 * Ray Packet Functions
 *
 * ********************************************************/

DF_API void df_release_ray_packet(df_ray_packet *rays)
{
#ifdef _MSC_VER
    _aligned_free(rays->simd_mem);
#else
    free(rays->simd_mem);
#endif
    free(rays->ray_uvs);

    memset(rays, 0, sizeof(*rays));
}

/* ********************************************************
 *
 * Perspective Camrea
 *
 * ********************************************************/

typedef struct
{
    float focal_dist;
    float lens_radius;
    
    /* Raster space, i.e. the space for which we generate rays */
    float raster_width;
    float raster_height;

    /* Image(=screen) space, i.e. the space in which we store pixels */
    float image_width;
    float image_height;

    df_m4 screen_to_raster;
    df_m4 raster_to_screen;
    df_m4 raster_to_camera;
} df_perspective_camera;

static void pc_release(void *o) { df_perspective_camera *camera = o; }

static df_ray_packet pc_build_ray_packet(void *o)
{
    df_perspective_camera *camera = o;
    df_ray_packet packet;
    memset(&packet, 0, sizeof(packet));

    /* Generate 1 ray per pixel (ignore lens for now) */
    size_t count = (size_t)camera->raster_width * (size_t)camera->raster_height;

    size_t alloc_count = count;
    /* Round up to nearest multiple of 4, for SSE.
     * Also satisfies 16 byte alignment (assuming simd_mem is 16 byte aligned)
     * because 4 * sizeof(float) = 16.
     */
    if ((alloc_count % 4) != 0) {
        alloc_count = ((alloc_count + 3) / 4) * 4;
    }

#ifdef _MSC_VER
    packet.simd_mem = _aligned_malloc(sizeof(float) * alloc_count * 6, 16);
#elif defined(_ISOC11_SOURCE) /* Feature test macro for GCC */
    packet.simd_mem = aligned_alloc(16, sizeof(float) * alloc_count * 6);
#else
    /* Fall back to regular malloc and hope for the best */
    packet.simd_mem = malloc(sizeof(float) * alloc_count * 6);
#endif

    packet.base_x = packet.simd_mem;
    packet.base_y = packet.base_x + alloc_count;
    packet.base_z = packet.base_y + alloc_count;
    packet.dir_x = packet.base_z + alloc_count;
    packet.dir_y = packet.dir_x + alloc_count;
    packet.dir_z = packet.dir_y + alloc_count;

    packet.ray_uvs = malloc(sizeof(df_v2) * count);

    packet.ray_count = count;

    size_t i = 0;
    for (float y = 0; y < camera->raster_height; y += 1.f) {
        for (float x = 0; x < camera->raster_width; x += 1.f) {
            packet.base_x[i] = 0.f;
            packet.base_y[i] = 0.f;
            packet.base_z[i] = 0.f;

            df_v3 raster_p = {x, y, 0.f};
            df_v3 camera_p = df_transform_v3(camera->raster_to_camera, raster_p);
            df_v3 dir = df_normalize_v3(camera_p);

            packet.dir_x[i] = dir.x;
            packet.dir_y[i] = dir.y;
            packet.dir_z[i] = dir.z;

            df_v3 img_p = df_transform_v3(camera->raster_to_screen, raster_p);
            packet.ray_uvs[i].u = img_p.x / camera->image_width;
            packet.ray_uvs[i].v = img_p.y / camera->image_height;

            ++i;
            assert(i <= count);
        }
    }

    return packet;
}

DF_API df_camera_i df_create_perspective_camera(float image_width,
                                                float image_height,
                                                float raster_width,
                                                float raster_height,
                                                float far_dist,
                                                float near_dist,
                                                float focal_dist,
                                                float lens_radius)
{
    df_perspective_camera *camera = malloc(sizeof(*camera));

    camera->screen_to_raster = df_scale(raster_width, raster_height, 1.f);
    camera->screen_to_raster =
        df_mul_m4(camera->screen_to_raster, df_scale(1.f / image_width, -1.f / image_height, 1.f));
    camera->screen_to_raster = df_mul_m4(camera->screen_to_raster, df_translate(0.f, -image_height, 0.f));
    camera->raster_to_screen = df_inverse_transform(camera->screen_to_raster);

    /* Perspective projection matrix */
    df_m4 persp = {0.f};
    DF_M4_AT(persp, 0, 0) = 1.f;
    DF_M4_AT(persp, 1, 1) = 1.f;
    DF_M4_AT(persp, 2, 2) = far_dist / (far_dist - near_dist);
    DF_M4_AT(persp, 2, 3) = -1.f * far_dist * near_dist / (far_dist - near_dist);
    DF_M4_AT(persp, 3, 2) = 1.f;

    camera->raster_to_camera = df_mul_m4(df_inverse(persp), camera->raster_to_screen);

    camera->focal_dist = focal_dist;
    camera->lens_radius = lens_radius;
    camera->raster_width = raster_width;
    camera->raster_height = raster_height;
    camera->image_width = image_width;
    camera->image_height = image_height;

    df_camera_i iface = {.release = pc_release, .build_ray_packet = pc_build_ray_packet, .o = camera};
    return iface;
}