1#import bevy_render::maths::PI
2#import bevy_pbr::{
3    lighting,
4    utils::{sample_cosine_hemisphere, dir_to_cube_uv, sample_cube_dir, hammersley_2d, rand_vec2f}
5}
6
7struct FilteringConstants {
8    mip_level: f32,
9    sample_count: u32,
10    roughness: f32,
11    noise_size_bits: vec2u,
12}
13
14@group(0) @binding(0) var input_texture: texture_2d_array<f32>;
15@group(0) @binding(1) var input_sampler: sampler;
16@group(0) @binding(2) var output_texture: texture_storage_2d_array<rgba16float, write>;
17@group(0) @binding(3) var<uniform> constants: FilteringConstants;
18@group(0) @binding(4) var blue_noise_texture: texture_2d_array<f32>;
19
20// Sample an environment map with a specific LOD
21fn sample_environment(dir: vec3f, level: f32) -> vec4f {
22    let cube_uv = dir_to_cube_uv(dir);
23    return textureSampleLevel(input_texture, input_sampler, cube_uv.uv, cube_uv.face, level);
24}
25
26// Blue noise randomization
27#ifdef HAS_BLUE_NOISE
28fn sample_noise(pixel_coords: vec2u) -> vec4f {
29    let noise_size = vec2u(1) << constants.noise_size_bits;
30    let noise_size_mask = noise_size - vec2u(1u);
31    let noise_coords = pixel_coords & noise_size_mask;
32    let uv = vec2f(noise_coords) / vec2f(noise_size);
33    return textureSampleLevel(blue_noise_texture, input_sampler, uv, 0u, 0.0);
34}
35#else
36// pseudo-random numbers using RNG
37fn sample_noise(pixel_coords: vec2u) -> vec4f {
38    var rng_state: u32 = (pixel_coords.x * 3966231743u) ^ (pixel_coords.y * 3928936651u);
39    let rnd = rand_vec2f(&rng_state);
40    return vec4f(rnd, 0.0, 0.0);
41}
42#endif
43
44// Calculate LOD for environment map lookup using filtered importance sampling
45fn calculate_environment_map_lod(pdf: f32, width: f32, samples: f32) -> f32 {
46    // Solid angle of current sample
47    let omega_s = 1.0 / (samples * pdf);
48    
49    // Solid angle of a texel in the environment map
50    let omega_p = 4.0 * PI / (6.0 * width * width);
51    
52    // Filtered importance sampling: compute the correct LOD
53    return 0.5 * log2(omega_s / omega_p);
54}
55
56@compute
57@workgroup_size(8, 8, 1)
58fn generate_radiance_map(@builtin(global_invocation_id) global_id: vec3u) {
59    let size = textureDimensions(output_texture).xy;
60    let invSize = 1.0 / vec2f(size);
61    
62    let coords = vec2u(global_id.xy);
63    let face = global_id.z;
64    
65    if (any(coords >= size)) {
66        return;
67    }
68    
69    // Convert texture coordinates to direction vector
70    let uv = (vec2f(coords) + 0.5) * invSize;
71    let normal = sample_cube_dir(uv, face);
72    
73    // For radiance map, view direction = normal for perfect reflection
74    let view = normal;
75    
76    // Convert perceptual roughness to physical microfacet roughness
77    let perceptual_roughness = constants.roughness;
78    let roughness = lighting::perceptualRoughnessToRoughness(perceptual_roughness);
79    
80    // Get blue noise offset for stratification
81    let vector_noise = sample_noise(coords);
82    
83    var radiance = vec3f(0.0);
84    var total_weight = 0.0;
85    
86    // Skip sampling for mirror reflection (roughness = 0)
87    if (roughness < 0.01) {
88        radiance = sample_environment(normal, 0.0).rgb;
89        textureStore(output_texture, coords, face, vec4f(radiance, 1.0));
90        return;
91    }
92    
93    // For higher roughness values, use importance sampling
94    let sample_count = constants.sample_count;
95    
96    for (var i = 0u; i < sample_count; i++) {
97        // Get sample coordinates from Hammersley sequence with blue noise offset
98        var xi = hammersley_2d(i, sample_count);
99        xi = fract(xi + vector_noise.rg); // Apply Cranley-Patterson rotation
100        
101        // Sample the GGX distribution with the spherical-cap VNDF method
102        let light_dir = lighting::sample_visible_ggx(xi, roughness, normal, view);
103        
104        // Calculate weight (N·L)
105        let NdotL = dot(normal, light_dir);
106        
107        if (NdotL > 0.0) {
108            // Reconstruct the microfacet half-vector from view and light and compute PDF terms
109            let half_vector = normalize(view + light_dir);
110            let NdotH = dot(normal, half_vector);
111            let NdotV = dot(normal, view);
112            
113            // Get the geometric shadowing term
114            let G = lighting::G_Smith(NdotV, NdotL, roughness);
115            
116            // PDF that matches the bounded-VNDF sampling
117            let pdf = lighting::ggx_vndf_pdf(view, NdotH, roughness);
118            
119            // Calculate LOD using filtered importance sampling
120            // This is crucial to avoid fireflies and improve quality
121            let width = f32(size.x);
122            let lod = calculate_environment_map_lod(pdf, width, f32(sample_count));
123            
124            // Get source mip level - ensure we don't go negative
125            let source_mip = max(0.0, lod);
126            
127            // Sample environment map with the light direction
128            var sample_color = sample_environment(light_dir, source_mip).rgb;
129            
130            // Accumulate weighted sample, including geometric term
131            radiance += sample_color * NdotL * G;
132            total_weight += NdotL * G;
133        }
134    }
135    
136    // Normalize by total weight
137    if (total_weight > 0.0) {
138        radiance = radiance / total_weight;
139    }
140    
141    // Write result to output texture
142    textureStore(output_texture, coords, face, vec4f(radiance, 1.0));
143}
144
145@compute
146@workgroup_size(8, 8, 1)
147fn generate_irradiance_map(@builtin(global_invocation_id) global_id: vec3u) {
148    let size = textureDimensions(output_texture).xy;
149    let invSize = 1.0 / vec2f(size);
150    
151    let coords = vec2u(global_id.xy);
152    let face = global_id.z;
153    
154    if (any(coords >= size)) {
155        return;
156    }
157    
158    // Convert texture coordinates to direction vector
159    let uv = (vec2f(coords) + 0.5) * invSize;
160    let normal = sample_cube_dir(uv, face);
161    
162    var irradiance = vec3f(0.0);
163    
164    // Use uniform sampling on a hemisphere
165    for (var i = 0u; i < constants.sample_count; i++) {
166        // Build a deterministic RNG seed for this pixel / sample
167        // 4 randomly chosen 32-bit primes
168        var rng: u32 = (coords.x * 2131358057u) ^ (coords.y * 3416869721u) ^ (face * 1199786941u) ^ (i * 566200673u);
169
170        // Sample a direction from the upper hemisphere around the normal
171        var sample_dir = sample_cosine_hemisphere(normal, &rng);
172
173        // Sample environment with level 0 (no mip)
174        var sample_color = sample_environment(sample_dir, 0.0).rgb;
175        
176        // Accumulate the contribution
177        irradiance += sample_color;
178    }
179
180    // Normalize by number of samples (cosine-weighted sampling already accounts for PDF)
181    irradiance = irradiance / f32(constants.sample_count);
182    
183    // Write result to output texture
184    textureStore(output_texture, coords, face, vec4f(irradiance, 1.0));
185}