#include "ALgorithm/BSDF.h"

float power_heuristic(float pdf_a, float pdf_b)
{
    float a2 = pdf_a * pdf_a;
    float b2 = pdf_b * pdf_b;
    return a2 / (a2 + b2);
}

float roughness_to_blinn_phong_specular_exponent(float roughness)
{
    return glm_clamp(2.0f * 1.0f / (fmaxf(roughness * roughness, FLT_EPSILON)) - 2.0f, FLT_EPSILON, 1.0f / FLT_EPSILON);
}

vec3s fresnel_schlick_vec3(vec3s f0, float cos_theta)
{
    float x = 1.0f - cos_theta;
    float x5 = x * x * x * x * x;
    return glms_vec3_adds(glms_vec3_scale(f0, (1.0f - x5)), x5);
}

// BSDF sampling functions

float pdf_cosine_weighted_hemisphere(vec3s normal, vec3s wi)
{
    return fmaxf(glms_vec3_dot(wi, normal), 0.0f) / (float)M_PI;
}

float pdf_blinn_phong_lobe(vec3s normal, vec3s wi, vec3s wo, float roughness)
{
    // Check if wo and wi are on the same side of the surface normal geometry
    if (glms_vec3_dot(wo, normal) <= 0.0f || glms_vec3_dot(wi, normal) <= 0.0f)
    {
        return 0.0f; // Cannot scatter from below horizon to above, or vice versa
    }

    // Calculate the half-vector h based on input wo and wi
    vec3s wo_n = glms_vec3_normalize(wo); // Ensure normalized inputs if not guaranteed
    vec3s wi_n = glms_vec3_normalize(wi);
    vec3s h = glms_vec3_add(wo_n, wi_n);
    float h_len_sq = glms_vec3_norm2(h);
    if (h_len_sq < FLT_EPSILON)
    {
        return 0.0f; // wo and wi are opposite, highly unlikely for reflection
    }
    h = glms_vec3_scale(h, 1.0f / sqrtf(h_len_sq)); // Normalize h

    // Calculate Blinn-Phong specular exponent
    float specular_exponent = roughness_to_blinn_phong_specular_exponent(roughness);

    // PDF of sampling h (Blinn-Phong distribution)
    // D(h) = (specular_exponent + 1) / (2 * PI) * pow(max(0, dot(n, h)), specular_exponent)
    float n_dot_h = fmaxf(0.0f, glms_vec3_dot(normal, h));
    float pdf_h = (specular_exponent + 1.0f) / (2.0f * (float)M_PI) * powf(n_dot_h, specular_exponent);

    // Jacobian of the transformation from h to wi
    // jacobian = 1 / (4 * dot(wo, h))
    float wo_dot_h = fmaxf(FLT_EPSILON, glms_vec3_dot(wo_n, h)); // Use normalized wo, ensure > 0
    float jacobian = 1.0f / (4.0f * wo_dot_h);

    // PDF of sampling wi is pdf(h) * jacobian
    float pdf_spec = pdf_h * jacobian;

    return pdf_spec;
}

vec3s sample_cosine_weighted_hemisphere_z_angular(float angular, uint32_t index, uint32_t d1, uint32_t d2)
{
    float r1 = sobol_sample(index, d1);
    float r2 = sobol_sample(index, d2);

    float phi = 2.0f * (float)M_PI * r1;

    float cos_angular = cosf(angular);
    // Correctly sample cos(theta) for cosine weighting within the cone [cos_angular, 1]
    // cos_theta = sqrt(cos(angular)^2 + r2 * (1 - cos(angular)^2))
    float cos_theta = sqrtf(cos_angular * cos_angular + r2 * (1.0f - cos_angular * cos_angular));
    float sin_theta = sqrtf(1.0f - cos_theta * cos_theta);

    // Convert spherical coordinates (1, theta, phi) to Cartesian (Z-up)
    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s local_dir = {{x, y, z}};
    return local_dir;
}

// Function to generate a direction with cosine weighting around (0, 0, 1)
// This is the local coordinate sample.
vec3s sample_cosine_weighted_hemisphere_z(uint32_t index, uint32_t d1, uint32_t d2)
{
    float r1 = sobol_sample(index, d1);
    float r2 = sobol_sample(index, d2);

    float r = sqrtf(r1);
    float phi = 2.0f * (float)M_PI * r2;

    float disk_x = r * cosf(phi);
    float disk_y = r * sinf(phi);

    // Map point (disk_x, disk_y) on disk to hemisphere (Z-up)
    float x = disk_x;
    float y = disk_y;
    float z = sqrtf(1.0f - disk_x * disk_x - disk_y * disk_y); // z = sqrt(1 - r*r) = sqrt(1 - r1)

    vec3s local_dir = {{x, y, z}};
    return local_dir;
}

// Function to create an orthonormal basis (coordinate system) from a single vector (normal)
// w will be aligned with normal, u and v will be perpendicular.
void create_orthonormal_basis(vec3s direction, vec3s* u, vec3s* v)
{
    vec3s a;
    if (fabsf(direction.x) > 0.9f)
    {
        a = (vec3s){{0.0f, 1.0f, 0.0f}}; // Use y-axis
    }
    else
    {
        a = (vec3s){{1.0f, 0.0f, 0.0f}}; // Use x-axis
    }

    *u = glms_vec3_normalize(glms_vec3_cross(a, direction));
    *v = glms_vec3_normalize(glms_vec3_cross(direction, *u));
}

vec3s random_cosine_direction_angular(vec3s direction, float angular, uint32_t index, uint32_t d1, uint32_t d2)
{
    vec3s local_dir = sample_cosine_weighted_hemisphere_z_angular(angular, index, d1, d2);

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, local_dir.x);
    vec3s term_v = glms_vec3_scale(v, local_dir.y);
    vec3s term_w = glms_vec3_scale(direction, local_dir.z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);

    return world_dir;
}

// Samples a direction from the hemisphere oriented along 'normal'
// with a cosine-weighted distribution.
vec3s random_cosine_direction(vec3s direction, uint32_t index, uint32_t d1, uint32_t d2)
{
    vec3s local_dir = sample_cosine_weighted_hemisphere_z(index, d1, d2);

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, local_dir.x);
    vec3s term_v = glms_vec3_scale(v, local_dir.y);
    vec3s term_w = glms_vec3_scale(direction, local_dir.z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}

vec3s random_uniform_cdf_direction(vec3s direction, uint32_t index, uint32_t d1, uint32_t d2)
{
    float r1 = sobol_sample(index, d1);
    float r2 = sobol_sample(index, d2);

    float phi = 2.0f * (float)M_PI * r1;
    float cos_theta = 1.0f - r2 * 2.0f;
    float sin_theta = sqrtf(fmaxf(0.0f, 1.0f - cos_theta * cos_theta));

    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, x);
    vec3s term_v = glms_vec3_scale(v, y);
    vec3s term_w = glms_vec3_scale(direction, z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}

vec3s random_uniform_cdf_direction_angular(vec3s direction, uint32_t index, float angular, uint32_t d1, uint32_t d2)
{
    float r1 = sobol_sample(index, d1);
    float r2 = sobol_sample(index, d2);

    float cos_alpha = cosf(angular);
    float cos_theta = 1.0f - r1 * (1.0f - cos_alpha);
    float sin_theta = sqrtf(fmaxf(0.0f, 1.0f - cos_theta * cos_theta));

    float phi = 2.0f * (float)M_PI * r2;

    float x = sin_theta * cosf(phi);
    float y = sin_theta * sinf(phi);
    float z = cos_theta;

    vec3s u, v;
    create_orthonormal_basis(direction, &u, &v);

    vec3s term_u = glms_vec3_scale(u, x);
    vec3s term_v = glms_vec3_scale(v, y);
    vec3s term_w = glms_vec3_scale(direction, z);

    vec3s world_dir = glms_vec3_add(glms_vec3_add(term_u, term_v), term_w);
    return world_dir;
}