/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

/*
    [Notes]:
        See "shaders/include/fragment/heitz/heitz.glsl" for references and notes.
*/

// Projected roughness in the incident direction wi
float alphaI(vec3 wi, vec2 alpha) {
    vec3 wiSq    = wi * wi; 
    vec2 alphaSq = alpha * alpha;

    float invSinThetaSq = 1.0 / (1.0 - wiSq.z);
    return sqrt(wiSq.x * invSinThetaSq * alphaSq.x + wiSq.y * invSinThetaSq * alphaSq.y);
}

float D(vec3 wm, vec2 alpha) {
    if (wm.z <= 0.0) return 0.0;
    return P22(-wm.xy / wm.z, alpha) / (wm.z * wm.z * wm.z * wm.z);
}

float D_wi(vec3 wi, vec3 wm, vec2 alpha) {
    if (wm.z <= 0.0) return 0.0;

    float projectedArea = projectedArea(wi, alpha);
    if (projectedArea == 0.0) return 0.0;

    return (max0(dot(wi, wm)) * D(wm, alpha)) / projectedArea;
}

vec3 sampleD_wi(vec3 wi, vec2 U, vec2 alpha) {

    // Stretch to match configuration with alpha = 1.0
    vec3 wi_11    = normalize(vec3(alpha * wi.xy, wi.z));
    vec2 slope_11 = sampleP22_11(acos(wi_11.z), U);

    // Align with view direction
    float phi    = atan(wi_11.y, wi_11.x);
    float cosPhi = cos(phi), sinPhi = sin(phi);
    vec2  slope  = vec2(cosPhi * slope_11.x - sinPhi * slope_11.y, sinPhi * slope_11.x + cosPhi * slope_11.y);
    
    // Unstretch
    slope *= alpha;

    // Preventing numerical instability
    if (isnan(slope.x) || isinf(slope.x)) {
        return wi.z > 0.0 ? vec3(0.0, 0.0, 1.0) : normalize(vec3(wi.xy, 0.0));
    }

    // Compute normal
    return normalize(vec3(-slope, 1.0));
}
