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

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

float alphaI(vec3 wi, vec2 alpha);

float lambdaSmith(vec3 wi, vec2 alpha) {
    if (wi.z >  0.9999) return  0.0;
    if (wi.z < -0.9999) return -1.0;

    float a = 1.0 / (alphaI(wi, alpha) * tan(acos(wi.z)));
    return 0.5 * (-1.0 + sign(a) * sqrt(1.0 + 1.0 / (a * a)));
}

float projectedArea(vec3 wi, vec2 alpha) {
    if (wi.z >  0.9999) return 1.0;
    if (wi.z < -0.9999) return 0.0;

    float sinThetaI = sin(acos(wi.z));
    float alphaI    = alphaI(wi, alpha);

    return 0.5 * (wi.z + sqrt(wi.z * wi.z + sinThetaI * sinThetaI * alphaI * alphaI));
}

float P22(vec2 slopes, vec2 alpha) {
    vec2 alphaSq = alpha * alpha; vec2 slopesSq = slopes * slopes;
    return 1.0 / (PI * alpha.x * alpha.y) / pow2(1.0 + slopesSq.x / alphaSq.x + slopesSq.y / alphaSq.y);
}

// Importance sampling the distribution of slopes
vec2 sampleP22_11(float thetaI, vec2 U) {
    vec2 slope = vec2(0.0);

    if (thetaI < EPS) return sqrt(U.x / (1.0 - U.x)) * sincos(TAU * U.y).yx;

    float cosThetaI = cos(thetaI), sinThetaI = sin(thetaI), tanThetaI = sinThetaI / cosThetaI;

    float projectedArea = 0.5 * (cosThetaI + 1.0);
    if (projectedArea < EPS || isnan(projectedArea)) return vec2(0.0);

    float A   = (2.0 * U.x * projectedArea) / cosThetaI - 1.0;
    float B   = tanThetaI;
    float tmp = 1.0 / (A * A - 1.0);
    float D   = sqrt(max0(B * B * tmp * tmp - (A * A - B * B) * tmp));

    vec2 slopeX  = vec2(B * tmp - D, B * tmp + D);
         slope.x = A < 0.0 || slopeX.y > 1.0 / tanThetaI ? slopeX.x : slopeX.y;

    float S  = sign(U.y - 0.5);
    float U2 = 2.0 * abs(U.y - 0.5);

    float Z = (U2 * (U2 * (U2 * 0.27385 - 0.73369) + 0.46341)) / (U2 * (U2 * (U2 * 0.093073 + 0.309420) - 1.0) + 0.597999);
    slope.y = S * Z * sqrt(1.0 + slope.x * slope.x);

    return slope;
}
