#include "Lighting/SkyLight.h"
#include "Common.h"

// Constant Sky

sky_light_t sky_create_constant_sky(const constant_sky_data_t* data)
{
    sky_light_t light = {
        .data_size = sizeof(constant_sky_data_t),
        .evaluate_bsdf_sky = evaluate_bsdf_const_sky,
    };

    light.data = malloc(sizeof(constant_sky_data_t));
    if (light.data == NULL)
    {
        return light;
    }

    memcpy(light.data, data, sizeof(constant_sky_data_t));
    return light;
}

path_output evaluate_bsdf_const_sky(const void* data, const light_shading_context_t* context, vec3s throughput, uint32_t sample_index)
{
    const constant_sky_data_t* sky_data = (const constant_sky_data_t*)data;

    path_output output = {0.0f};
    output.state = PS_TERMINATE;

    if (sky_data->intensity <= 0.0f)
    {
        return output;
    }

    vec3s sky_light = glms_vec3_scale(sky_data->color, sky_data->intensity);
    float pdf = 0.25f * INV_PI;
    output.pdf = pdf;

    if (context->bvh_tree == NULL)
    {
        output.direct_lighting = glms_vec3_mul(sky_light, throughput);
        return output;
    }

    uint16_t d1 = sobol_get_dimension(context->bounce_depth, PRNG_LIGHT_U);
    uint16_t d2 = sobol_get_dimension(context->bounce_depth, PRNG_LIGHT_V);

    vec3s wi = random_uniform_cdf_direction(context->normal, sample_index, d1, d2);

    ray_t shadow_ray = ray_create(offset_ray_origin(context->position, context->normal, context->wo), wi);

    hit_result_t shadow_hit = {0};
    ray_intersect_bvh_any(&shadow_ray, context->bvh_tree->nodes, context->bvh_tree->primitive_indices, context->bvh_tree->triangles, 0, &shadow_hit);
    if (shadow_hit.hit)
    {
        return output;
    }

    float cos_theta = fmaxf(glms_vec3_dot(wi, context->normal), 0.0f);
    output.direct_lighting = glms_vec3_scale(glms_vec3_mul(sky_light, throughput), cos_theta / pdf);

    output.wi = wi;
    output.state = PS_NORMAL;
    return output;
}


static uint32_t lower_bound_float(const float* arr, uint32_t n, float target)
{
    uint32_t low = 0;
    uint32_t high = n; // Search range [low, high)
    while (low < high)
    {
        uint32_t mid = low + (high - low) / 2;
        if (arr[mid] < target)
        {
            low = mid + 1;
        }
        else
        {
            high = mid;
        }
    }

    return low; // This is the index where target would be inserted to maintain order
}

// HDR Sky

sky_light_t sky_create_hdr_sky(const texture_collection_t* textures, texture_entity_t hdri_entity, float intensity)
{
    sky_light_t light = {
        .data_size = sizeof(hdr_sky_data_t),
        .evaluate_bsdf_sky = evaluate_bsdf_hdr_sky,
        .free_sky_data = (sky_free_f)hdr_sky_free,
    };

    const texture_t* hdri = get_texture(textures, hdri_entity);
    if (hdri == NULL)
    {
        return light;
    }

    hdr_sky_data_t data = {
        .width = hdri->width,
        .height = hdri->height,
        .texture = hdri_entity,
        .intensity = intensity,
    };

    size_t size_hw = (size_t)data.height * data.width; // height * width elements
    size_t size_w = (size_t)data.width;                // width elements
    size_t size_wh = (size_t)data.width * data.height; // width * height elements (same as hw)

    // Total size for intermediate buffers: p(x,y)/cond_cdf + p(x) + cdf_x + cond_transposed_cdf
    size_t intermediate_buffer_size = size_hw + size_w + size_w + size_hw + size_wh;

    float* intermediate_buffer = (float*)malloc(sizeof(float) * intermediate_buffer_size);
    if (intermediate_buffer == NULL)
    {
        free(light.data);
        return light;
    }

    vec3s* cache = (vec3s*)malloc(sizeof(vec3s) * size_hw);
    if (cache == NULL)
    {
        free(light.data);
        free(intermediate_buffer);
        return light;
    }

    // These pointers partition the single allocated intermediate_buffer
    float* p_xy_original_pdf = intermediate_buffer;               // size hw (Stores the normalized p(x,y) values)
    float* p_x_buffer = p_xy_original_pdf + size_hw;              // size w   (Stores p(x))
    float* cdf_x_buffer = p_x_buffer + size_w;                    // size w   (Stores cdf_x_margin)
    float* cdf_y_conditional = cdf_x_buffer + size_w;             // size hw  (Stores cdf(y|x))
    float* cdf_y_transposed_buffer = cdf_y_conditional + size_hw; // size wh (Stores transposed cdf(y|x))

    // --- Calculate Initial Luminance (PDF) ---
    float lumSum = 0.0f;
    for (uint32_t i = 0; i < data.height; ++i)
    {
        // Note: If equiareal mapping requires an area element factor like sin(theta),
        // calculate it here and multiply the luminance before summing and normalizing.
        // float v = (i + 0.5f) / (float)data.height;
        // float sin_theta = sinf(v * PI);

        for (uint32_t j = 0; j < data.width; ++j)
        {
            vec2s uv = {(float)j / (float)data.width, (float)i / (float)data.height};
            vec4s pixel = texture_get_pixel(hdri, uv, 0);
            float lum = luminance(glms_vec3(pixel)); // * sin_theta;

            p_xy_original_pdf[i * data.width + j] = lum;
            lumSum += lum; // Sum of luminances (or lum * area_element)
        }
    }

    float inv_lumSum = (lumSum > FLT_EPSILON) ? (1.0f / lumSum) : 0.0f;
    for (size_t i = 0; i < size_hw; ++i)
    {
        p_xy_original_pdf[i] *= inv_lumSum;
    }

    // --- Calculate Marginal PDF p(x) ---
    for (uint32_t j = 0; j < data.width; ++j)
    {
        p_x_buffer[j] = 0.0f;
        for (uint32_t i = 0; i < data.height; ++i)
        {
            p_x_buffer[j] += p_xy_original_pdf[i * data.width + j];
        }
    }

    // --- Calculate Marginal CDF (prefix sum on marginal PDF) ---
    memcpy(cdf_x_buffer, p_x_buffer, sizeof(float) * size_w);
    if (data.width > 0)
    {
        for (uint32_t i = 1; i < data.width; ++i)
        {
            cdf_x_buffer[i] += cdf_x_buffer[i - 1]; // Prefix sum on cdf_x_buffer
        }
        // The last element cdf_x_buffer[width-1] should be close to 1.0 (total probability)
    }

    // --- Calculate Conditional PDF and CDF ---
    // p_xy_buffer currently holds p(x,y)
    // It will be modified in place to hold cdf(y|x)
    memcpy(cdf_y_conditional, p_xy_original_pdf, sizeof(float) * size_hw);
    for (uint32_t j = 0; j < data.width; ++j)
    {
        float marginal_x_val = p_x_buffer[j]; // Get p(x=j) from the stored p_x_buffer
        float inv_marginal_x_val = (marginal_x_val > 1e-6f) ? (1.0f / marginal_x_val) : 0.0f;

        // Normalize p(x,y) column to get conditional PDF p(y|x) and calculate conditional CDF cdf(y|x)
        if (data.height > 0)
        {
            // First element is p(y=0|x) = p(x,y=0) / p(x), then becomes cdf(y=0|x)
            cdf_y_conditional[0 * data.width + j] *= inv_marginal_x_val;
            // Calculate prefix sum (conditional CDF) in place
            for (uint32_t i = 1; i < data.height; ++i)
            {
                // p(y=i|x) = p(x,y=i) / p(x)
                cdf_y_conditional[i * data.width + j] *= inv_marginal_x_val;
                // cdf(y=i|x) = cdf(y=i-1|x) + p(y=i|x)
                cdf_y_conditional[i * data.width + j] += cdf_y_conditional[(i - 1) * data.width + j];
            }
            // The last element of this column (at row height-1) should be close to 1.0
        }
    }

    // --- Transpose Conditional CDF ---
    // Transpose cdf_y_conditional (size hw) into cdf_y_transposed_buffer (size wh)
    for (uint32_t i = 0; i < data.height; ++i)
    {
        for (uint32_t j = 0; j < data.width; ++j)
        {
            cdf_y_transposed_buffer[j * data.height + i] = cdf_y_conditional[i * data.width + j];
        }
    }

    // --- Precompute Samples using Inverse Transform Sampling ---
    // Populate the 'cache' array directly, which is now vec3s*
    // cache[i * width + j] represents the output pixel for input random numbers (i/height, j/width)
    for (uint32_t i = 0; i < data.height; ++i)
    { // Corresponds to xi_1 (vertical random number)
        for (uint32_t j = 0; j < data.width; ++j)
        {
            // Input random numbers [0, 1)
            float xi_1 = i / (float)data.height;
            float xi_2 = j / (float)data.width;

            // Sample x index using xi_1 and marginal CDF (cdf_x_buffer)
            uint32_t x_sample_idx = lower_bound_float(cdf_x_buffer, data.width, xi_1);
            x_sample_idx = (x_sample_idx == data.width) ? data.width - 1 : x_sample_idx;

            // Sample y index using xi_2 and conditional CDF for x=x_sample_idx (cdf_y_transposed_buffer)
            // The conditional CDF for x_sample_idx is a 1D array starting at cdf_y_transposed_buffer[x_sample_idx * height]
            uint32_t y_sample_idx = lower_bound_float(cdf_y_transposed_buffer + x_sample_idx * data.height, data.height, xi_2);
            y_sample_idx = (y_sample_idx == data.height) ? data.height - 1 : y_sample_idx;

            // Get the normalized PDF value at the sampled texture coordinates (x_sample_idx, y_sample_idx)
            // This is stored in p_xy_original_pdf (which holds the normalized p(x,y))
            float pdf_at_sample = p_xy_original_pdf[y_sample_idx * data.width + x_sample_idx];

            // Store the sampled texture coordinates (normalized [0, 1)) and PDF in the cache (vec3s)
            int cache_idx = i * data.width + j;
            cache[cache_idx].x = (float)x_sample_idx / (float)data.width;  // Sampled U
            cache[cache_idx].y = (float)y_sample_idx / (float)data.height; // Sampled V
            cache[cache_idx].z = pdf_at_sample;                            // Normalized PDF at (U, V)
        }
    }

    free(intermediate_buffer);

    light.data = malloc(sizeof(hdr_sky_data_t));
    if (light.data == NULL)
    {
        free(cache);
        return light;
    }

    data.total_weight = lumSum;
    data.cdf_cache = cache;

    memcpy(light.data, &data, sizeof(hdr_sky_data_t));
    return light;
}

static inline float hdr_pdf(const vec3s* cdf, vec3s v, uint32_t height, uint32_t width)
{
    vec2s uv = direction_to_equirectangular(v);
    uint32_t x, y;
    uv_to_index(uv, width, height, &x, &y);
    float pdf = cdf[y * width + x].z;

    float theta = uv.y * PI;
    float sin_theta = fmaxf(sinf(theta), FLT_EPSILON);

    float convert = (float)(height * width) / (2.0f * PI_TWO * sin_theta);

    return pdf * convert;
}

path_output evaluate_bsdf_hdr_sky(const void* data, const light_shading_context_t* context, vec3s throughput, uint32_t sample_index)
{
    const hdr_sky_data_t* sky_data = (const hdr_sky_data_t*)data;

    path_output output = {0.0f};
    output.state = PS_TERMINATE;

    if (sky_data->intensity <= 0.0f)
    {
        return output;
    }

    if (context->bvh_tree == NULL)
    {
        vec2s uv = direction_to_equirectangular(context->wo);
        vec4s sky_light = texture_sample_lod(get_texture(context->textures, sky_data->texture), uv, 0);
        output.direct_lighting = glms_vec3_scale(glms_vec3_mul(glms_vec3(sky_light), throughput), sky_data->intensity);
        output.pdf = 1.0f / (4.0f * PI);
        return output;
    }

    uint32_t i,j;
    // Should we also use sobol numbers for the sky sampling?
    uint16_t d1 = sobol_get_dimension(context->bounce_depth, PRNG_LIGHT_U);
    uint16_t d2 = sobol_get_dimension(context->bounce_depth, PRNG_LIGHT_V);
    float r1 = sobol_sample(sample_index, d1);
    float r2 = sobol_sample(sample_index, d2);
    //float r1 = random_float();
    //float r2 = random_float();

    uv_to_index((vec2s){r1, r2}, sky_data->width, sky_data->height, &j, &i);
    vec3s cdf = sky_data->cdf_cache[i * sky_data->width + j];
    vec3s wi = equirectangular_to_direction(cdf.x,-cdf.y);

    float n_dot_l = glms_vec3_dot(wi, context->normal);
    if (n_dot_l <= 0.0f)
    {
        return output;
    }

    ray_t shadow_ray = ray_create(offset_ray_origin(context->position, context->normal, context->wo), wi);
    hit_result_t shadow_hit = {0};
    ray_intersect_bvh_any(&shadow_ray, context->bvh_tree->nodes, context->bvh_tree->primitive_indices, context->bvh_tree->triangles, 0, &shadow_hit);
    if (shadow_hit.hit)
    {
        return output;
    }

    vec2s uv = direction_to_equirectangular(wi);

    const texture_t* hdri = get_texture(context->textures, sky_data->texture);
    vec4s pixel = texture_sample_lod(hdri, uv, 0);
    vec3s sky_light = glms_vec3_scale(glms_vec3(pixel), sky_data->intensity);

    float pdf = hdr_pdf(sky_data->cdf_cache, wi, sky_data->height, sky_data->width);
    output.direct_lighting = glms_vec3_scale(glms_vec3_mul(sky_light, throughput), n_dot_l / pdf);

    output.wi = wi;
    output.pdf = pdf;
    output.state = PS_NORMAL;
    return output;
}

void hdr_sky_free(hdr_sky_data_t* data)
{
    free(data->cdf_cache);

    data->cdf_cache = NULL;
}