brdf.slang

#pragma once

#include "shared.h"

static const f32 PI = 3.14159265359;

// --- Low-discrepancy sampling helpers ---

f32 RadicalInverseVdC(u32 bits) {
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return f32(bits) * 2.3283064365386963e-10;
}

vec2 Hammersley(u32 i, u32 n) {
    return vec2(f32(i) / f32(n), RadicalInverseVdC(i));
}

vec3 ImportanceSampleGGX(vec2 xi, vec3 N, f32 roughness) {
    f32 a = roughness * roughness;
    f32 phi = 2.0 * PI * xi.x;
    f32 cosTheta = sqrt((1.0 - xi.y) / (1.0 + (a * a - 1.0) * xi.y));
    f32 sinTheta = sqrt(max(1.0 - cosTheta * cosTheta, 0.0));

    vec3 H;
    H.x = cos(phi) * sinTheta;
    H.y = sin(phi) * sinTheta;
    H.z = cosTheta;

    vec3 up = abs(N.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    vec3 T = normalize(cross(up, N));
    vec3 B = cross(N, T);

    return normalize(T * H.x + B * H.y + N * H.z);
}

// --- Normal distribution function (GGX / Trowbridge-Reitz) ---

f32 D_GGX(f32 NdotH, f32 roughness) {
    f32 a = roughness * roughness;
    f32 a2 = a * a;
    f32 denom = (NdotH * NdotH) * (a2 - 1.0) + 1.0;
    return a2 / max(PI * denom * denom, 1e-7);
}

// --- Geometry (Smith-GGX with Schlick-GGX approximation) ---

f32 G_SchlickGGX(f32 NdotV, f32 k) {
    return NdotV / (NdotV * (1.0 - k) + k);
}

f32 G_Smith_Direct(f32 NdotV, f32 NdotL, f32 roughness) {
    f32 r = roughness + 1.0;
    f32 k = (r * r) / 8.0;
    return G_SchlickGGX(NdotV, k) * G_SchlickGGX(NdotL, k);
}

f32 G_Smith_IBL(f32 NdotV, f32 NdotL, f32 roughness) {
    f32 k = (roughness * roughness) / 2.0;
    return G_SchlickGGX(NdotV, k) * G_SchlickGGX(NdotL, k);
}

// --- Fresnel (Schlick approximation) ---

vec3 F_Schlick(f32 cosTheta, vec3 F0) {
    return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}

vec3 F_SchlickRoughness(f32 cosTheta, vec3 F0, f32 roughness) {
    return F0 + (max(vec3(1.0 - roughness), F0) - F0) * pow(1.0 - cosTheta, 5.0);
}

// --- BRDF integration (split-sum approximation for the IBL LUT) ---

vec2 IntegrateBRDF(f32 NdotV, f32 roughness) {
    const u32 SAMPLE_COUNT = 1024u;

    vec3 V;
    V.x = sqrt(max(1.0 - NdotV * NdotV, 0.0));
    V.y = 0.0;
    V.z = NdotV;

    vec3 N = vec3(0.0, 0.0, 1.0);

    f32 A = 0.0;
    f32 B = 0.0;

    for (u32 i = 0u; i < SAMPLE_COUNT; ++i) {
        vec2 xi = Hammersley(i, SAMPLE_COUNT);
        vec3 H = ImportanceSampleGGX(xi, N, roughness);
        vec3 L = normalize(2.0 * dot(V, H) * H - V);

        f32 NdotL = saturate(L.z);
        f32 NdotH = saturate(H.z);
        f32 VdotH = saturate(dot(V, H));

        if (NdotL > 0.0) {
            f32 G = G_Smith_IBL(NdotV, NdotL, roughness);
            f32 G_Vis = (G * VdotH) / max(NdotH * NdotV, 0.0001);
            f32 Fc = pow(1.0 - VdotH, 5.0);

            A += (1.0 - Fc) * G_Vis;
            B += Fc * G_Vis;
        }
    }

    return vec2(A, B) / f32(SAMPLE_COUNT);
}