#ifndef BSDF_H
#define BSDF_H

#include "cglm/struct/vec3.h"

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

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

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

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

inline float pdf_blinn_phong_lobe(const vec3s normal, const vec3s wi, const vec3s wo, const 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;
}

#endif // BSDF_H