#include "Lighting/SkyLight.h"
#include "Common.h"
#include <stdio.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;
    }

    uint32_t scramble = hash_position(context->position);
    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, scramble);

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

    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_SUCCESS;
    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_handle_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,
        .pdf_uv_mass = NULL,
        .cdf_x = NULL,
        .cdf_y_transposed = NULL,
    };

    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)
    {
        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)
    {
        // Calculate the sine of the polar angle theta
        // Map row i to V [0,1], then to theta [0, PI]
        float v = (i + 0.5f) / (float)data.height;
        float theta = v * PI;
        float sin_theta = sinf(theta);

        for (uint32_t j = 0; j < data.width; ++j)
        {
            vec2s uv = {((float)j + 0.5f) / (float)data.width, ((float)i + 0.5f) / (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;
    }

    // Keep a copy of the normalized discrete probability mass per texel for MIS queries by direction.
    float* pdf_uv_mass = (float*)malloc(sizeof(float) * size_hw);
    if (pdf_uv_mass == NULL)
    {
        free(intermediate_buffer);
        return light;
    }
    memcpy(pdf_uv_mass, p_xy_original_pdf, sizeof(float) * size_hw);

    // --- 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];
        }
    }

    // Persist CDF tables for correct runtime sampling (keeps wi/pdf consistent)
    float* cdf_x_persist = (float*)malloc(sizeof(float) * size_w);
    float* cdf_y_transposed_persist = (float*)malloc(sizeof(float) * size_wh);
    if (cdf_x_persist == NULL || cdf_y_transposed_persist == NULL)
    {
        free(cdf_x_persist);
        free(cdf_y_transposed_persist);
        free(pdf_uv_mass);
        free(intermediate_buffer);
        return light;
    }
    memcpy(cdf_x_persist, cdf_x_buffer, sizeof(float) * size_w);
    memcpy(cdf_y_transposed_persist, cdf_y_transposed_buffer, sizeof(float) * size_wh);

    free(intermediate_buffer);

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

    data.total_weight = lumSum;
    data.pdf_uv_mass = pdf_uv_mass;
    data.cdf_x = cdf_x_persist;
    data.cdf_y_transposed = cdf_y_transposed_persist;

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

static inline float hdr_pdf_from_uv_mass(float uv_mass, float v, uint32_t height, uint32_t width)
{
    // v in [0,1] maps to theta in [0, PI]
    float theta = v * PI;
    float sin_theta = fmaxf(sinf(theta), 1e-4f);

    // Δω per texel for equirectangular map:
    // Δω = (2π/width) * (π/height) * sinθ = (2π^2 sinθ)/(W H)
    // pdf(ω) = uv_mass / Δω = uv_mass * (W H)/(2π^2 sinθ)
    float convert = (float)(height * width) / (2.0f * PI_TWO * sin_theta);
    float pdf = uv_mass * convert;

    return fminf(pdf, 1e6f);
}

static inline float hdr_sky_pdf_direction(const hdr_sky_data_t* sky_data, vec3s wi)
{
    vec2s uv = direction_to_equirectangular(wi);

    uint32_t x, y;
    uv_to_index(uv, sky_data->width, sky_data->height, &x, &y);

    float mass = 0.0f;
    if (sky_data->pdf_uv_mass != NULL)
    {
        mass = sky_data->pdf_uv_mass[y * sky_data->width + x];
    }

    return hdr_pdf_from_uv_mass(mass, uv.y, sky_data->height, sky_data->width);
}

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;
    }

    // Last hit was no geometry, directly sample the sky
    if (context->bvh_tree == NULL)
    {
        vec2s uv = direction_to_equirectangular(context->wo);
        
        // Calculate LOD based on ray spread angle
        float lod = 0.0f;
        if (context->spread_angle > 0.0f)
        {
            // Approximate texel footprint: spread_angle * resolution / (2*PI)
            // This is a rough estimation for equirectangular mapping
            float texels = context->spread_angle * (float)fmax(sky_data->width, sky_data->height) * INV_PI;
            lod = log2f(fmaxf(texels, 1.0f));
        }

        vec4s sky_light = texture_sample_lod(get_texture(context->textures, sky_data->texture), uv, lod);
        output.direct_lighting = glms_vec3_scale(glms_vec3_mul(glms_vec3(sky_light), throughput), sky_data->intensity);
        // Return the correct environment PDF for MIS when the BSDF-sampled ray escapes to the sky.
        output.pdf = hdr_sky_pdf_direction(sky_data, context->wo);
        return output;
    }

    uint32_t x_idx, y_idx;
    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);

    if (sky_data->cdf_x == NULL || sky_data->cdf_y_transposed == NULL || sky_data->pdf_uv_mass == NULL)
    {
        return output;
    }

    x_idx = lower_bound_float(sky_data->cdf_x, sky_data->width, r1);
    x_idx = (x_idx == sky_data->width) ? (sky_data->width - 1) : x_idx;

    y_idx = lower_bound_float(sky_data->cdf_y_transposed + x_idx * sky_data->height, sky_data->height, r2);
    y_idx = (y_idx == sky_data->height) ? (sky_data->height - 1) : y_idx;

    float u = ((float)x_idx + 0.5f) / (float)sky_data->width;
    float v = ((float)y_idx + 0.5f) / (float)sky_data->height;
    vec3s wi = equirectangular_to_direction(u, v);

    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, 0.0f, 0.0f);
    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 = (vec2s){u, v};
    const texture_t* hdri = get_texture(context->textures, sky_data->texture);
    vec4s pixel = texture_get_pixel(hdri, uv, 0);
    vec3s sky_light = glms_vec3_scale(glms_vec3(pixel), sky_data->intensity);

    float mass = sky_data->pdf_uv_mass[y_idx * sky_data->width + x_idx];
    float pdf = hdr_pdf_from_uv_mass(mass, v, sky_data->height, sky_data->width);
    if (pdf < 1e-12f)
    {
        return output;
    }

    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_SUCCESS;
    return output;
}

void hdr_sky_free(hdr_sky_data_t* data)
{
    free(data->pdf_uv_mass);
    free(data->cdf_x);
    free(data->cdf_y_transposed);

    data->pdf_uv_mass = NULL;
    data->cdf_x = NULL;
    data->cdf_y_transposed = NULL;
}