1// Incident-light, BRDF, and environment-sampling primitives. Final water
2// composition remains in the terminal material to keep imports one-way.
3
4#define_import_path aqua::light::incident
5
6#import bevy_pbr::{
7    clustered_forward as clustering,
8    lighting,
9    mesh_view_types,
10    shadows,
11    atmosphere::functions::{calculate_visible_sun_ratio, clamp_to_surface},
12    atmosphere::bruneton_functions::transmittance_lut_r_mu_to_uv,
13    mesh_view_bindings::{light_probes, lights, view},
14}
15#import bevy_pbr::mesh_view_bindings as view_bindings
16#import aqua::cascade::{LUMINANCE_EPSILON, LocalLightSample, SAFE_LENGTH_SQUARED, surface}
17#import bevy_aqua_core::material::{PrimaryLightState, SurfaceVertexOutput}
18
19const PI: f32 = 3.14159265359;
20
21const LUMINANCE_WEIGHTS: vec3<f32> = vec3(0.2126, 0.7152, 0.0722);
22
23const GODOT_NORMAL_FADE_RATE: f32 = 0.0175;
24const GODOT_NORMAL_MINIMUM_STRENGTH: f32 = 0.015;
25const GODOT_WATER_ALBEDO: vec3<f32> = vec3(
26    0.010022826, 0.01960665, 0.02721178,
27);
28const GODOT_SSS_MODIFIER: vec3<f32> = vec3(0.9, 1.15, 0.85);
29
30fn smith_masking_shadowing(cos_theta: f32, alpha: f32) -> f32 {
31    let sine = sqrt(max(1.0 - cos_theta * cos_theta, SAFE_LENGTH_SQUARED));
32    let a = cos_theta / max(alpha * sine, SAFE_LENGTH_SQUARED);
33    let a_squared = a * a;
34    return select(
35        0.0,
36        (1.0 - 1.259 * a + 0.396 * a_squared)
37            / (3.535 * a + 2.181 * a_squared),
38        a < 1.6,
39    );
40}
41
42fn ggx_distribution(cos_theta: f32, alpha: f32) -> f32 {
43    let alpha_squared = alpha * alpha;
44    let denominator = 1.0
45        + (alpha_squared - 1.0) * cos_theta * cos_theta;
46    return alpha_squared / (PI * denominator * denominator);
47}
48
49fn safe_normalize(value: vec3<f32>, fallback: vec3<f32>) -> vec3<f32> {
50    let length_squared = dot(value, value);
51    let normalized = value * inverseSqrt(max(length_squared, SAFE_LENGTH_SQUARED));
52    return select(fallback, normalized, length_squared > SAFE_LENGTH_SQUARED);
53}
54
55struct LocalLightContribution {
56    body: vec3<f32>,
57    reflection: vec3<f32>,
58}
59
60fn sample_local_light(
61    light_id: u32,
62    is_spot: bool,
63    world_position: vec3<f32>,
64    surface_normal: vec3<f32>,
65    frag_coord: vec2<f32>,
66) -> LocalLightSample {
67    let light = view_bindings::clustered_lights.data[light_id];
68    let to_light = light.position_radius.xyz - world_position;
69    let distance_squared = dot(to_light, to_light);
70    let direction = safe_normalize(to_light, vec3(0.0, 1.0, 0.0));
71    var attenuation = lighting::getDistanceAttenuation(
72        distance_squared,
73        light.color_inverse_square_range.w,
74    );
75    if is_spot {
76        var spot_direction = vec3(
77            light.light_custom_data.x,
78            0.0,
79            light.light_custom_data.y,
80        );
81        spot_direction.y = sqrt(max(
82            0.0,
83            1.0 - spot_direction.x * spot_direction.x
84                - spot_direction.z * spot_direction.z,
85        ));
86        if (light.flags
87            & mesh_view_types::POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE) != 0u {
88            spot_direction.y = -spot_direction.y;
89        }
90        let cosine = dot(-spot_direction, direction);
91        let cone = clamp(
92            cosine * light.light_custom_data.z + light.light_custom_data.w,
93            0.0,
94            1.0,
95        );
96        attenuation *= cone * cone;
97    }
98    var shadow = 1.0;
99    if (light.flags
100        & mesh_view_types::POINT_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u {
101        if is_spot {
102            shadow = shadows::fetch_spot_shadow(
103                light_id,
104                vec4(world_position, 1.0),
105                surface_normal,
106                light.shadow_map_near_z,
107                frag_coord,
108            );
109        } else {
110            shadow = shadows::fetch_point_shadow(
111                light_id,
112                vec4(world_position, 1.0),
113                surface_normal,
114                frag_coord,
115            );
116        }
117    }
118    return LocalLightSample(
119        direction,
120        light.color_inverse_square_range.rgb
121            * attenuation * shadow * view.exposure,
122    );
123}
124
125fn local_light_contribution(
126    sample: LocalLightSample,
127    lighting_normal: vec3<f32>,
128    to_view: vec3<f32>,
129    wave_height: f32,
130    sun_roughness: f32,
131    grazing: f32,
132    crest_transmission: f32,
133    sss_enabled: bool,
134) -> LocalLightContribution {
135    let dot_nl = max(dot(lighting_normal, sample.direction), 0.0);
136    let dot_nv = max(dot(lighting_normal, to_view), 2e-5);
137    var body = 0.5 * dot_nl * sample.radiance * GODOT_WATER_ALBEDO;
138    if sss_enabled {
139        let towards_light = pow(
140            max(dot(sample.direction, -to_view), 0.0),
141            surface.sss.z,
142        );
143        let light_mask = smith_masking_shadowing(surface.sun.y, dot_nv);
144        let sss_near = 0.5 * pow(dot_nv, 2.0);
145        let sss_height = max(0.0, wave_height + 2.5)
146            * pow(max(dot(sample.direction, -to_view), 0.0), 4.0)
147            * pow(
148                0.5 - 0.5 * dot(sample.direction, lighting_normal),
149                3.0,
150            );
151        body += (sss_height + sss_near)
152            * GODOT_SSS_MODIFIER
153            / (1.0 + light_mask)
154            * sample.radiance
155            * GODOT_WATER_ALBEDO;
156        body += (surface.sss.x + surface.sss.y * towards_light)
157            * surface.sss_tint.rgb
158            * sample.radiance
159            * grazing
160            * crest_transmission
161            * max(sample.direction.y, 0.0);
162    }
163
164    var reflection = vec3(0.0);
165    if dot_nl > 0.0 {
166        let halfway = safe_normalize(
167            sample.direction + to_view,
168            lighting_normal,
169        );
170        let light_mask = smith_masking_shadowing(sun_roughness, dot_nv);
171        let view_mask = smith_masking_shadowing(sun_roughness, max(dot_nl, 2e-5));
172        let distribution = ggx_distribution(
173            clamp(dot(lighting_normal, halfway), 0.0, 1.0),
174            sun_roughness,
175        );
176        let geometric_attenuation = 1.0 / (1.0 + light_mask + view_mask);
177        reflection = distribution
178            * geometric_attenuation
179            / (4.0 * dot_nv + 0.1)
180            * surface.sun.x
181            * sample.radiance;
182    }
183    return LocalLightContribution(body, reflection);
184}
185
186fn quat_rotate(rotation: vec4<f32>, value: vec3<f32>) -> vec3<f32> {
187    return value + 2.0 * cross(rotation.xyz, cross(rotation.xyz, value) + rotation.w * value);
188}
189
190fn view_direction(world_position: vec3<f32>) -> vec3<f32> {
191    let orthographic = view.clip_from_view[3].w == 1.0;
192    let orthographic_direction = vec3(
193        view.clip_from_world[0].z,
194        view.clip_from_world[1].z,
195        view.clip_from_world[2].z,
196    );
197    let perspective_direction = view.world_position.xyz - world_position;
198    return safe_normalize(
199        select(perspective_direction, orthographic_direction, orthographic),
200        vec3(0.0, 1.0, 0.0),
201    );
202}
203
204fn sample_environment(
205    reflection: vec3<f32>,
206    surface_normal: vec3<f32>,
207    perceptual_roughness: f32,
208) -> vec3<f32> {
209    // Godot renderer substrate: bend rough reflections toward the normal and
210    // suppress the below-horizon lobe before sampling the environment.
211    let bent_reflection = mix(
212        reflection,
213        surface_normal,
214        perceptual_roughness * perceptual_roughness,
215    );
216    let horizon = min(1.0 + dot(bent_reflection, surface_normal), 1.0);
217    // With no scene environment there is no radiance to reflect.
218    var radiance = vec3(0.0);
219#ifdef ENVIRONMENT_MAP
220    if light_probes.view_cubemap_index >= 0 {
221        // Work in probe-local space: arbitrary probe rotation must not tilt this
222        // guard back into the cubemap's ground hemisphere. Raise the rough lobe
223        // by its RMS slope cone. Where the original cone crosses the horizon,
224        // blend to a ground-free mip-0 sample (or zero if no probe exists)
225        // instead of a contaminated rough mip.
226        var probe_reflection = quat_rotate(
227            light_probes.view_rotation,
228            safe_normalize(bent_reflection, surface_normal),
229        );
230        let horizon_sine = perceptual_roughness
231            / sqrt(1.0 + perceptual_roughness * perceptual_roughness);
232        let roughness_enabled = perceptual_roughness > 0.0;
233        let guard_threshold = max(horizon_sine, LUMINANCE_EPSILON);
234        let ground_risk = select(
235            0.0,
236            1.0 - smoothstep(
237                guard_threshold,
238                2.0 * guard_threshold,
239                probe_reflection.y,
240            ),
241            roughness_enabled,
242        );
243        probe_reflection.y = select(
244            probe_reflection.y,
245            max(probe_reflection.y, horizon_sine),
246            roughness_enabled,
247        );
248        var sample_direction = safe_normalize(
249            probe_reflection,
250            vec3(0.0, 1.0, 0.0),
251        );
252        sample_direction.z = -sample_direction.z;
253        let mip = sqrt(perceptual_roughness)
254            * f32(light_probes.smallest_specular_mip_level_for_view);
255#ifdef MULTIPLE_LIGHT_PROBES_IN_ARRAY
256        radiance = textureSampleLevel(
257            view_bindings::specular_environment_maps[u32(light_probes.view_cubemap_index)],
258            view_bindings::environment_map_sampler,
259            sample_direction,
260            mip,
261        ).rgb * light_probes.intensity_for_view * view.exposure;
262#else
263        radiance = textureSampleLevel(
264            view_bindings::specular_environment_map,
265            view_bindings::environment_map_sampler,
266            sample_direction,
267            mip,
268        ).rgb * light_probes.intensity_for_view * view.exposure;
269#endif
270        if ground_risk > 0.0 {
271            var ground_free_sky = vec3(0.0);
272#ifdef MULTIPLE_LIGHT_PROBES_IN_ARRAY
273            ground_free_sky = textureSampleLevel(
274                view_bindings::specular_environment_maps[u32(light_probes.view_cubemap_index)],
275                view_bindings::environment_map_sampler,
276                sample_direction,
277                0.0,
278            ).rgb * light_probes.intensity_for_view * view.exposure;
279#else
280            ground_free_sky = textureSampleLevel(
281                view_bindings::specular_environment_map,
282                view_bindings::environment_map_sampler,
283                sample_direction,
284                0.0,
285            ).rgb * light_probes.intensity_for_view * view.exposure;
286#endif
287            radiance = mix(radiance, ground_free_sky, ground_risk);
288        }
289    }
290#endif
291    return radiance * horizon * horizon;
292}
293
294fn sample_diffuse_environment(surface_normal: vec3<f32>) -> vec3<f32> {
295    // Material colours are reflectance coefficients, never fallback emitters.
296    var irradiance = vec3(0.0);
297#ifdef ENVIRONMENT_MAP
298    if light_probes.view_cubemap_index >= 0 {
299        var sample_direction = quat_rotate(
300            light_probes.view_rotation,
301            surface_normal,
302        );
303        sample_direction.y = max(sample_direction.y, 0.0);
304        sample_direction = safe_normalize(
305            sample_direction,
306            vec3(0.0, 1.0, 0.0),
307        );
308        sample_direction.z = -sample_direction.z;
309#ifdef MULTIPLE_LIGHT_PROBES_IN_ARRAY
310        irradiance = textureSampleLevel(
311            view_bindings::diffuse_environment_maps[u32(light_probes.view_cubemap_index)],
312            view_bindings::environment_map_sampler,
313            sample_direction,
314            0.0,
315        ).rgb * light_probes.intensity_for_view * view.exposure;
316#else
317        irradiance = textureSampleLevel(
318            view_bindings::diffuse_environment_map,
319            view_bindings::environment_map_sampler,
320            sample_direction,
321            0.0,
322        ).rgb * light_probes.intensity_for_view * view.exposure;
323#endif
324    }
325#endif
326    return irradiance;
327}
328
329fn filtered_primary_light_color(
330    world_position: vec4<f32>,
331    direction_to_light: vec3<f32>,
332    sun_disk_angular_size: f32,
333    raw_color: vec3<f32>,
334) -> vec3<f32> {
335    var color = raw_color;
336#ifdef ATMOSPHERE
337    if surface.sun.z > 0.5 {
338        let atmosphere = view_bindings::atmosphere;
339        let atmosphere_position = (
340            atmosphere.world_to_atmosphere * world_position
341        ).xyz;
342        let clamped_position = clamp_to_surface(
343            atmosphere,
344            atmosphere_position,
345        );
346        let radius = length(clamped_position);
347        let local_up = normalize(clamped_position);
348        let light_direction = safe_normalize(
349            direction_to_light,
350            vec3(0.0, 1.0, 0.0),
351        );
352        let mu_light = dot(light_direction, local_up);
353        let uv = transmittance_lut_r_mu_to_uv(
354            atmosphere,
355            radius,
356            mu_light,
357        );
358        let transmittance = textureSampleLevel(
359            view_bindings::atmosphere_transmittance_texture,
360            view_bindings::atmosphere_transmittance_sampler,
361            uv,
362            0.0,
363        ).rgb;
364        let sun_visibility = calculate_visible_sun_ratio(
365            atmosphere,
366            radius,
367            mu_light,
368            sun_disk_angular_size,
369        );
370        color *= transmittance * sun_visibility;
371    }
372#endif
373    return color;
374}
375
376fn resolve_primary_light(
377    in: SurfaceVertexOutput,
378    normal: vec3<f32>,
379) -> PrimaryLightState {
380    let view_z = (view.view_from_world * in.world_position).z;
381    let is_orthographic = view.clip_from_view[3].w == 1.0;
382    let cluster_index = clustering::view_fragment_cluster_index(
383        in.position.xy,
384        view_z,
385        is_orthographic,
386    );
387    let ranges = clustering::unpack_clusterable_object_index_ranges(cluster_index);
388    var primary_light_shadow = 1.0;
389    var primary_light_color = vec3(0.0);
390    if lights.n_directional_lights > 0u {
391        let light = lights.directional_lights[0u];
392        primary_light_color = filtered_primary_light_color(
393            in.world_position,
394            light.direction_to_light,
395            light.sun_disk_angular_size,
396            light.color.rgb,
397        );
398        if (light.flags
399            & mesh_view_types::DIRECTIONAL_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u {
400            primary_light_shadow = shadows::fetch_directional_shadow(
401                0u,
402                in.world_position,
403                normal,
404                view_z,
405                in.position.xy,
406            );
407        }
408    }
409    let primary_light_radiance =
410        primary_light_color * view.exposure * primary_light_shadow;
411    return PrimaryLightState(
412        view_z,
413        ranges.first_point_light_index_offset,
414        ranges.first_spot_light_index_offset,
415        ranges.first_reflection_probe_index_offset,
416        primary_light_shadow,
417        primary_light_color,
418        primary_light_radiance,
419    );
420}
421
422struct IncidentDirectionalLight {
423    valid: bool,
424    direction: vec3<f32>,
425    color: vec3<f32>,
426    shadow: f32,
427}
428
429fn strongest_incident_directional_light(
430    in: SurfaceVertexOutput,
431    shadow_normal: vec3<f32>,
432    view_z: f32,
433) -> IncidentDirectionalLight {
434    var selected = false;
435    var selected_index = 0u;
436    var selected_direction = vec3(0.0, -1.0, 0.0);
437    var best_surface_illuminance = 0.0;
438    for (var index = 0u; index < lights.n_directional_lights; index += 1u) {
439        let light = lights.directional_lights[index];
440        let direction = safe_normalize(
441            light.direction_to_light,
442            vec3(0.0, -1.0, 0.0),
443        );
444        if direction.y <= 0.0 {
445            continue;
446        }
447        let surface_illuminance = direction.y * dot(
448            max(light.color.rgb, vec3(0.0)),
449            vec3(0.2126, 0.7152, 0.0722),
450        );
451        if surface_illuminance > best_surface_illuminance {
452            selected = true;
453            selected_index = index;
454            selected_direction = direction;
455            best_surface_illuminance = surface_illuminance;
456        }
457    }
458    if !selected {
459        return IncidentDirectionalLight(
460            false,
461            selected_direction,
462            vec3(0.0),
463            0.0,
464        );
465    }
466    let light = lights.directional_lights[selected_index];
467    let color = filtered_primary_light_color(
468        in.world_position,
469        light.direction_to_light,
470        light.sun_disk_angular_size,
471        light.color.rgb,
472    );
473    var shadow = 1.0;
474    if (light.flags
475        & mesh_view_types::DIRECTIONAL_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u {
476        shadow = shadows::fetch_directional_shadow(
477            selected_index,
478            in.world_position,
479            shadow_normal,
480            view_z,
481            in.position.xy,
482        );
483    }
484    return IncidentDirectionalLight(
485        true,
486        selected_direction,
487        color,
488        shadow,
489    );
490}
491